Add Files

app.py

@@ -4,208 +4,170 @@ import numpy as np
 from PIL import Image
 import torchvision.transforms as transforms
 from ultralytics import YOLO
+import tempfile
 import time
 import os
-import tempfile
-from flask import Flask, request, jsonify
+import json
 import gradio as gr
+from fastapi import FastAPI, UploadFile, File, HTTPException
+import uvicorn
 
-# Initialize Flask app and Gradio interface
-app = Flask(__name__)
+# Initialize FastAPI
+app = FastAPI()
 
-# Global variable to store detection history
-detection_history = []
+# Global variable for face detections
+largest_face_detections = []
 
-# Emotion labels
-emotions = ['Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral']
+# Load models
+yolo_model_path = "yolov8n-face.pt"
+emotion_model_path = "best_emotion_model.pth"
 
-# Load models (cache in Hugging Face Space)
-def load_models():
-    # Face detection model
-    face_model = YOLO('yolov8n-face.pt')
-    
-    # Emotion model (simplified version of your CNN)
-    class EmotionCNN(torch.nn.Module):
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Check if models exist
+if os.path.exists(yolo_model_path):
+    yolo_model = YOLO(yolo_model_path)
+else:
+    raise FileNotFoundError(f"YOLO model not found at {yolo_model_path}")
+
+if os.path.exists(emotion_model_path):
+    from torch import nn
+
+    class EmotionCNN(nn.Module):
         def __init__(self, num_classes=7):
-            super().__init__()
-            self.features = torch.nn.Sequential(
-                torch.nn.Conv2d(1, 64, 3, padding=1),
-                torch.nn.ReLU(),
-                torch.nn.MaxPool2d(2),
-                torch.nn.Conv2d(64, 128, 3, padding=1),
-                torch.nn.ReLU(),
-                torch.nn.MaxPool2d(2),
-                torch.nn.Conv2d(128, 256, 3, padding=1),
-                torch.nn.ReLU(),
-                torch.nn.MaxPool2d(2)
-            )
-            self.classifier = torch.nn.Sequential(
-                torch.nn.Dropout(0.5),
-                torch.nn.Linear(256*6*6, 1024),
-                torch.nn.ReLU(),
-                torch.nn.Dropout(0.5),
-                torch.nn.Linear(1024, num_classes)
-            )
-        
+            super(EmotionCNN, self).__init__()
+            self.conv1 = nn.Sequential(nn.Conv2d(1, 64, kernel_size=3, padding=1),
+                                       nn.BatchNorm2d(64),
+                                       nn.ReLU(),
+                                       nn.MaxPool2d(kernel_size=2, stride=2))
+
+            self.fc = nn.Sequential(nn.Linear(64 * 24 * 24, 1024),
+                                    nn.ReLU(),
+                                    nn.Linear(1024, num_classes))
+
         def forward(self, x):
-            x = self.features(x)
-            x = torch.flatten(x, 1)
-            x = self.classifier(x)
+            x = self.conv1(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
             return x
-    
-    emotion_model = EmotionCNN()
-    # Load your pretrained weights here
-    # emotion_model.load_state_dict(torch.load('emotion_model.pth'))
+
+    emotion_model = EmotionCNN(num_classes=7)
+    checkpoint = torch.load(emotion_model_path, map_location=device)
+    emotion_model.load_state_dict(checkpoint['model_state_dict'])
+    emotion_model.to(device)
     emotion_model.eval()
-    
-    return face_model, emotion_model
+else:
+    raise FileNotFoundError(f"Emotion model not found at {emotion_model_path}")
 
-face_model, emotion_model = load_models()
+# Emotion labels
+emotions = ['Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral']
 
-# Preprocessing function
 def preprocess_face(face_img):
+    """Preprocess face image for emotion detection"""
     transform = transforms.Compose([
         transforms.Resize((48, 48)),
-        transforms.Grayscale(),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5], std=[0.5])
     ])
-    face_pil = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB))
-    return transform(face_pil).unsqueeze(0)
-
-# Process video function
-def process_video(video_path):
-    global detection_history
-    detection_history = []
     
+    face_img = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)).convert('L')
+    face_tensor = transform(face_img).unsqueeze(0)
+    return face_tensor
+
+def process_video(video_path: str):
+    """Process video and return emotion results"""
+    global largest_face_detections
+    largest_face_detections = []
+
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        return {"error": "Could not open video"}
-    
-    frame_count = 0
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    frame_skip = int(fps / 3)  # Process ~3 frames per second
-    
+        return {"success": False, "message": "Could not open video file"}
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
-            
-        frame_count += 1
-        if frame_count % frame_skip != 0:
-            continue
-            
-        # Face detection
-        results = face_model(frame)
-        
+
+        largest_face_area = 0
+        current_detection = None
+
+        results = yolo_model(frame, stream=True)
         for result in results:
             boxes = result.boxes
