Rework multithreading

main.py

@@ -2,7 +2,7 @@ import os
 import multiprocessing
 import threading
 from flask import Flask, request, render_template, jsonify
-from timelapse import process_faces
+from timelapse import process_faces, ProcessConfig
 
 app = Flask(__name__)
 
@@ -21,15 +21,36 @@ AVAILABLE_CORES = multiprocessing.cpu_count()
 # Global progress dictionary – only one job at a time is assumed here
 progress_info = {"completed": 0, "total": 0, "status": "idle"}
 
+
 def update_progress(current, total):
+    """
+    Updates the global progress dictionary.
+
+    Args:
+        current (int): Number of completed tasks.
+        total (int): Total number of tasks.
+    """
     progress_info["completed"] = current
     progress_info["total"] = total
     progress_info["status"] = "running" if current < total else "done"
 
+
 def background_process(person_id, padding_percent, resize_size, face_resolution_threshold, pose_threshold, max_workers):
+    """
+    Background process that creates a configuration object and calls process_faces.
+
+    Args:
+        person_id (str): The target person ID.
+        padding_percent (float): Padding percentage for face cropping.
+        resize_size (int): Desired width and height for the aligned face image.
+        face_resolution_threshold (int): Minimum required face resolution.
+        pose_threshold (float): Maximum allowed head pose deviation.
+        max_workers (int): Number of concurrent worker processes.
+    """
     try:
         progress_info["status"] = "running"
