Upload 4 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo_video.mp4 filter=lfs diff=lfs merge=lfs -text

demo_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cb921ae36d4d7ce8f9e13586166aaaba55ef9f3f384fcc43bffffae39f71cb0
+size 4432717

process_emotion_video.py

@@ -0,0 +1,119 @@
+import os
+import tempfile
+import torch
+import cv2
+from PIL import Image
+from torchvision import transforms
+from facenet_pytorch import MTCNN
+
+def process_emotion(input_video):
+    try:
+        # Load the pre-trained CNN model
+        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
+        num_ftrs = model.fc.in_features
+        model.fc = torch.nn.Linear(num_ftrs, 7)  # 7 classes for emotions
+
+        # Load the pre-trained weights
+        model.eval()
+
+        # Define emotions
+        emotions = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
+
+        # Initialize MTCNN for face detection
+        mtcnn = MTCNN()
+
+        # Open the video capture
+        cap = cv2.VideoCapture(input_video.name)
+
+        if not cap.isOpened():
+            return "Error: Could not open video capture."
+
+        # Create a temporary directory to store the frames
+        temp_dir = tempfile.mkdtemp()
+        frame_paths = []
+
+        # Define image transformations
+        preprocess = transforms.Compose([
+            transforms.Resize(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+        # Process each frame of the video
+        frame_count = 0
+        while cap.isOpened():
+            ret, frame = cap.read()
+
+            if not ret:
+                break
+
+            if frame is None:
+                print("Error: Empty frame.")
+                continue
+
+            frame_count += 1
+
+            pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            face_boxes, _ = mtcnn.detect(pil_frame)
+
+            if face_boxes is None:
+                print(f"No face detected in frame {frame_count}.")
+            else:
+                for box in face_boxes:
+                    x1, y1, x2, y2 = box.astype(int)
+                    w = x2 - x1
+                    h = y2 - y1
+
+                    face_roi = frame[y1:y2, x1:x2]
+                    pil_face = Image.fromarray(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB))
+
+                    input_tensor = preprocess(pil_face)
+                    input_batch = input_tensor.unsqueeze(0)
+
+                    with torch.no_grad():
+                        output = model(input_batch)
+
+                    predicted_emotion = output.argmax().item()
+
+                    emotion_label = emotions[predicted_emotion]
+
+                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                    cv2.putText(frame, emotion_label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+
+                    frame_path = os.path.join(temp_dir, f"{frame_count:04d}.png")
+                    cv2.imwrite(frame_path, frame)
+                    frame_paths.append(frame_path)
+
+        cap.release()
+        cv2.destroyAllWindows()
+
+        if not frame_paths:
+            return "No faces detected in the video."
+
+        # Convert the frames to a video
+        output_path = "output_video.mp4"
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+
+        # Get dimensions from the first frame
+        sample_frame = cv2.imread(frame_paths[0])
+        if sample_frame is None:
+            return "Error: Unable to read sample frame."
+
+        frame_height, frame_width, _ = sample_frame.shape
+
+        out = cv2.VideoWriter(output_path, fourcc, 30, (frame_width, frame_height))
+
+        for frame_path in frame_paths:
+            frame = cv2.imread(frame_path)
+            if frame is not None:
+                out.write(frame)
+                os.remove(frame_path)
+            else:
+                print(f"Warning: Unable to read frame {frame_path}")
+
+        out.release()
+
+        return output_path
+
+    except Exception as e:
+        return f"An error occurred: {str(e)}"

requirements.txt

@@ -0,0 +1,7 @@
+ffmpeg
+facenet-pytorch
+torch 
+torchvision
+streamlit 
+opencv-python-headless 
+pyngrok

web.py

@@ -0,0 +1,41 @@
+import streamlit as st
+from process_emotion_video import process_emotion
+
+def main():
+    st.title("Emotion Recognition in Video")
+    option = st.radio("Choose an option:", ("Upload a video file", "Example Video"))
+    if option == "Upload a video file":
+        # Upload video file
+        uploaded_file = st.file_uploader("Upload a video file", type=["mp4"])
+    
+        if uploaded_file is not None:
+            st.write("Processing...")
+    
+            # Process the video and get the output path
+            output_path = process_emotion(uploaded_file)
+    
+            if output_path.endswith(".mp4"):
+                st.video(output_path)
+    
+                # Add a download button for the processed video
+                st.download_button(
+                    label="Download Processed Video",
+                    data=output_path,
+                    key="download_processed_video",
+                    file_name="processed_video.mp4",
+                )
+            else:
+                st.error("An error occurred during video processing.")
+    else:
+        example_video_path = 'demo_video.mp4'
+        st.video(example_video_path)
+        st.download_button(
+            label="Download Processed Video",
+            data=example_video_path,
+            key="download_demo_video",
+            file_name="demo_video.mp4",
+            )
+    
+
+if __name__ == "__main__":
+    main()