468 lines
16 KiB
Python
468 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
|
|
) -> 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}
|
|
|
|
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'
|
|
|
|
# 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''
|