-        process_faces(
+        # Build the configuration object for processing
+        config = ProcessConfig(
             api_key=API_KEY,
             base_url=BASE_URL,
             person_id=person_id,
@@ -41,22 +62,29 @@ def background_process(person_id, padding_percent, resize_size, face_resolution_
             min_face_height=face_resolution_threshold,
             pose_threshold=pose_threshold,
             desired_left_eye=LEFT_EYE_POS,
-            max_workers=max_workers,
             face_detect_model_path=FACE_DETECT_MODEL,
-            landmark_model_path=LANDMARK_MODEL,
+            landmark_model_path=LANDMARK_MODEL
-            progress_callback=update_progress
         )
+        process_faces(config, max_workers=max_workers, progress_callback=update_progress)
     except Exception as e:
         progress_info["status"] = f"error: {e}"
     else:
         progress_info["status"] = "done"
 
+
 @app.route("/progress")
 def progress():
+    """
+    Endpoint to return current progress as JSON.
+    """
     return jsonify(progress_info)
 
+
 @app.route("/", methods=["GET", "POST"])
 def index():
+    """
+    Index route that displays the form and starts processing in a background thread on POST.
+    """
     result = None
     error = None
     # Create max_workers_options as a list from 1 to AVAILABLE_CORES
@@ -69,10 +97,12 @@ def index():
             face_resolution_threshold = int(request.form.get("face_resolution_threshold", 128))
             pose_threshold = float(request.form.get("pose_threshold", 25))
             max_workers = int(request.form.get("max_workers", 1))  # default is 1
+
             # Reset progress info before starting
             progress_info["completed"] = 0
             progress_info["total"] = 0
             progress_info["status"] = "idle"
+
             # Start the processing in a background thread
             threading.Thread(
                 target=background_process,
@@ -83,5 +113,6 @@ def index():
             error = f"Error processing request: {e}"
     return render_template("index.html", result=result, error=error, max_workers_options=max_workers_options)
 
+
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=5000)

timelapse.py

@@ -4,6 +4,7 @@ import io
 import requests
 import concurrent.futures
 from datetime import datetime
+from dataclasses import dataclass
 from PIL import Image, ImageOps
 import numpy as np
 import cv2
@@ -11,10 +12,11 @@ import dlib
 from tqdm import tqdm
 import logging
 
-# Custom logging handler that works with tqdm
+
 class TqdmLoggingHandler(logging.Handler):
     def __init__(self, level=logging.NOTSET):
         super().__init__(level)
+
     def emit(self, record):
         try:
             msg = self.format(record)
@@ -22,6 +24,7 @@ class TqdmLoggingHandler(logging.Handler):
         except Exception:
             self.handleError(record)
 
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 tqdm_handler = TqdmLoggingHandler()
@@ -29,7 +32,50 @@ formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s', datef
 tqdm_handler.setFormatter(formatter)
 logger.addHandler(tqdm_handler)
 
+face_detector = None
+face_predictor = None
+
+@dataclass
+class ProcessConfig:
+    """
+    Dataclass to hold configuration parameters for processing assets.
+    """
+    api_key: str
+    base_url: str
+    person_id: str
+    output_folder: str = "output"
+    padding_percent: float = 0.3
+    resize_width: int = 512
+    resize_height: int = 512
+    min_face_width: int = 128
+    min_face_height: int = 128
+    pose_threshold: float = 25
+    desired_left_eye: tuple = (0.35, 0.45)
+    face_detect_model_path: str = "mmod_human_face_detector.dat"
+    landmark_model_path: str = "shape_predictor_68_face_landmarks.dat"
+
+
+def initialize_worker(face_detect_model_path, landmark_model_path):
+    """
+    Initializes the face detector and predictor in each worker process.
+    """
+    global face_detector, face_predictor
+    face_detector = dlib.cnn_face_detection_model_v1(face_detect_model_path)
+    face_predictor = dlib.shape_predictor(landmark_model_path)
+
+
 def get_assets_with_person(api_key, base_url, person_id):
+    """
+    Retrieve all image assets containing the specified person by querying the API.
+
+    Args:
+        api_key (str): API key for authentication.
+        base_url (str): Base URL of the API.
+        person_id (str): ID of the person to search for.
+
+    Returns:
+        list: List of asset dictionaries.
+    """
     headers = {
         'Content-Type': 'application/json',
         'Accept': 'application/json',
@@ -60,17 +106,51 @@ def get_assets_with_person(api_key, base_url, person_id):
         payload["page"] = data['assets'].get('nextPage')
     return all_assets
 
+
 def download_asset(api_key, base_url, asset_id):
+    """
+    Downloads the original image asset from the API.
+
+    Args:
+        api_key (str): API key for authentication.
+        base_url (str): Base URL of the API.
+        asset_id (str): The asset's ID.
+
+    Returns:
+        bytes: The content of the downloaded image.
+    """
     headers = {'x-api-key': api_key}
     response = requests.get(f'{base_url}/assets/{asset_id}/original', headers=headers)
     response.raise_for_status()
     return response.content
 
+
 def format_timestamp(timestamp):
+    """
+    Converts an ISO formatted timestamp to a custom string format.
+
+    Args:
+        timestamp (str): The timestamp string.
+
+    Returns:
+        str: Formatted timestamp.
+    """
     dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
     return dt.strftime("%Y%m%d_%H%M%S")
 
+
 def crop_face_from_metadata(image, face_data, padding_percent):
+    """
+    Crops the face from the image using metadata and applies padding.
+
+    Args:
+        image (PIL.Image): The original image.
+        face_data (dict): Metadata containing face bounding box info.
+        padding_percent (float): Padding as a percentage of face dimensions.
+
+    Returns:
+        PIL.Image: The cropped face image.
+    """
     face_img_width = face_data.get("imageWidth")
     face_img_height = face_data.get("imageHeight")
     img_width, img_height = image.size
@@ -89,23 +169,36 @@ def crop_face_from_metadata(image, face_data, padding_percent):
     new_y2 = min(y2 + padding, img_height)
     return image.crop((new_x1, new_y1, new_x2, new_y2))
 
+
 def get_head_pose(shape, img_size):
+    """
+    Estimates the head pose (pitch, yaw, roll) using facial landmarks.
+
+    Args:
+        shape (dlib.full_object_detection): Detected facial landmarks.
+        img_size (tuple): The size of the image (width, height).
+
+    Returns:
+        tuple or None: (pitch, yaw, roll) in degrees if successful; otherwise None.
+    """
     image_points = np.array([
-        (shape.part(30).x, shape.part(30).y),
+        (shape.part(30).x, shape.part(30).y),  # Nose tip
-        (shape.part(8).x, shape.part(8).y),
+        (shape.part(8).x, shape.part(8).y),  # Chin
-        (shape.part(36).x, shape.part(36).y),
+        (shape.part(36).x, shape.part(36).y),  # Left eye left corner
-        (shape.part(45).x, shape.part(45).y),
+        (shape.part(45).x, shape.part(45).y),  # Right eye right corner
-        (shape.part(48).x, shape.part(48).y),
+        (shape.part(48).x, shape.part(48).y),  # Left Mouth corner
-        (shape.part(54).x, shape.part(54).y)
+        (shape.part(54).x, shape.part(54).y)  # Right mouth corner
     ], dtype="double")
+
     model_points = np.array([
-        (0.0, 0.0, 0.0),
+        (0.0, 0.0, 0.0),  # Nose tip
-        (0.0, -330.0, -65.0),
+        (0.0, -330.0, -65.0),  # Chin
-        (-225.0, 170.0, -135.0),
+        (-225.0, 170.0, -135.0),  # Left eye left corner
-        (225.0, 170.0, -135.0),
+        (225.0, 170.0, -135.0),  # Right eye right corner
-        (-150.0, -150.0, -125.0),
+        (-150.0, -150.0, -125.0),  # Left Mouth corner
-        (150.0, -150.0, -125.0)
+        (150.0, -150.0, -125.0)  # Right mouth corner
     ])
+
     w, h = img_size
     focal_length = w
     center = (w / 2, h / 2)
@@ -118,27 +211,45 @@ def get_head_pose(shape, img_size):
     success, rotation_vector, translation_vector = cv2.solvePnP(
         model_points, image_points, camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE
     )
-    rotation_mat, _ = cv2.Rodrigues(rotation_vector)
+    if not success:
-    proj_matrix = np.hstack((rotation_mat, translation_vector))
+        logger.info("Head pose estimation failed in solvePnP.")
-    _, _, _, _, _, _, eulerAngles = cv2.decomposeProjectionMatrix(proj_matrix)
+        return None
-    pitch, yaw, roll = [float(angle) for angle in eulerAngles]
+    rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
+    proj_matrix = np.hstack((rotation_matrix, translation_vector))
+    _, _, _, _, _, _, euler_angles = cv2.decomposeProjectionMatrix(proj_matrix)
+    pitch, yaw, roll = [float(angle) for angle in euler_angles]
     return pitch, yaw, roll
 
-def align_face(image, predictor, detector, desired_face_width, desired_face_height,
+
-               desired_left_eye, pose_threshold):
+def align_face(image, desired_face_width, desired_face_height, desired_left_eye, pose_threshold):
+    """
+    Aligns the face in the image using facial landmarks and head pose estimation.
+
+    Args:
+        image (PIL.Image): The image containing the face.
+        desired_face_width (int): The desired output face width.
+        desired_face_height (int): The desired output face height.
+        desired_left_eye (tuple): The desired relative position of the left eye.
+        pose_threshold (float): The maximum allowable head pose deviation.
+
+    Returns:
+        PIL.Image or None: The aligned face image if successful, otherwise None.
+    """
     image_np = np.array(image)
     gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
-    detections = detector(gray)
+    detections = face_detector(gray)
     if not detections:
         logger.info("No face detected in the crop. Discarding.")
         return None
-    if hasattr(detections[0], "rect"):
+    # Use the first detection; handle both dlib rectangle and CNN detection type.
-        rect = detections[0].rect
+    detection = detections[0]
-    else:
+    rect = detection.rect if hasattr(detection, "rect") else detection
-        rect = detections[0]
+    shape = face_predictor(gray, rect)
-    shape = predictor(gray, rect)
     img_size = (image_np.shape[1], image_np.shape[0])
-    pitch, yaw, roll = get_head_pose(shape, img_size)
+    head_pose = get_head_pose(shape, img_size)
+    if head_pose is None:
+        return None
+    pitch, yaw, roll = head_pose
     if abs(abs(pitch) - 180) > pose_threshold or abs(yaw) > pose_threshold:
         logger.info(f"Face not frontal enough: pitch={pitch:.2f}, yaw={yaw:.2f}, roll={roll:.2f}. Discarding.")
         return None
@@ -154,9 +265,10 @@ def align_face(image, predictor, detector, desired_face_width, desired_face_heig
     scale = desired_eye_distance / eye_distance
     eyes_center = ((left_eye_center[0] + right_eye_center[0]) / 2.0,
                    (left_eye_center[1] + right_eye_center[1]) / 2.0)
+    # Adjust scale factor if needed
     adjusted_scale = scale * 0.8
     M = cv2.getRotationMatrix2D(eyes_center, angle, adjusted_scale)
-    extra_offset_x = 10
+    extra_offset_x = 10  # Could be parameterized if necessary
     tX = desired_face_width * 0.5 + extra_offset_x
     tY = desired_face_height * desired_left_eye[1]
     M[0, 2] += (tX - eyes_center[0])
@@ -170,23 +282,33 @@ def align_face(image, predictor, detector, desired_face_width, desired_face_heig
     )
     return Image.fromarray(aligned_face_np)
 
-def process_asset_worker(asset, api_key, base_url, person_id, output_folder,
+
-                         padding_percent, min_face_width, min_face_height,
+def process_asset_worker(asset, config: ProcessConfig):
-                         resize_width, resize_height, pose_threshold, desired_left_eye,
+    """
-                         cnn_model_path, predictor_model_path):
+    Worker function to process a single asset.
-    detector = dlib.cnn_face_detection_model_v1(cnn_model_path)
+
-    local_predictor = dlib.shape_predictor(predictor_model_path)
+    This function downloads the asset, crops the face based on metadata,
+    verifies resolution, aligns the face, and then saves the aligned face.
+
+    Args:
+        asset (dict): The asset metadata.
+        config (ProcessConfig): Configuration parameters.
+
+    Returns:
+        str or None: The file path of the saved image if processing is successful; otherwise None.
+    """
     try:
         asset_id = asset['id']
         timestamp = format_timestamp(asset['fileCreatedAt'])
-        image_bytes = download_asset(api_key, base_url, asset_id)
+        image_bytes = download_asset(config.api_key, config.base_url, asset_id)
         image = Image.open(io.BytesIO(image_bytes))
         image = ImageOps.exif_transpose(image)
         image = image.convert("RGB")
     except Exception as e:
         logger.info(f"Error processing asset {asset.get('id')}: {e}")
         return None
-    matching_person = next((p for p in asset.get('people', []) if p.get('id') == person_id), None)
+
+    matching_person = next((p for p in asset.get('people', []) if p.get('id') == config.person_id), None)
     if not matching_person:
         logger.info("Subject not in image.")
         return None
@@ -195,60 +317,67 @@ def process_asset_worker(asset, api_key, base_url, person_id, output_folder,
         logger.info("No face data available.")
         return None
     face_data = faces[0]
-    cropped_face = crop_face_from_metadata(image, face_data, padding_percent)
+    cropped_face = crop_face_from_metadata(image, face_data, config.padding_percent)
     face_width, face_height = cropped_face.size
-    if face_width < min_face_width or face_height < min_face_height:
+    if face_width < config.min_face_width or face_height < config.min_face_height:
         logger.info(f"Face resolution too low ({face_width}x{face_height}).")
         return None
-    aligned_face = align_face(cropped_face, local_predictor, detector,
+    aligned_face = align_face(cropped_face,
-                              desired_face_width=resize_width,
+                              desired_face_width=config.resize_width,
-                              desired_face_height=resize_height,
+                              desired_face_height=config.resize_height,
-                              desired_left_eye=desired_left_eye,
+                              desired_left_eye=config.desired_left_eye,
-                              pose_threshold=pose_threshold)
+                              pose_threshold=config.pose_threshold)
     if aligned_face is None:
         return None
-    filename = os.path.join(output_folder, f"{timestamp}.jpg")
+    os.makedirs(config.output_folder, exist_ok=True)
+    filename = os.path.join(config.output_folder, f"{timestamp}.jpg")
     aligned_face.save(filename)
     return filename
 
-def process_asset_wrapper(asset, process_args):
-    return process_asset_worker(asset, *process_args)
 
-def process_faces(
+def process_asset_wrapper(args):
-    api_key,
+    """
-    base_url,
+    Wrapper to unpack arguments for the worker function.
-    person_id,
+    """
-    output_folder="output",
+    asset, config = args
-    padding_percent=0.3,
+    return process_asset_worker(asset, config)
-    resize_width=512,
+
-    resize_height=512,
+
-    min_face_width=128,
+def process_faces(config: ProcessConfig, max_workers=1, progress_callback=None):
-    min_face_height=128,
+    """
-    pose_threshold=25,
+    Processes assets containing the person and saves aligned face images.
-    desired_left_eye=(0.35, 0.45),
+
-    max_workers=1,
+    This function retrieves assets from the API, then uses a process pool to
-    face_detect_model_path="mmod_human_face_detector.dat",
+    concurrently download, crop, and align faces.
-    landmark_model_path="shape_predictor_68_face_landmarks.dat",
+
-    progress_callback=None  # New optional parameter
+    Args:
-):
+        config (ProcessConfig): Configuration parameters.
-    os.makedirs(output_folder, exist_ok=True)
+        max_workers (int): Number of worker processes.
-    assets = get_assets_with_person(api_key, base_url, person_id)
+        progress_callback (callable, optional): A callback function for progress updates.
+
+    Returns:
+        list: A list of file paths of the saved images.
+    """
+    os.makedirs(config.output_folder, exist_ok=True)
+    assets = get_assets_with_person(config.api_key, config.base_url, config.person_id)
     logger.info(f"Found {len(assets)} assets containing the person.")
     total_assets = len(assets)
     if progress_callback:
         progress_callback(0, total_assets)
-    process_args = (
+
-        api_key, base_url, person_id, output_folder,
+    processed_files = []
-        padding_percent, min_face_width, min_face_height,
+    # Initialize workers with the face detection models
-        resize_width, resize_height, pose_threshold, desired_left_eye,
+    initializer_args = (config.face_detect_model_path, config.landmark_model_path)
-        face_detect_model_path, landmark_model_path
+    with concurrent.futures.ProcessPoolExecutor(
-    )
+            max_workers=max_workers,
-    results = []
+            initializer=initialize_worker,
-    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+            initargs=initializer_args) as executor:
-        for result in tqdm(executor.map(process_asset_wrapper, assets, [process_args]*total_assets), total=total_assets):
+        # Pack arguments for each asset
-            results.append(result)
+        tasks = ((asset, config) for asset in assets)
+        for result in tqdm(executor.map(process_asset_wrapper, tasks), total=total_assets):
+            if result is not None:
+                processed_files.append(result)
             if progress_callback:
-                progress_callback(len(results), total_assets)
+                progress_callback(len(processed_files), total_assets)
-    processed_files = [r for r in results if r is not None]
     logger.info(f"Finished processing. {len(processed_files)} images saved.")
     return processed_files