First commit

app.py

@@ -41,18 +41,57 @@ from models.pcs400   import cal_pcs
 from ultralytics import YOLO
 import supervision as sv
 
+import gradio as gr
+import cv2
+import os
+import tempfile
+from ultralytics import YOLO
+from moviepy.editor import ImageSequenceClip
 
+# Load face detector
+model = YOLO("yolov8n-face.pt").cuda()  # assumes CUDA available
 
-def dummy_fn(video):
-    return video
-
-model = YOLO("yolov8n.pt")  # or a face-specific checkpoint
-results = model.predict("input_frame.jpg")[0]
-
-# Filter only face (class 0 in COCO = 'person'; need a face-detection specific model if you want more precision)
-faces = [b for b in results.boxes if b.cls == 0]
+@spaces.GPU
+def extract_faces(video_file):
+    cap = cv2.VideoCapture(video_file)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frames = []
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Inference
+        results = model(frame, verbose=False)[0]
+        for box in results.boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            face_crop = frame[y1:y2, x1:x2]
+            if face_crop.size != 0:
+                resized = cv2.resize(face_crop, (224, 224))
+                frames.append(resized)
+            break  # only one face per frame
+
+    cap.release()
+
+    # Save as video
+    tmpdir = tempfile.mkdtemp()
+    output_path = os.path.join(tmpdir, "face_only_video.mp4")
+    clip = ImageSequenceClip([cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in frames], fps=fps)
+    clip.write_videofile(output_path, codec="libx264", audio=False, verbose=False, logger=None)
+
+    return output_path
+
+iface = gr.Interface(
+    fn=extract_faces,
+    inputs=gr.Video(source="upload", label="Upload your video or record"),
+    outputs=gr.Video(label="Detected Face Only Video"),
+    title="Face Detector",
+    description="Upload or record a video. We'll crop face regions and return a face-only video."
+)
+
+iface.launch()
 
-gr.Interface(fn=dummy_fn, inputs=gr.Video(source="webcam"), outputs="video").launch()
 
 
 ckpt = "ckpts/SEMamba_advanced.pth"