Update app/main.py

app/main.py

@@ -1,58 +1,70 @@
+# app/main.py
 import os
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # silence TF/MediaPipe logs
-
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # quiet TF/MediaPipe logs
 from absl import logging as absl_logging
 absl_logging.set_verbosity(absl_logging.ERROR)
 
-from fastapi import FastAPI, Request, UploadFile, File, Form
-from fastapi.responses import HTMLResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-
 import uuid
+from pathlib import Path
+
 import cv2
 import numpy as np
 import mediapipe as mp
-from pathlib import Path
+from fastapi import FastAPI, Request, UploadFile, File
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
 
-# ------------ Paths ------------
+# ---------------- Paths & FastAPI ----------------
 BASE_DIR = Path(__file__).resolve().parent
 TEMPLATES_DIR = BASE_DIR / "templates"
 STATIC_DIR = BASE_DIR / "static"
-PROCESSED_DIR = STATIC_DIR / "processed"
-PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
 
-# ------------ FastAPI -----------
+# Writable location on Hugging Face Spaces
+RESULTS_DIR = Path("/tmp/faceblur")
+RESULTS_DIR.mkdir(parents=True, exist_ok=True)
+
 app = FastAPI(title="Face Blur Web")
-app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
+app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
+app.mount("/results", StaticFiles(directory=str(RESULTS_DIR)), name="results")
 templates = Jinja2Templates(directory=str(TEMPLATES_DIR))
 
-# ------------ Detection/Blur settings ------------
-SCALES   = [1.0, 1.5, 2.0, 2.5, 3.0]   # multi-scale improves small-face recall
+# --------------- Detection/Blur settings ----------------
+# Multi-scale + both MediaPipe models helps with small/far faces
+SCALES   = [1.0, 1.5, 2.0, 2.5, 3.0]
 MIN_CONF = 0.25
 IOU_NMS  = 0.35
 
-OVAL_SCALE_X = 1.35     # widen ellipse vs. face box
-OVAL_SCALE_Y = 1.55     # height scale
-FEATHER_FRAC = 0.12     # feather edge ~ % of face size
+# Oval size around each detected face (relative to bbox)
+OVAL_SCALE_X = 1.35
+OVAL_SCALE_Y = 1.55
 
-KERNEL_FRAC = 0.9       # blur kernel ~ % of face size (strong)
-MIN_KERNEL  = 55        # lower bound kernel (odd)
-MAX_KERNEL  = 301       # upper bound kernel (odd)
-MAX_SIDE    = 1800      # optional resize for very large uploads (performance)
+# Edge feathering (as fraction of face size)
+FEATHER_FRAC = 0.12
 
-# ------------ Utils ------------
+# Blur strength scales with face size
+KERNEL_FRAC = 0.9
+MIN_KERNEL  = 55
+MAX_KERNEL  = 301
+
+# Downscale very large uploads for speed
+MAX_SIDE = 1800
+
+# ----------------- Utils -----------------
 def odd(n: int) -> int:
+    """Ensure odd kernel size >= 3."""
     n = max(3, int(n))
     return n if n % 2 == 1 else n + 1
 
 def nms(boxes, scores, iou_thresh=0.35):
+    """Non-max suppression for [x,y,w,h] boxes."""
     if not boxes:
         return []
     idxs = np.argsort(scores)[::-1]
     keep = []
     while len(idxs) > 0:
-        i = idxs[0]; keep.append(i)
+        i = idxs[0]
+        keep.append(i)
         xx1 = np.maximum(boxes[i][0], np.array([boxes[j][0] for j in idxs[1:]]))
         yy1 = np.maximum(boxes[i][1], np.array([boxes[j][1] for j in idxs[1:]]))
         xx2 = np.minimum(boxes[i][0]+boxes[i][2], np.array([boxes[j][0]+boxes[j][2] for j in idxs[1:]]))
