Files
tfm_ainventory/backend/services/image_processing.py

474 lines
16 KiB
Python

"""
OpenCV-based image processing pipeline for smart photo handling.
Handles:
- EXIF orientation detection and auto-rotation
- Smart cropping using OpenCV contour detection
- Text orientation detection using Hough lines
- Resize and compression to 1200px
- Thumbnail generation (200px square)
- Fallback to Pillow for basic processing if OpenCV fails
"""
import io
import logging
from typing import Dict, Optional, Tuple
from PIL import Image
from PIL.ExifTags import TAGS
import piexif
import cv2
import numpy as np
logger = logging.getLogger(__name__)
class ImageProcessor:
"""Service for processing uploaded images with smart features."""
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB
LONG_SIDE = 1200
THUMBNAIL_SIZE = 200
JPEG_QUALITY = 85
# Algorithm parameters (Issue 4: DRY Violation - Magic Numbers)
CANNY_CROP_THRESHOLDS = (100, 200)
CANNY_TEXT_THRESHOLDS = (50, 150)
HOUGH_THRESHOLD = 100
CROP_PADDING_FACTOR = 0.1
ANGLE_UPSIDE_DOWN_THRESHOLD = 80 # degrees
ANGLE_SIDEWAYS_THRESHOLD = 45 # degrees
def __init__(self):
"""Initialize the image processor."""
self.logger = logger
def process_photo(
self, file_bytes: bytes, crop_bounds: Optional[Dict] = None, rotation_degrees: float = 0
) -> Dict:
"""
Process a photo with EXIF rotation, smart cropping, and compression.
Args:
file_bytes: Raw image file bytes
crop_bounds: Optional manual crop bounds {x, y, width, height}
rotation_degrees: Optional manual rotation in degrees (applied after crop)
Returns:
{
'status': 'success' | 'error',
'cropped_image_bytes': bytes or None,
'thumbnail_bytes': bytes or None,
'original_size': (width, height),
'crop_size': (width, height) or None,
'text_angle': float or None,
'metadata': {
'exif_orientation': int,
'crop_method': 'manual' | 'opencv' | 'pillow' | 'none',
'file_size_bytes': int
}
}
"""
try:
# Validate file size
if len(file_bytes) > self.MAX_FILE_SIZE:
return {
'status': 'error',
'error': f'File too large: {len(file_bytes)} > {self.MAX_FILE_SIZE}',
'cropped_image_bytes': None,
'thumbnail_bytes': None,
}
# Open image with PIL
image = Image.open(io.BytesIO(file_bytes))
original_size = image.size
# Extract and apply EXIF orientation
exif_orientation = self._extract_exif_orientation(image)
if exif_orientation and exif_orientation > 1:
image = self._rotate_by_orientation(image, exif_orientation)
self.logger.info(f"Applied EXIF rotation: {exif_orientation}")
# Smart cropping
cropped_image = image
crop_size = None
text_angle = None
crop_method = 'none'
if crop_bounds:
# Manual crop bounds provided
cropped_image = image.crop(
(
crop_bounds['x'],
crop_bounds['y'],
crop_bounds['x'] + crop_bounds['width'],
crop_bounds['y'] + crop_bounds['height'],
)
)
crop_size = cropped_image.size
crop_method = 'manual'
self.logger.info(f"Applied manual crop: {crop_size}")
else:
# Try OpenCV smart crop
try:
crop_result = self._smart_crop_opencv(image)
if crop_result is not None:
cropped_image, crop_size = crop_result
crop_method = 'opencv'
self.logger.info(f"Applied OpenCV crop: {crop_size}")
# Detect text orientation within the cropped region
text_angle, angle_status = self._detect_text_orientation(
cropped_image
)
if text_angle is not None:
self.logger.info(
f"Detected text angle: {text_angle}° ({angle_status})"
)
if angle_status in ['upside_down', 'sideways']:
cropped_image = self._rotate_image(
cropped_image, text_angle
)
else:
crop_method = 'pillow'
except (IOError, ValueError, cv2.error) as e:
# Fallback to Pillow if OpenCV fails
self.logger.warning(
f"OpenCV crop failed, falling back to Pillow: {e}"
)
crop_method = 'pillow'
# Apply manual rotation if provided
if abs(rotation_degrees) > 0.5:
cropped_image = cropped_image.rotate(-rotation_degrees, expand=True)
self.logger.info(f"Applied manual rotation: {rotation_degrees}°")
# Resize and compress
compressed_bytes = self._resize_and_compress(cropped_image)
# Generate thumbnail
thumbnail_bytes = self._generate_thumbnail(image)
return {
'status': 'success',
'cropped_image_bytes': compressed_bytes,
'thumbnail_bytes': thumbnail_bytes,
'original_size': original_size,
'crop_size': crop_size,
'text_angle': text_angle,
'metadata': {
'exif_orientation': exif_orientation or 1,
'crop_method': crop_method,
'file_size_bytes': len(file_bytes),
},
}
except (IOError, ValueError, cv2.error) as e:
self.logger.error(f"Image processing failed: {e}")
return {
'status': 'error',
'error': str(e),
'cropped_image_bytes': None,
'thumbnail_bytes': None,
}
def _extract_exif_orientation(self, image: Image.Image) -> Optional[int]:
"""
Extract EXIF orientation tag from image.
Returns:
Orientation value (1-8) or None if not present
"""
try:
# Use piexif for EXIF extraction (avoiding private PIL API)
if hasattr(image, 'info') and 'exif' in image.info:
exif_dict = piexif.load(image.info['exif'])
orientation = exif_dict['0th'].get(piexif.ImageIFD.Orientation)
if orientation:
return orientation
return None
except (piexif.InvalidImageData, ValueError, IOError) as e:
self.logger.debug(f"Could not extract EXIF orientation: {e}")
return None
def _rotate_by_orientation(
self, image: Image.Image, orientation: int
) -> Image.Image:
"""
Rotate image based on EXIF orientation tag.
Args:
image: PIL Image
orientation: EXIF orientation value (1-8)
Returns:
Rotated PIL Image
"""
if orientation == 1:
return image
elif orientation == 2:
return image.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
elif orientation == 3:
return image.transpose(Image.Transpose.ROTATE_180)
elif orientation == 4:
return image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
elif orientation == 5:
return image.transpose(Image.Transpose.TRANSPOSE)
elif orientation == 6:
return image.transpose(Image.Transpose.ROTATE_270)
elif orientation == 7:
return image.transpose(Image.Transpose.TRANSVERSE)
elif orientation == 8:
return image.transpose(Image.Transpose.ROTATE_90)
return image
def _smart_crop_opencv(
self, image: Image.Image
) -> Optional[Tuple[Image.Image, Tuple[int, int]]]:
"""
Use OpenCV to detect and crop the main object in the image.
Args:
image: PIL Image
Returns:
Tuple of (cropped PIL Image, crop size) or None if no contours found
"""
try:
# Convert PIL image to OpenCV format
cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
# Edge detection
edges = cv2.Canny(gray, *self.CANNY_CROP_THRESHOLDS)
# Find contours
contours, _ = cv2.findContours(
edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)
if not contours:
self.logger.debug("No contours found in image")
return None
# Get bounding box of largest contour
largest_contour = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(largest_contour)
# Apply padding around bounds
pad_x = int(w * self.CROP_PADDING_FACTOR)
pad_y = int(h * self.CROP_PADDING_FACTOR)
x1 = max(0, x - pad_x)
y1 = max(0, y - pad_y)
x2 = min(cv_image.shape[1], x + w + pad_x)
y2 = min(cv_image.shape[0], y + h + pad_y)
# Crop image
cropped = image.crop((x1, y1, x2, y2))
crop_size = cropped.size
self.logger.debug(
f"OpenCV crop bounds: ({x1}, {y1}, {x2}, {y2}), size: {crop_size}"
)
return cropped, crop_size
except (IOError, ValueError, cv2.error) as e:
self.logger.warning(f"OpenCV smart crop failed: {e}")
return None
def _detect_text_orientation(
self, image: Image.Image
) -> Tuple[Optional[float], str]:
"""
Detect text orientation using Hough line transform.
Args:
image: PIL Image
Returns:
Tuple of (angle in degrees, status string)
status: 'normal', 'upside_down', 'sideways', 'not_detected'
"""
try:
# Convert to OpenCV format
cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
# DoS prevention: check resolution
if cv_image.shape[0] * cv_image.shape[1] > 2000 * 2000: # >4MP
self.logger.warning("ROI too large for text detection, skipping")
return None, 'not_detected'
gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
# Edge detection
edges = cv2.Canny(gray, *self.CANNY_TEXT_THRESHOLDS)
# Hough line detection
lines = cv2.HoughLines(edges, 1, np.pi / 180, self.HOUGH_THRESHOLD)
if lines is None or len(lines) == 0:
self.logger.debug("No lines detected for text orientation")
return None, 'not_detected'
# Extract angles from lines
angles = []
for line in lines:
rho, theta = line[0]
angle = np.degrees(theta)
angles.append(angle)
# Normalize angles to 0-180 range
angles = np.array(angles)
angles = np.where(angles > 90, angles - 180, angles)
# Find dominant angle
mean_angle = np.mean(angles)
# Determine orientation status
status = 'normal'
corrected_angle = mean_angle
# Check for upside-down text (~180°)
if abs(mean_angle) > self.ANGLE_UPSIDE_DOWN_THRESHOLD:
status = 'upside_down'
corrected_angle = mean_angle + 180 if mean_angle > 0 else mean_angle - 180
# Check for sideways text (~90°)
elif abs(mean_angle) > self.ANGLE_SIDEWAYS_THRESHOLD:
status = 'sideways'
self.logger.debug(
f"Text orientation: angle={mean_angle:.1f}°, status={status}"
)
return corrected_angle, status
except (IOError, ValueError, cv2.error) as e:
self.logger.warning(f"Text orientation detection failed: {e}")
return None, 'not_detected'
def _rotate_image(self, image: Image.Image, angle: float) -> Image.Image:
"""
Rotate image by specified angle.
Args:
image: PIL Image
angle: Rotation angle in degrees
Returns:
Rotated PIL Image
"""
return image.rotate(angle, expand=False, fillcolor='white')
def _resize_and_compress(self, image: Image.Image) -> bytes:
"""
Resize image to 1200px on long side and compress to JPEG.
Args:
image: PIL Image
Returns:
Compressed JPEG bytes
"""
# Get current size
width, height = image.size
max_dim = max(width, height)
# Only resize if necessary
if max_dim > self.LONG_SIDE:
scale = self.LONG_SIDE / max_dim
new_width = int(width * scale)
new_height = int(height * scale)
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
self.logger.debug(
f"Resized from {(width, height)} to {(new_width, new_height)}"
)
# Convert to RGB if necessary (for JPEG)
if image.mode in ('RGBA', 'LA', 'P'):
# Extract alpha channel if present
mask = None
if image.mode == 'RGBA':
alpha = image.split()[3]
mask = alpha
elif image.mode == 'LA':
alpha = image.split()[1]
mask = alpha
rgb_image = Image.new('RGB', image.size, (255, 255, 255))
rgb_image.paste(image, mask=mask)
image = rgb_image
# Compress to JPEG
output = io.BytesIO()
image.save(output, format='JPEG', quality=self.JPEG_QUALITY, optimize=True)
compressed_bytes = output.getvalue()
self.logger.debug(
f"Compressed to JPEG: {len(compressed_bytes)} bytes, "
f"quality={self.JPEG_QUALITY}"
)
return compressed_bytes
def _generate_thumbnail(self, image: Image.Image) -> bytes:
"""
Generate 200px square thumbnail with center crop.
Args:
image: PIL Image
Returns:
Thumbnail JPEG bytes
"""
try:
# Get current size
width, height = image.size
# Center crop to square
min_dim = min(width, height)
left = (width - min_dim) // 2
top = (height - min_dim) // 2
right = left + min_dim
bottom = top + min_dim
square = image.crop((left, top, right, bottom))
# Resize to thumbnail size
thumbnail = square.resize(
(self.THUMBNAIL_SIZE, self.THUMBNAIL_SIZE),
Image.Resampling.LANCZOS,
)
# Convert to RGB if necessary
if thumbnail.mode in ('RGBA', 'LA', 'P'):
# Extract alpha channel if present
mask = None
if thumbnail.mode == 'RGBA':
alpha = thumbnail.split()[3]
mask = alpha
elif thumbnail.mode == 'LA':
alpha = thumbnail.split()[1]
mask = alpha
rgb_thumbnail = Image.new('RGB', thumbnail.size, (255, 255, 255))
rgb_thumbnail.paste(thumbnail, mask=mask)
thumbnail = rgb_thumbnail
# Compress
output = io.BytesIO()
thumbnail.save(output, format='JPEG', quality=self.JPEG_QUALITY, optimize=True)
thumbnail_bytes = output.getvalue()
self.logger.debug(
f"Generated thumbnail: {self.THUMBNAIL_SIZE}x{self.THUMBNAIL_SIZE}, "
f"{len(thumbnail_bytes)} bytes"
)
return thumbnail_bytes
except (IOError, ValueError, cv2.error) as e:
self.logger.error(f"Thumbnail generation failed: {e}")
return b''