initial add files commit

app.py

@@ -0,0 +1,210 @@
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse
+from ultralytics import YOLO
+from PIL import Image
+import cv2
+import mediapipe as mp
+import numpy as np
+import io
+import tempfile
+import os
+from fastapi.middleware.cors import CORSMiddleware
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, specify your allowed domains
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Health Check Endpoint
+@app.get("/")
+def root():
+    """
+    Health check endpoint to verify that the API is running.
+    """
+    return {"message": "YOLO11 Emotion Detection API is live!"}
+
+# Load your custom YOLO emotion detection model
+try:
+    emotion_model = YOLO("model/yolo11m_affectnet_best.pt")
+except Exception as e:
+    raise Exception(f"Error loading emotion model: {e}")
+
+def detect_emotions(image):
+    """
+    Given an OpenCV BGR image, detect faces using Mediapipe and perform emotion detection
+    on each face crop using the YOLO model.
+    Returns a list of detections with bounding box and emotion details.
+    """
+    height, width, _ = image.shape
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    detections_list = []
+    
+    mp_face_detection = mp.solutions.face_detection
+    with mp_face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5) as face_detection:
+        results = face_detection.process(image_rgb)
+    
+    if results.detections:
+        for detection in results.detections:
+            bbox = detection.location_data.relative_bounding_box
+            x_min = int(bbox.xmin * width)
+            y_min = int(bbox.ymin * height)
+            box_width = int(bbox.width * width)
+            box_height = int(bbox.height * height)
+            x_max = x_min + box_width
+            y_max = y_min + box_height
+
+            x_min = max(0, x_min)
+            y_min = max(0, y_min)
+            x_max = min(width, x_max)
+            y_max = min(height, y_max)
+
+            face_crop = image[y_min:y_max, x_min:x_max]
+            if face_crop.size == 0:
+                continue
+            face_crop_rgb = cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB)
+            face_pil = Image.fromarray(face_crop_rgb)
+
+            emotion_results = emotion_model.predict(source=face_pil, conf=0.5)
+            if len(emotion_results) > 0 and len(emotion_results[0].boxes) > 0:
+                box_detect = emotion_results[0].boxes[0]
+                emotion_label = emotion_results[0].names[int(box_detect.cls)]
+                confidence = float(box_detect.conf)
+            else:
+                emotion_label = "N/A"
+                confidence = 0.0
+
+            detection_info = {
+                "bbox": {
+                    "x_min": x_min,
+                    "y_min": y_min,
+                    "x_max": x_max,
+                    "y_max": y_max
+                },
+                "emotion": emotion_label,
+                "confidence": confidence
+            }
+            detections_list.append(detection_info)
+    return detections_list
+
+@app.post("/predict_frame")
+async def predict_frame(file: UploadFile = File(...)):
+    """
+    Accept an image file, run face and emotion detection, annotate the image with bounding boxes 
+    and emotion labels, and return the annotated image as PNG.
+    """
+    if not file.filename.lower().endswith(('.jpg', '.jpeg', '.png')):
+        raise HTTPException(status_code=400, detail="Invalid file format. Only JPG, JPEG, and PNG are allowed.")
+    try:
+        contents = await file.read()
+        nparr = np.frombuffer(contents, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        if image is None:
+            raise HTTPException(status_code=400, detail="Invalid image file.")
+        
+        detections = detect_emotions(image)
+        
+        for det in detections:
+            bbox = det["bbox"]
+            label_text = f'{det["emotion"]} ({det["confidence"]:.2f})'
+            cv2.rectangle(image, (bbox["x_min"], bbox["y_min"]), (bbox["x_max"], bbox["y_max"]), (0, 255, 0), 2)
+            cv2.putText(image, label_text, (bbox["x_min"], bbox["y_min"] - 10),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
+        
+        is_success, im_buf_arr = cv2.imencode(".png", image)
+        if not is_success:
+            raise HTTPException(status_code=500, detail="Error encoding image.")
+        byte_im = im_buf_arr.tobytes()
+        buf = io.BytesIO(byte_im)
+        buf.seek(0)
+        return StreamingResponse(buf, media_type="image/png")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+    finally:
+        await file.close()
+
+@app.post("/predict_emotion")
+async def predict_emotion(file: UploadFile = File(...)):
+    """
+    Accept an image file, run face and emotion detection, and return the results as JSON.
+    The JSON response includes a list of detections with bounding box coordinates, emotion label, and confidence score.
+    """
+    if not file.filename.lower().endswith(('.jpg', '.jpeg', '.png')):
+        raise HTTPException(status_code=400, detail="Invalid file format. Only JPG, JPEG, and PNG are allowed.")
+    try:
+        contents = await file.read()
+        nparr = np.frombuffer(contents, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        if image is None:
+            raise HTTPException(status_code=400, detail="Invalid image file.")
+        detections = detect_emotions(image)
+        return {"detections": detections}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+    finally:
+        await file.close()
+
+@app.post("/predict_video")
+async def predict_video(file: UploadFile = File(...)):
+    """
+    Accept a video file, process it frame-by-frame with face and emotion detection,
+    annotate each frame with bounding boxes and emotion labels, and return the annotated
+    video as an MP4 file.
+    """
+    if not file.filename.lower().endswith(('.mp4', '.avi', '.mov')):
+        raise HTTPException(status_code=400, detail="Invalid file format. Only MP4, AVI, and MOV are allowed.")
+    try:
+        contents = await file.read()
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_input:
+            tmp_input.write(contents)
+            input_video_path = tmp_input.name
+        
+        cap = cv2.VideoCapture(input_video_path)
+        if not cap.isOpened():
+            raise HTTPException(status_code=400, detail="Could not open video file.")
+        
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_output:
+            output_video_path = tmp_output.name
+        
+        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+        
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            
+            detections = detect_emotions(frame)
+            for det in detections:
+                bbox = det["bbox"]
+                label_text = f'{det["emotion"]} ({det["confidence"]:.2f})'
+                cv2.rectangle(frame, (bbox["x_min"], bbox["y_min"]), (bbox["x_max"], bbox["y_max"]), (0, 255, 0), 2)
+                cv2.putText(frame, label_text, (bbox["x_min"], bbox["y_min"] - 10),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
+            out.write(frame)
+        
+        cap.release()
+        out.release()
+        
+        with open(output_video_path, "rb") as f:
+            video_bytes = f.read()
+        
+        os.remove(input_video_path)
+        os.remove(output_video_path)
+        
+        return StreamingResponse(io.BytesIO(video_bytes), media_type="video/mp4")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing video: {str(e)}")
+    finally:
+        await file.close()
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

model/yolo11m_affectnet_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d9103ead25cfa9307e1580c95d19d744251b3cbc03b8945a7c74b239309d105
+size 40481573

requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn
+python-multipart
+Pillow
+opencv-python
+mediapipe
+numpy
+ultralytics