-            if len(boxes) == 0:
-                continue
-                
             for box in boxes:
-                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+                x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
                 face_img = frame[y1:y2, x1:x2]
                 
                 if face_img.size == 0:
                     continue
-                    
-                # Emotion prediction
-                face_tensor = preprocess_face(face_img)
+                
+                face_tensor = preprocess_face(face_img).to(device)
+
                 with torch.no_grad():
                     output = emotion_model(face_tensor)
-                    prob = torch.nn.functional.softmax(output, dim=1)[0]
-                    pred_idx = torch.argmax(output).item()
-                    confidence = prob[pred_idx].item()
-                
-                detection_history.append({
-                    "frame": frame_count,
-                    "time": frame_count / fps,
-                    "emotion": emotions[pred_idx],
-                    "confidence": confidence,
-                    "box": [x1, y1, x2, y2]
-                })
-    
+                    probabilities = torch.nn.functional.softmax(output, dim=1)
+                    emotion_idx = torch.argmax(output, dim=1).item()
+                    confidence = probabilities[0][emotion_idx].item()
+
+                emotion = emotions[emotion_idx]
+
+                if (x2 - x1) * (y2 - y1) > largest_face_area:
+                    largest_face_area = (x2 - x1) * (y2 - y1)
+                    current_detection = {"emotion": emotion, "confidence": confidence}
+
+        if current_detection:
+            largest_face_detections.append(current_detection)
+
     cap.release()
     
-    if not detection_history:
-        return {"error": "No faces detected"}
-    
+    if not largest_face_detections:
+        return {"success": True, "message": "No faces detected", "results": []}
+
     return {
-        "detections": detection_history,
-        "summary": {
-            "total_frames": frame_count,
-            "fps": fps,
-            "duration": frame_count / fps
-        }
+        "success": True,
+        "message": "Video processed",
+        "results": largest_face_detections
     }
 
-# Flask API endpoint
-@app.route('/api/predict', methods=['POST'])
-def api_predict():
-    if 'file' not in request.files:
-        return jsonify({"error": "No file provided"}), 400
-        
-    file = request.files['file']
-    if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
-    
-    # Save to temp file
-    temp_path = os.path.join(tempfile.gettempdir(), file.filename)
-    file.save(temp_path)
-    
-    # Process video
-    result = process_video(temp_path)
-    
-    # Clean up
-    os.remove(temp_path)
-    
-    return jsonify(result)
+@app.post("/api/video")
+async def handle_video(file: UploadFile = File(...)):
+    """API endpoint for video emotion detection"""
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
+            tmp.write(await file.read())
+            video_path = tmp.name
 
-# Gradio interface
-def gradio_predict(video):
-    temp_path = os.path.join(tempfile.gettempdir(), video.name)
-    with open(temp_path, 'wb') as f:
-        f.write(video.read())
-    
-    result = process_video(temp_path)
-    os.remove(temp_path)
-    
-    if "error" in result:
-        return result["error"]
-    
-    # Create visualization
-    cap = cv2.VideoCapture(video.name)
-    ret, frame = cap.read()
-    cap.release()
+        result = process_video(video_path)
+        os.remove(video_path)
+        return result
+
+    except Exception as e:
+        return {"success": False, "message": "Error processing video", "error": str(e)}
+
+# Gradio UI
+def gradio_process(video):
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
+        tmp.write(video)
+        video_path = tmp.name
+
+    result = process_video(video_path)
+    os.remove(video_path)
+    return result
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Video Emotion Analysis")
     
-    if ret:
-        # Draw last detection on frame
-        last_det = result["detections"][-1]
-        x1, y1, x2, y2 = last_det["box"]
-        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        cv2.putText(frame, f"{last_det['emotion']} ({last_det['confidence']:.2f})", 
-                   (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.File(label="Upload a video", file_types=[".mp4"])
+            submit_btn = gr.Button("Analyze")
         
-        # Convert to RGB for Gradio
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        return frame, result
-    return result
+        with gr.Column():
+            output = gr.JSON(label="Results")
+    
+    submit_btn.click(fn=gradio_process, inputs=video_input, outputs=output)
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=gradio_predict,
-    inputs=gr.Video(label="Upload Video"),
-    outputs=[
-        gr.Image(label="Detection Preview"), 
-        gr.JSON(label="Results")
-    ],
-    title="Video Emotion Detection",
-    description="Upload a video to detect emotions in faces"
-)
-
-# Mount Gradio app
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,8 +1,8 @@
+ultralytics
 torch
 torchvision
-opencv-python
-ultralytics
 gradio
-flask
-numpy
-Pillow
+fastapi
+uvicorn
+opencv-python
+pillow