ilia/punimtag
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Enhance face processing with EXIF orientation handling and database updates

Browse Source 
This commit introduces a comprehensive EXIF orientation handling system to improve face processing accuracy. Key changes include the addition of an `exif_orientation` field in the database schema, updates to the `FaceProcessor` class for applying orientation corrections before face detection, and the implementation of a new `EXIFOrientationHandler` utility for managing image orientation. The README has been updated to document these enhancements, including recent fixes for face orientation issues and improved face extraction logic. Additionally, tests for EXIF orientation handling have been added to ensure functionality and reliability.
This commit is contained in:
tanyar09
2025-10-17 13:50:47 -04:00
parent 2828b9966b
commit 5db41b63ef
5 changed files with 484 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/database.py
+9 -6
View File
@@ -88,6 +88,7 @@ class DatabaseManager:
detector_backend TEXT DEFAULT 'retinaface',
model_name TEXT DEFAULT 'ArcFace',
face_confidence REAL DEFAULT 0.0,
exif_orientation INTEGER DEFAULT NULL,
FOREIGN KEY (photo_id) REFERENCES photos (id),
FOREIGN KEY (person_id) REFERENCES people (id)
)
@@ -231,7 +232,8 @@ class DatabaseManager:
quality_score: float = 0.0, person_id: Optional[int] = None,
detector_backend: str = 'retinaface',
model_name: str = 'ArcFace',
face_confidence: float = 0.0) -> int:
face_confidence: float = 0.0,
exif_orientation: Optional[int] = None) -> int:
"""Add a face to the database and return its ID
Args:
@@ -244,6 +246,7 @@ class DatabaseManager:
detector_backend: DeepFace detector used (retinaface, mtcnn, opencv, ssd)
model_name: DeepFace model used (ArcFace, Facenet, etc.)
face_confidence: Confidence from DeepFace detector
exif_orientation: EXIF orientation value (1-8) for coordinate transformation
Returns:
Face ID
@@ -252,10 +255,10 @@ class DatabaseManager:
cursor = conn.cursor()
cursor.execute('''
INSERT INTO faces (photo_id, person_id, encoding, location, confidence,
quality_score, detector_backend, model_name, face_confidence)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
quality_score, detector_backend, model_name, face_confidence, exif_orientation)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
''', (photo_id, person_id, encoding, location, confidence, quality_score,
detector_backend, model_name, face_confidence))
detector_backend, model_name, face_confidence, exif_orientation))
return cursor.lastrowid
def update_face_person(self, face_id: int, person_id: Optional[int]):
@@ -388,11 +391,11 @@ class DatabaseManager:
return result[0] if result else None
def get_face_photo_info(self, face_id: int) -> Optional[Tuple]:
"""Get photo information for a specific face"""
"""Get photo information for a specific face including EXIF orientation"""
with self.get_db_connection() as conn:
cursor = conn.cursor()
cursor.execute('''
SELECT f.photo_id, p.filename, f.location
SELECT f.photo_id, p.filename, f.location, f.exif_orientation
FROM faces f
JOIN photos p ON f.photo_id = p.id
WHERE f.id = ?
src/core/face_processing.py
+66 -77
View File
@@ -31,6 +31,7 @@ from src.core.config import (
MAX_FACE_SIZE
)
from src.core.database import DatabaseManager
from src.utils.exif_utils import EXIFOrientationHandler
class FaceProcessor:
@@ -125,24 +126,55 @@ class FaceProcessor:
continue
try:
# Get EXIF orientation information
exif_orientation = EXIFOrientationHandler.get_exif_orientation(photo_path)
# Process with DeepFace
if self.verbose >= 1:
print(f"📸 Processing: {filename}")
if exif_orientation and exif_orientation != 1:
print(f" 📐 EXIF orientation: {exif_orientation} ({EXIFOrientationHandler._get_orientation_description(exif_orientation)})")
elif self.verbose == 0:
print(".", end="", flush=True)
if self.verbose >= 2:
print(f" 🔍 Using DeepFace: detector={self.detector_backend}, model={self.model_name}")
# Apply EXIF orientation correction before face detection
corrected_image, original_orientation = EXIFOrientationHandler.correct_image_orientation_from_path(photo_path)
if corrected_image is not None and original_orientation and original_orientation != 1:
# Save corrected image temporarily for DeepFace processing
import tempfile
temp_dir = tempfile.gettempdir()
temp_filename = f"corrected_{photo_id}_{filename}"
temp_path = os.path.join(temp_dir, temp_filename)
corrected_image.save(temp_path, "JPEG", quality=95)
# Use corrected image for face detection
face_detection_path = temp_path
if self.verbose >= 2:
print(f" 📐 Using EXIF-corrected image for face detection (orientation {original_orientation})")
else:
# Use original image if no correction needed
face_detection_path = photo_path
# Use DeepFace.represent() to get face detection and encodings
results = DeepFace.represent(
img_path=photo_path,
img_path=face_detection_path,
model_name=self.model_name,
detector_backend=self.detector_backend,
enforce_detection=DEEPFACE_ENFORCE_DETECTION,
align=DEEPFACE_ALIGN_FACES
)
# Clean up temporary file if created
if 'temp_path' in locals() and os.path.exists(temp_path):
try:
os.remove(temp_path)
except:
pass
if not results:
if self.verbose >= 1:
print(f" 👤 No faces found")
@@ -193,7 +225,7 @@ class FaceProcessor:
image_np = np.array(image)
quality_score = self._calculate_face_quality_score(image_np, face_location_dict)
# Store in database with DeepFace format
# Store in database with DeepFace format and EXIF orientation
self.db.add_face(
photo_id=photo_id,
encoding=embedding.tobytes(),
@@ -203,7 +235,8 @@ class FaceProcessor:
person_id=None,
detector_backend=self.detector_backend,
model_name=self.model_name,
face_confidence=face_confidence
face_confidence=face_confidence,
exif_orientation=exif_orientation
)
if self.verbose >= 3:
@@ -407,7 +440,7 @@ class FaceProcessor:
return 0.5 # Default medium quality on error
def _extract_face_crop(self, photo_path: str, location: dict, face_id: int) -> str:
"""Extract and save individual face crop for identification with caching"""
"""Extract and save individual face crop for identification with EXIF orientation correction"""
try:
# Check cache first
cache_key = f"{photo_path}_{location}_{face_id}"
@@ -429,6 +462,35 @@ class FaceProcessor:
if not isinstance(location, dict):
raise ValueError(f"Expected DeepFace dict format, got {type(location)}")
# Get EXIF orientation from database
face_info = self.db.get_face_photo_info(face_id)
exif_orientation = face_info[3] if face_info and len(face_info) > 3 else None
# Load the image with EXIF orientation correction
# Since we now apply correction before face detection, we need to apply it here too
corrected_image, original_orientation = EXIFOrientationHandler.correct_image_orientation_from_path(photo_path)
if corrected_image is not None and original_orientation and original_orientation != 1:
# Only apply correction if orientation is not 1 (normal)
image = corrected_image
# Use the detected orientation if not stored in database
if exif_orientation is None:
exif_orientation = original_orientation
else:
# Use original image if no correction needed or correction fails
image = Image.open(photo_path)
# Transform face coordinates if image was rotated
# TEMPORARILY DISABLED FOR TESTING - coordinate transformation might be causing issues
# if exif_orientation and exif_orientation != 1:
# # Get original image dimensions for coordinate transformation
# with Image.open(photo_path) as original_image:
# original_width, original_height = original_image.size
#
# # Transform coordinates based on orientation correction
# location = EXIFOrientationHandler.transform_face_coordinates(
# location, original_width, original_height, exif_orientation
# )
left = location.get('x', 0)
top = location.get('y', 0)
width = location.get('w', 0)
@@ -436,9 +498,6 @@ class FaceProcessor:
right = left + width
bottom = top + height
# Load the image
image = Image.open(photo_path)
# Add padding around the face (20% of face size)
face_width = right - left
face_height = bottom - top
@@ -784,76 +843,6 @@ class FaceProcessor:
"""Update person encodings when a face is identified"""
self.db.update_person_encodings(person_id)
def _extract_face_crop(self, photo_path: str, location: dict, face_id: int) -> str:
"""Extract and save individual face crop for identification with caching"""
try:
# Check cache first
cache_key = f"{photo_path}_{location}_{face_id}"
if cache_key in self._image_cache:
cached_path = self._image_cache[cache_key]
# Verify the cached file still exists
if os.path.exists(cached_path):
return cached_path
else:
# Remove from cache if file doesn't exist
del self._image_cache[cache_key]
# Parse location from string format (DeepFace format only)
if isinstance(location, str):
import ast
location = ast.literal_eval(location)
# DeepFace format: {x, y, w, h}
if not isinstance(location, dict):
raise ValueError(f"Expected DeepFace dict format, got {type(location)}")
left = location.get('x', 0)
top = location.get('y', 0)
width = location.get('w', 0)
height = location.get('h', 0)
right = left + width
bottom = top + height
# Load the image
image = Image.open(photo_path)
# Add padding around the face (20% of face size)
face_width = right - left
face_height = bottom - top
padding_x = int(face_width * 0.2)
padding_y = int(face_height * 0.2)
# Calculate crop bounds with padding
crop_left = max(0, left - padding_x)
crop_top = max(0, top - padding_y)
crop_right = min(image.width, right + padding_x)
crop_bottom = min(image.height, bottom + padding_y)
# Crop the face
face_crop = image.crop((crop_left, crop_top, crop_right, crop_bottom))
# Create temporary file for the face crop
temp_dir = tempfile.gettempdir()
face_filename = f"face_{face_id}_crop.jpg"
face_path = os.path.join(temp_dir, face_filename)
# Resize for better viewing (minimum 200px width)
if face_crop.width < 200:
ratio = 200 / face_crop.width
new_width = 200
new_height = int(face_crop.height * ratio)
face_crop = face_crop.resize((new_width, new_height), Image.Resampling.LANCZOS)
face_crop.save(face_path, "JPEG", quality=95)
# Cache the result
self._image_cache[cache_key] = face_path
return face_path
except Exception as e:
if self.verbose >= 1:
print(f"⚠️ Could not extract face crop: {e}")
return None
def _create_comparison_image(self, unid_crop_path: str, match_crop_path: str, person_name: str, confidence: float) -> str:
"""Create a side-by-side comparison image"""
src/utils/exif_utils.py
+252
View File
@@ -0,0 +1,252 @@
#!/usr/bin/env python3
"""
EXIF orientation handling utilities for PunimTag
Handles image orientation correction based on EXIF data
"""
from PIL import Image
from typing import Tuple, Optional, Dict
class EXIFOrientationHandler:
"""Handles EXIF orientation detection and correction"""
# EXIF orientation tag mapping
ORIENTATION_TAG = 274 # EXIF orientation tag ID
# Orientation values and their corresponding transformations
ORIENTATION_TRANSFORMS = {
1: None, # Normal (no rotation needed)
2: Image.Transpose.FLIP_LEFT_RIGHT, # Mirrored horizontally
3: Image.Transpose.ROTATE_180, # Rotated 180°
4: Image.Transpose.FLIP_TOP_BOTTOM, # Mirrored vertically
5: [Image.Transpose.FLIP_LEFT_RIGHT, Image.Transpose.ROTATE_90], # Mirrored horizontally + rotated 90° CCW
6: Image.Transpose.ROTATE_270, # Rotated 90° CW
7: [Image.Transpose.FLIP_LEFT_RIGHT, Image.Transpose.ROTATE_270], # Mirrored horizontally + rotated 90° CW
8: Image.Transpose.ROTATE_90, # Rotated 90° CCW
}
@staticmethod
def get_exif_orientation(image_path: str) -> Optional[int]:
"""Get EXIF orientation value from image
Args:
image_path: Path to the image file
Returns:
Orientation value (1-8) or None if not found/invalid
"""
try:
with Image.open(image_path) as image:
exif = image.getexif()
if exif is None:
return None
# Get orientation value
orientation = exif.get(EXIFOrientationHandler.ORIENTATION_TAG)
return orientation if orientation in EXIFOrientationHandler.ORIENTATION_TRANSFORMS else None
except Exception:
return None
@staticmethod
def get_orientation_info(image_path: str) -> Dict[str, any]:
"""Get comprehensive orientation information from image
Args:
image_path: Path to the image file
Returns:
Dictionary with orientation information
"""
orientation = EXIFOrientationHandler.get_exif_orientation(image_path)
return {
'orientation': orientation,
'needs_correction': orientation is not None and orientation != 1,
'transforms': EXIFOrientationHandler.ORIENTATION_TRANSFORMS.get(orientation, None),
'description': EXIFOrientationHandler._get_orientation_description(orientation)
}
@staticmethod
def _get_orientation_description(orientation: Optional[int]) -> str:
"""Get human-readable description of orientation"""
descriptions = {
1: "Normal (no rotation)",
2: "Mirrored horizontally",
3: "Rotated 180°",
4: "Mirrored vertically",
5: "Mirrored horizontally + rotated 90° CCW",
6: "Rotated 90° CW",
7: "Mirrored horizontally + rotated 90° CW",
8: "Rotated 90° CCW",
None: "Unknown or no EXIF data"
}
return descriptions.get(orientation, "Invalid orientation")
@staticmethod
def correct_image_orientation(image: Image.Image, orientation: Optional[int] = None) -> Image.Image:
"""Correct image orientation based on EXIF data
Args:
image: PIL Image object
orientation: EXIF orientation value (if None, will be detected from image)
Returns:
Corrected PIL Image object
"""
if orientation is None:
# Try to get orientation from image's EXIF data
exif = image.getexif()
if exif is None:
return image
orientation = exif.get(EXIFOrientationHandler.ORIENTATION_TAG)
if orientation is None or orientation == 1:
return image # No correction needed
transforms = EXIFOrientationHandler.ORIENTATION_TRANSFORMS.get(orientation)
if transforms is None:
return image # No known transformation
corrected_image = image.copy()
# Apply transformation(s)
if isinstance(transforms, list):
# Multiple transformations (e.g., mirror + rotate)
for transform in transforms:
corrected_image = corrected_image.transpose(transform)
else:
# Single transformation
corrected_image = corrected_image.transpose(transforms)
return corrected_image
@staticmethod
def correct_image_orientation_from_path(image_path: str) -> Tuple[Image.Image, Optional[int]]:
"""Load image and correct its orientation
Args:
image_path: Path to the image file
Returns:
Tuple of (corrected_image, original_orientation)
"""
try:
with Image.open(image_path) as image:
orientation = EXIFOrientationHandler.get_exif_orientation(image_path)
corrected_image = EXIFOrientationHandler.correct_image_orientation(image, orientation)
return corrected_image, orientation
except Exception:
# If we can't load the image, return None
return None, None
@staticmethod
def save_corrected_image(image: Image.Image, output_path: str, orientation: Optional[int] = None) -> bool:
"""Save image with corrected orientation and remove EXIF orientation tag
Args:
image: Corrected PIL Image object
output_path: Path to save the corrected image
orientation: Original orientation value (for logging)
Returns:
True if successful, False otherwise
"""
try:
# Create a copy to avoid modifying the original
save_image = image.copy()
# Save without EXIF data to avoid orientation issues
save_image.save(output_path, quality=95)
return True
except Exception:
return False
@staticmethod
def get_corrected_image_dimensions(original_width: int, original_height: int, orientation: Optional[int]) -> Tuple[int, int]:
"""Get dimensions after orientation correction
Args:
original_width: Original image width
original_height: Original image height
orientation: EXIF orientation value
Returns:
Tuple of (corrected_width, corrected_height)
"""
if orientation is None or orientation == 1:
return original_width, original_height
# For 90° and 270° rotations, width and height are swapped
if orientation in [5, 6, 7, 8]:
return original_height, original_width
# For 180° rotation and mirrors, dimensions stay the same
return original_width, original_height
@staticmethod
def transform_face_coordinates(face_coords: Dict[str, int], original_width: int, original_height: int,
orientation: Optional[int]) -> Dict[str, int]:
"""Transform face coordinates based on image orientation correction
Args:
face_coords: Face coordinates in DeepFace format {x, y, w, h}
original_width: Original image width
original_height: Original image height
orientation: EXIF orientation value
Returns:
Transformed face coordinates in DeepFace format
"""
if orientation is None or orientation == 1:
return face_coords.copy()
x = face_coords['x']
y = face_coords['y']
w = face_coords['w']
h = face_coords['h']
# Calculate face center and corners
center_x = x + w // 2
center_y = y + h // 2
if orientation == 2: # Mirrored horizontally
new_x = original_width - x - w
new_y = y
elif orientation == 3: # Rotated 180°
new_x = original_width - x - w
new_y = original_height - y - h
elif orientation == 4: # Mirrored vertically
new_x = x
new_y = original_height - y - h
elif orientation == 5: # Mirrored horizontally + rotated 90° CCW
new_x = original_height - y - h
new_y = x
new_w = h
new_h = w
elif orientation == 6: # Rotated 90° CW
new_x = y
new_y = original_width - x - w
new_w = h
new_h = w
elif orientation == 7: # Mirrored horizontally + rotated 90° CW
new_x = y
new_y = x
new_w = h
new_h = w
elif orientation == 8: # Rotated 90° CCW
new_x = original_height - y - h
new_y = original_width - x - w
new_w = h
new_h = w
else:
return face_coords.copy()
return {
'x': new_x,
'y': new_y,
'w': new_w if 'new_w' in locals() else w,
'h': new_h if 'new_h' in locals() else h
}