@@ -65,8 +77,15 @@ def nms(boxes, scores, iou_thresh=0.35):
         idxs = idxs[1:][iou <= iou_thresh]
     return keep
 
+def gentle_contrast_boost(img_bgr: np.ndarray) -> np.ndarray:
+    """CLAHE on L channel of LAB for small contrast lift (helps detection)."""
+    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
+    l, a, b = cv2.split(lab)
+    l = cv2.createCLAHE(2.0, (8, 8)).apply(l)
+    return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR)
+
 def detect_faces_mediapipe(img_bgr: np.ndarray):
-    """Detect faces across multiple scales and both MediaPipe models; return boxes [x,y,w,h]."""
+    """Detect faces across multiple scales and both MediaPipe models."""
     H, W = img_bgr.shape[:2]
     mp_fd = mp.solutions.face_detection
     detectors = [
@@ -85,7 +104,6 @@ def detect_faces_mediapipe(img_bgr: np.ndarray):
                 bb = d.location_data.relative_bounding_box
                 x, y = int(bb.xmin * Uw), int(bb.ymin * Uh)
                 w, h = int(bb.width * Uw), int(bb.height * Uh)
-                # map back to original coordinates
                 x, y, w, h = int(x/s), int(y/s), int(w/s), int(h/s)
                 x = max(0, x); y = max(0, y)
                 w = max(1, min(w, W - x)); h = max(1, min(h, H - y))
@@ -94,46 +112,39 @@ def detect_faces_mediapipe(img_bgr: np.ndarray):
     keep = nms(boxes, scores, iou_thresh=IOU_NMS)
     return [boxes[i] for i in keep]
 
-def gentle_contrast_boost(img_bgr: np.ndarray) -> np.ndarray:
-    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
-    l, a, b = cv2.split(lab)
-    l = cv2.createCLAHE(2.0, (8, 8)).apply(l)
-    return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR)
-
 def draw_unblurred_preview(img_bgr: np.ndarray, boxes):
