Rivalcoder/Video-Processing
0
1import cv22import torch3import numpy as np4from PIL import Image5import torchvision.transforms as transforms6import time7import os8import json9from typing import Dict, List, Any10from fastapi import FastAPI, UploadFile, File, HTTPException11from fastapi.responses import JSONResponse, HTMLResponse12import uuid13from pathlib import Path14import gradio as gr15import tempfile16 17app = FastAPI()18 19# Global variable to store the history of largest face detections20largest_face_detections = []21 22# EmotionCNN model definition23class EmotionCNN(torch.nn.Module):24 def __init__(self, num_classes=7):25 super(EmotionCNN, self).__init__()26 27 # First convolutional block28 self.conv1 = torch.nn.Sequential(29 torch.nn.Conv2d(1, 64, kernel_size=3, padding=1),30 torch.nn.BatchNorm2d(64),31 torch.nn.ReLU(),32 torch.nn.MaxPool2d(kernel_size=2, stride=2)33 )34 35 # Second convolutional block36 self.conv2 = torch.nn.Sequential(37 torch.nn.Conv2d(64, 128, kernel_size=3, padding=1),38 torch.nn.BatchNorm2d(128),39 torch.nn.ReLU(),40 torch.nn.MaxPool2d(kernel_size=2, stride=2)41 )42 43 # Third convolutional block44 self.conv3 = torch.nn.Sequential(45 torch.nn.Conv2d(128, 256, kernel_size=3, padding=1),46 torch.nn.BatchNorm2d(256),47 torch.nn.ReLU(),48 torch.nn.MaxPool2d(kernel_size=2, stride=2)49 )50 51 # Fourth convolutional block52 self.conv4 = torch.nn.Sequential(53 torch.nn.Conv2d(256, 512, kernel_size=3, padding=1),54 torch.nn.BatchNorm2d(512),55 torch.nn.ReLU(),56 torch.nn.MaxPool2d(kernel_size=2, stride=2)57 )58 59 # Fifth convolutional block with residual connection60 self.conv5 = torch.nn.Sequential(61 torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),62 torch.nn.BatchNorm2d(512),63 torch.nn.ReLU()64 )65 66 # Attention mechanism67 self.attention = torch.nn.Sequential(68 torch.nn.Conv2d(512, 1, kernel_size=1),69 torch.nn.Sigmoid()70 )71 72 # Fully connected layers73 self.fc = torch.nn.Sequential(74 torch.nn.Dropout(0.5),75 torch.nn.Linear(512 * 3 * 3, 1024),76 torch.nn.ReLU(),77 torch.nn.Dropout(0.5),78 torch.nn.Linear(1024, 512),79 torch.nn.ReLU(),80 torch.nn.Dropout(0.3),81 torch.nn.Linear(512, num_classes)82 )83 84 def forward(self, x):85 x = self.conv1(x)86 x = self.conv2(x)87 x = self.conv3(x)88 x = self.conv4(x)89 90 # Fifth conv block with residual connection91 x_res = x92 x = self.conv5(x)93 x = x + x_res94 95 # Apply attention96 attn = self.attention(x)97 x = x * attn98 99 # Flatten100 x = x.view(x.size(0), -1)101 102 # Fully connected103 x = self.fc(x)104 return x105 106def load_emotion_model(model_path, device='cuda' if torch.cuda.is_available() else 'cpu'):107 """Load the emotion recognition model"""108 checkpoint = torch.load(model_path, map_location=device)109 110 model = EmotionCNN(num_classes=7)111 model.load_state_dict(checkpoint['model_state_dict'])112 model.to(device)113 model.eval()114 115 return model116 117def preprocess_face(face_img, size=(48, 48)):118 """Preprocess face image for emotion detection"""119 transform = transforms.Compose([120 transforms.Resize(size),121 transforms.ToTensor(),122 transforms.Normalize(mean=[0.5], std=[0.5])123 ])124 125 # Convert to PIL Image126 if isinstance(face_img, np.ndarray):127 face_img = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB))128 129 # Convert to grayscale130 face_img = face_img.convert('L')131 132 # Apply transformations133 face_tensor = transform(face_img).unsqueeze(0)134 return face_tensor135 136def process_video(video_path: str) -> Dict[str, Any]:137 """138 Process a video file and return emotion detection results.139 140 Args:141 video_path (str): Path to the video file142 143 Returns:144 Dict containing:145 - success (bool): Whether processing was successful146 - message (str): Status message147 - results (List[Dict]): List of emotion detection results148 - error (str): Error message if any149 """150 global largest_face_detections151 largest_face_detections = [] # Reset detections for new video152 153 # Paths - adjust these paths according to your Hugging Face Space154 face_cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'155 emotion_model_path = "./models/best_emotion_model.pth" # Path in Hugging Face Space156 157 # Check if models exist158 if not os.path.exists(face_cascade_path):159 return {160 "success": False,161 "message": "Face cascade classifier not found",162 "results": [],163 "error": f"Error: Face cascade classifier not found at {face_cascade_path}"164 }165 166 if not os.path.exists(emotion_model_path):167 return {168 "success": False,169 "message": "Emotion model not found",170 "results": [],171 "error": f"Error: Emotion model not found at {emotion_model_path}"172 }173 174 # Set device175 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')176 177 # Load models178 try:179 face_cascade = cv2.CascadeClassifier(face_cascade_path)180 emotion_model = load_emotion_model(emotion_model_path, device)181 except Exception as e:182 return {183 "success": False,184 "message": "Error loading models",185 "results": [],186 "error": str(e)187 }188 189 # Emotion labels190 emotions = ['Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral']191 192 # Open video193 cap = cv2.VideoCapture(video_path)194 if not cap.isOpened():195 return {196 "success": False,197 "message": "Could not open video file",198 "results": [],199 "error": f"Error: Could not open video file at {video_path}"200 }201 202 frame_count = 0203 total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))204 205 while True:206 ret, frame = cap.read()207 if not ret:208 break209 210 frame_count += 1211 212 # Variables to track largest face213 largest_face_area = 0214 current_detection = None215 216 # Convert frame to grayscale for face detection217 gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)218 219 # Detect faces using Haar Cascade220 faces = face_cascade.detectMultiScale(221 gray,222 scaleFactor=1.1,223 minNeighbors=5,224 minSize=(30, 30)225 )226 227 # Process each detected face228 for (x, y, w, h) in faces:229 # Calculate face area230 face_area = w * h231 232 # Extract face region with margin233 margin = 20234 x1 = max(0, x - margin)235 y1 = max(0, y - margin)236 x2 = min(frame.shape[1], x + w + margin)237 y2 = min(frame.shape[0], y + h + margin)238 239 face_img = frame[y1:y2, x1:x2]240 241 # Skip if face is too small242 if face_img.size == 0 or face_img.shape[0] < 20 or face_img.shape[1] < 20:243 continue244 245 # Convert face to PIL Image and preprocess246 face_tensor = preprocess_face(face_img)247 248 # Predict emotion249 with torch.no_grad():250 face_tensor = face_tensor.to(device)251 output = emotion_model(face_tensor)252 probabilities = torch.nn.functional.softmax(output, dim=1)253 emotion_idx = torch.argmax(output, dim=1).item()254 confidence = probabilities[0][emotion_idx].item()255 256 # Get emotion label257 emotion = emotions[emotion_idx]258 259 # Update largest face if current face is larger260 if face_area > largest_face_area:261 largest_face_area = face_area262 current_detection = {263 'emotion': emotion,264 'confidence': confidence,265 'timestamp': time.time(),266 'frame_number': frame_count267 }268 269 # Add current detection to history if a face was detected270 if current_detection:271 largest_face_detections.append(current_detection)272 273 # Release resources274 cap.release()275 276 # Process results277 if not largest_face_detections:278 return {279 "success": True,280 "message": "No faces detected in video",281 "results": {282 "average_emotions": {},283 "dominant_emotion": None,284 "detections": [],285 "summary": {286 "total_frames": total_frames,287 "total_detections": 0288 }289 },290 "error": None291 }292 293 emotion_scores = {e: [] for e in emotions} # Initialize with all emotion types294 295 for detection in largest_face_detections:296 emotion = detection['emotion']297 confidence = detection['confidence']298 emotion_scores[emotion].append(confidence)299 300 # Calculate summary statistics301 average_emotions = {302 e: sum(scores)/len(scores) if scores else 0 303 for e, scores in emotion_scores.items()304 }305 306 # Get dominant emotion based on average confidence307 dominant_emotion = max(average_emotions.items(), key=lambda x: x[1])[0]308 309 return {310 "success": True,311 "message": "Video processed successfully",312 "results": {313 "average_emotions": average_emotions,314 "dominant_emotion": dominant_emotion,315 # "detections": largest_face_detections, # Optional: include all detections316 # "summary": {317 # "total_frames": total_frames,318 # "total_detections": len(largest_face_detections),319 # "emotions_count": {e: len(s) for e, s in emotion_scores.items()},320 # "dominant_emotion": dominant_emotion321 # }322 },323 "error": None324 }325 326# Gradio Interface Functions327def gradio_analyze_video(video_path: str):328 """Wrapper function for Gradio interface"""329 result = process_video(video_path)330 if not result["success"]:331 return {"error": result.get("error", "Processing failed")}332 333 # Format results for better Gradio display334 summary = result["results"]["summary"]335 detections = result["results"]["detections"]336 337 # output = {338 # "summary": {339 # "total_frames": summary["total_frames"],340 # "faces_detected": summary["total_detections"],341 # "dominant_emotion": summary["dominant_emotion"],342 # "emotion_distribution": summary["emotions_count"]343 # },344 # "sample_detections": detections[:5] # Show first 5 detections345 # }346 # return output347 348 output = {349 "average_emotions": result["results"]["average_emotions"],350 "dominant_emotion": result["results"]["dominant_emotion"],351 # "frames_analyzed": result["results"]["summary"]["total_frames"],352 # "faces_detected": result["results"]["summary"]["total_detections"]353 }354 return output355 356def save_upload_file_tmp(upload_file: UploadFile) -> str:357 """Save uploaded file to temporary location"""358 try:359 suffix = Path(upload_file.filename).suffix360 with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:361 tmp.write(upload_file.file.read())362 return tmp.name363 finally:364 upload_file.file.close()365 366# Gradio Interface367with gr.Blocks(title="Video Emotion Detection", theme=gr.themes.Soft()) as demo:368 gr.Markdown("""369 # 🎭 Video Emotion Detection370 Upload a video to analyze facial emotions frame by frame371 """)372 373 with gr.Row():374 with gr.Column():375 video_input = gr.Video(376 label="Upload Video", 377 sources=["upload"] # Corrected line378 )379 submit_btn = gr.Button("Analyze Video", variant="primary")380 381 with gr.Column():382 output_json = gr.JSON(label="Analysis Results")383 gr.Markdown("""384 ### Results Interpretation385 - **Dominant Emotion**: Most frequently detected emotion386 - **Emotion Distribution**: Count of each emotion detected387 - **Sample Detections**: First 5 emotion detections388 """)389 390 submit_btn.click(391 fn=gradio_analyze_video,392 inputs=video_input,393 outputs=output_json,394 api_name="predict"395 )396 397 398# FastAPI Endpoints399@app.post("/api/analyze-video")400async def analyze_video(file: UploadFile = File(...)):401 """Original FastAPI endpoint"""402 try:403 temp_path = save_upload_file_tmp(file)404 result = process_video(temp_path)405 os.unlink(temp_path)406 407 if not result["success"]:408 raise HTTPException(status_code=400, detail=result.get("error", "Processing failed"))409 return JSONResponse(content=result)410 411 except Exception as e:412 if 'temp_path' in locals() and os.path.exists(temp_path):413 os.unlink(temp_path)414 raise HTTPException(status_code=500, detail=str(e))415 416@app.get("/", response_class=HTMLResponse)417async def root():418 """Redirect root to Gradio interface"""419 return """420 <html>421 <head>422 <title>Video Emotion Detection</title>423 <meta http-equiv="refresh" content="0; url=/gradio/" />424 </head>425 <body>426 <p>Redirecting to Gradio interface... <a href="/gradio">Click here</a> if not redirected.</p>427 </body>428 </html>429 """430 431# Mount Gradio app to FastAPI432app = gr.mount_gradio_app(app, demo, path="/gradio")433 434if __name__ == "__main__":435 import uvicorn436 uvicorn.run(app, host="0.0.0.0", port=7860)