fix: improve AI tool reliability for face detection and background removal (#25)

- Replace OpenCV Haar Cascades with MediaPipe for face detection, using
  short-range model first with full-range fallback for better accuracy
- Add auto-orient to remove-background route for EXIF-rotated photos
- Change default background removal model from u2net to birefnet-general-lite
- Fix flaky test by setting SQLite busy_timeout before journal_mode pragma

Co-authored-by: Siddharth Kumar Sah <siddharth123sk@gmail.com>

apps/api/src/db/index.ts

@@ -10,9 +10,11 @@ mkdirSync(dirname(env.DB_PATH), { recursive: true });
 
 const sqlite: DatabaseType = new Database(env.DB_PATH);
 
-// Critical SQLite pragmas for reliability
-sqlite.pragma("journal_mode = WAL");
+// Critical SQLite pragmas for reliability.
+// busy_timeout must be set first so journal_mode = WAL can retry
+// if another connection holds the lock (e.g. parallel test files).
 sqlite.pragma("busy_timeout = 5000");
+sqlite.pragma("journal_mode = WAL");
 sqlite.pragma("synchronous = NORMAL");
 sqlite.pragma("foreign_keys = ON");
 

apps/api/src/routes/tools/blur-faces.ts

@@ -10,10 +10,7 @@ import { createWorkspace } from "../../lib/workspace.js";
 import { updateSingleFileProgress } from "../progress.js";
 import { registerToolProcessFn } from "../tool-factory.js";
 
-/**
- * Face detection and blurring route.
- * Uses MediaPipe for detection, PIL for blurring.
- */
+/** Face detection and blurring route. */
 export function registerBlurFaces(app: FastifyInstance) {
   app.post("/api/v1/tools/blur-faces", async (request: FastifyRequest, reply: FastifyReply) => {
     let fileBuffer: Buffer | null = null;

apps/api/src/routes/tools/remove-background.ts

@@ -4,6 +4,7 @@ import { basename, join } from "node:path";
 import { removeBackground } from "@stirling-image/ai";
 import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
 import { z } from "zod";
+import { autoOrient } from "../../lib/auto-orient.js";
 import { validateImageBuffer } from "../../lib/file-validation.js";
 import { createWorkspace } from "../../lib/workspace.js";
 import { updateSingleFileProgress } from "../progress.js";
@@ -56,6 +57,10 @@ export function registerRemoveBackground(app: FastifyInstance) {
 
       try {
         const settings = settingsRaw ? JSON.parse(settingsRaw) : {};
+
+        // Auto-orient to fix EXIF rotation before processing
+        fileBuffer = await autoOrient(fileBuffer);
+
         request.log.info(
           { toolId: "remove-background", imageSize: fileBuffer.length, model: settings.model },
           "Starting background removal",
@@ -126,9 +131,10 @@ export function registerRemoveBackground(app: FastifyInstance) {
     }),
     process: async (inputBuffer, settings, filename) => {
       const s = settings as { model?: string; backgroundColor?: string };
+      const orientedBuffer = await autoOrient(inputBuffer);
       const jobId = randomUUID();
       const workspacePath = await createWorkspace(jobId);
-      const resultBuffer = await removeBackground(inputBuffer, join(workspacePath, "output"), {
+      const resultBuffer = await removeBackground(orientedBuffer, join(workspacePath, "output"), {
         model: s.model,
         backgroundColor: s.backgroundColor,
       });

packages/ai/python/detect_faces.py

@@ -1,4 +1,4 @@
-"""Face detection and blurring using OpenCV."""
+"""Face detection and blurring using MediaPipe."""
 import sys
 import json
 
@@ -23,37 +23,47 @@ def main():
         img = Image.open(input_path).convert("RGB")
 
         try:
-            import cv2
+            import mediapipe as mp
             import numpy as np
 
             emit_progress(20, "Ready")
 
-            # Load Haar cascade for face detection
-            haar_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
-            face_cascade = cv2.CascadeClassifier(haar_path)
+            # Map sensitivity (0.1-0.9) to MediaPipe confidence threshold.
+            # Higher sensitivity = lower confidence threshold = more detections.
+            min_confidence = max(0.1, 1.0 - sensitivity)
 
-            # Convert to grayscale for detection
             img_array = np.array(img)
-            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
-
-            # Map sensitivity (0.1-0.9) to minNeighbors (8-2)
-            # Higher sensitivity = fewer required neighbors = more detections
-            min_neighbors = max(2, int(8 - sensitivity * 7))
+            mp_face = mp.solutions.face_detection
 
+            # Try short-range model first (model_selection=0, best for faces
+            # within ~2m which covers most photos), then fall back to
+            # full-range model (model_selection=1) for distant/group shots.
             emit_progress(25, "Scanning for faces")
-            faces_detected = face_cascade.detectMultiScale(
-                gray,
-                scaleFactor=1.1,
-                minNeighbors=min_neighbors,
-                minSize=(30, 30),
-            )
+            results = None
+            for model_sel in [0, 1]:
+                detector = mp_face.FaceDetection(
+                    model_selection=model_sel,
+                    min_detection_confidence=min_confidence,
+                )
+                results = detector.process(img_array)
+                detector.close()
+                if results.detections:
+                    break
 
             faces = []
-            num_faces = len(faces_detected)
+            detections = results.detections or []
+            num_faces = len(detections)
             emit_progress(50, f"Found {num_faces} face{'s' if num_faces != 1 else ''}")
 
             if num_faces > 0:
-                for i, (x, y, w, h) in enumerate(faces_detected):
+                ih, iw = img_array.shape[:2]
+                for i, detection in enumerate(detections):
+                    bbox = detection.location_data.relative_bounding_box
+                    x = int(bbox.xmin * iw)
+                    y = int(bbox.ymin * ih)
+                    w = int(bbox.width * iw)
+                    h = int(bbox.height * ih)
+
                     # Add padding around the face
                     pad = int(max(w, h) * 0.1)
                     x1 = max(0, x - pad)
@@ -66,7 +76,7 @@ def main():
                         ImageFilter.GaussianBlur(blur_radius)
                     )
                     img.paste(blurred, (x1, y1))
-                    faces.append({"x": int(x), "y": int(y), "w": int(w), "h": int(h)})
+                    faces.append({"x": x, "y": y, "w": w, "h": h})
                     emit_progress(
                         50 + int((i + 1) / num_faces * 40),
                         f"Blurring face {i + 1} of {num_faces}",
@@ -89,7 +99,7 @@ def main():
                 json.dumps(
                     {
                         "success": False,
-                        "error": "Face detection requires OpenCV. Install with: pip install opencv-python-headless",
+                        "error": "Face detection requires MediaPipe. Install with: pip install mediapipe",
                     }
                 )
             )

packages/ai/python/dispatcher.py

@@ -37,7 +37,7 @@ def _try_import(name, import_fn):
 
 
 _try_import("PIL", lambda: __import__("PIL"))
-_try_import("cv2", lambda: __import__("cv2"))
+_try_import("mediapipe", lambda: __import__("mediapipe"))
 _try_import("numpy", lambda: __import__("numpy"))
 _try_import("gpu", lambda: __import__("gpu"))
 

packages/ai/python/remove_bg.py

@@ -14,7 +14,7 @@ def main():
     output_path = sys.argv[2]
     settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
 
-    model = settings.get("model", "u2net")
+    model = settings.get("model", "birefnet-general-lite")
     bg_color = settings.get("backgroundColor", "")
 
     # Redirect stdout to stderr so library download/progress output