-    preview = img_bgr.copy()
+    """Draw green ovals (no blur)."""
+    out = img_bgr.copy()
     for (x, y, w, h) in boxes:
         center = (x + w // 2, y + h // 2)
         axes   = (int((w/2) * OVAL_SCALE_X), int((h/2) * OVAL_SCALE_Y))
-        cv2.ellipse(preview, center, axes, 0, 0, 360, (0, 255, 0), 3)
-    return preview
+        cv2.ellipse(out, center, axes, 0, 0, 360, (0, 255, 0), 3)
+    return out
 
 def blur_faces_oval(img_bgr: np.ndarray, boxes):
+    """Strong oval blur with feathered edges; kernel scales per face size."""
     H, W = img_bgr.shape[:2]
     out = img_bgr.copy()
     for (x, y, w, h) in boxes:
-        # ellipse parameters
         cx, cy = x + w // 2, y + h // 2
         ax = int((w / 2) * OVAL_SCALE_X)
         ay = int((h / 2) * OVAL_SCALE_Y)
 
-        # rectangular ROI around ellipse
         x0 = max(0, cx - ax); y0 = max(0, cy - ay)
         x1 = min(W, cx + ax); y1 = min(H, cy + ay)
         roi = out[y0:y1, x0:x1]
         rh, rw = roi.shape[:2]
 
-        # kernel proportional to face size (clamped)
         k_base = int(KERNEL_FRAC * max(w, h))
+        # clamp to ROI and ensure odd
         k = odd(max(MIN_KERNEL, min(MAX_KERNEL, k_base, rh - (rh % 2 == 0), rw - (rw % 2 == 0))))
         if k < 9:
-            # fall back to pixelation if ROI is tiny
+            # fallback to pixelation if ROI too small
             small = cv2.resize(roi, (max(1, rw // 10), max(1, rh // 10)), interpolation=cv2.INTER_LINEAR)
             roi_blur = cv2.resize(small, (rw, rh), interpolation=cv2.INTER_NEAREST)
         else:
             roi_blur = cv2.GaussianBlur(roi, (k, k), 0)
 
-        # oval mask + feather
         mask = np.zeros((rh, rw), dtype=np.uint8)
         cv2.ellipse(mask, (rw // 2, rh // 2), (ax, ay), 0, 0, 360, 255, -1)
         feather = odd(int(max(w, h) * FEATHER_FRAC))
@@ -145,68 +156,57 @@ def blur_faces_oval(img_bgr: np.ndarray, boxes):
         out[y0:y1, x0:x1] = roi_out
     return out
 
-# ------------ Routes ------------
+# ----------------- Routes -----------------
 @app.get("/", response_class=HTMLResponse)
 async def index(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request, "result": None})
+    return templates.TemplateResponse("index.html", {"request": request, "result": None, "error": None})
 
 @app.post("/upload", response_class=HTMLResponse)
-async def upload_image(
-    request: Request,
-    file: UploadFile = File(...),
-):
-    # Validate extension
-    name = file.filename or "upload"
+async def upload_image(request: Request, file: UploadFile = File(...)):
+    name = (file.filename or "upload").strip()
     if not name.lower().endswith((".jpg", ".jpeg", ".png", ".bmp", ".webp")):
         return templates.TemplateResponse("index.html", {
-            "request": request,
-            "error": "Unsupported file type. Please upload JPG/PNG/BMP/WEBP.",
-            "result": None
+            "request": request, "error": "Unsupported file type. Use JPG/PNG/BMP/WEBP.", "result": None
         })
 
-    # Read file into OpenCV image
     data = await file.read()
     npbuf = np.frombuffer(data, np.uint8)
     img = cv2.imdecode(npbuf, cv2.IMREAD_COLOR)
     if img is None:
         return templates.TemplateResponse("index.html", {
-            "request": request,
-            "error": "Could not decode image. Try another file.",
-            "result": None
+            "request": request, "error": "Could not decode image.", "result": None
         })
 
-    # Optional: downscale very large images for performance
+    # Optional downscale for very large images
     H, W = img.shape[:2]
     scale = min(1.0, float(MAX_SIDE) / max(H, W))
     if scale < 1.0:
         img = cv2.resize(img, (int(W * scale), int(H * scale)), interpolation=cv2.INTER_AREA)
 
-    # Small contrast boost helps detection
+    # Detection pipeline
     img_proc = gentle_contrast_boost(img)
-
-    # Detect faces
     boxes = detect_faces_mediapipe(img_proc)
     faces_count = len(boxes)
 
-    # Generate outputs
     preview = draw_unblurred_preview(img_proc, boxes)
     blurred  = blur_faces_oval(img_proc, boxes)
 
-    # Save results with unique IDs
+    # Save to /tmp and serve via /results
     uid = uuid.uuid4().hex
-    annot_rel = f"static/processed/{uid}_annot.jpg"
-    blur_rel  = f"static/processed/{uid}_blur.jpg"
+    annot_name = f"{uid}_annot.jpg"
+    blur_name  = f"{uid}_blur.jpg"
 
-    cv2.imwrite(str(BASE_DIR / annot_rel), preview)
-    cv2.imwrite(str(BASE_DIR / blur_rel),  blurred)
+    cv2.imwrite(str(RESULTS_DIR / annot_name), preview)
+    cv2.imwrite(str(RESULTS_DIR / blur_name),  blurred)
 
     return templates.TemplateResponse("index.html", {
         "request": request,
         "error": None,
         "result": {
             "faces": faces_count,
-            "annot_url": f"/{annot_rel}",
-            "blur_url":  f"/{blur_rel}",
+            "annot_url": f"/results/{annot_name}",
+            "blur_url":  f"/results/{blur_name}",
             "filename": name
         }
     })
+