fixing

Dockerfile

@@ -1,6 +1,3 @@
-# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
-# You will also find guides on how best to write your Dockerfile
-
 FROM python:3.12-slim
 
 ENV TORCH_HOME=/tmp/torch_cache
@@ -9,7 +6,12 @@ WORKDIR /code
 
 COPY ./requirements.txt /code/requirements.txt
 
-RUN apt-get update && apt-get install -y libgl1 libglib2.0-0
+# Install ffmpeg along with other dependencies
+RUN apt-get update && apt-get install -y \
+    libgl1 \
+    libglib2.0-0 \
+    ffmpeg \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
 
 RUN pip install --no-cache-dir --upgrade pip \
     && pip install --no-cache-dir -r /code/requirements.txt

services/summarizer.py

@@ -2,6 +2,9 @@ import cv2
 import torch
 from config import SCORE_THRESHOLD
 from services.model_loader import load_model
+import subprocess
+import os
+import numpy as np
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = load_model("Model/epoch-199.pkl")
@@ -20,13 +23,12 @@ def get_scores(features):
     return scores
 
 def get_selected_indices(scores, picks, threshold=SCORE_THRESHOLD):
+    selected = [picks[i] for i, score in enumerate(scores) if score >= threshold]
     print("Threshold for selection:", threshold)
-    print("Scores:", scores.shape, scores)
-    print("Picks:", picks.shape, picks)
-    return [picks[i] for i, score in enumerate(scores) if score >= threshold]
-
-import subprocess
-import os
+    print("Scores:", len(scores), scores)
+    print("Picks:", len(picks), picks)
+    print("Selected indices:", len(selected), selected)
+    return selected
 
 def save_summary_video(video_path, selected_indices, output_path, fps=15):
     import cv2
@@ -34,24 +36,36 @@ def save_summary_video(video_path, selected_indices, output_path, fps=15):
     cap = cv2.VideoCapture(video_path)
     selected = set(selected_indices)
     frame_id = 0
-    frames = {}
+    frames = []
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
         if frame_id in selected:
-            frames[frame_id] = frame
+            frames.append(frame)
         frame_id += 1
     cap.release()
 
-    if not frames:
+    if len(frames) == 0:
         print("No frames selected.")
         return
 
-    h, w, _ = list(frames.values())[0].shape
+    h, w, _ = frames[0].shape
+    print("Video dimensions:", w, h)
+
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (h, w))
+    for frame in frames:
+        out.write(frame)
+    out.release()
+
+    print("Fixing the video with ffmpeg")
+    # fix_video_with_ffmpeg(output_path)
 
-    writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-    for fid in sorted(frames.keys()):
-        writer.write(frames[fid])
-    writer.release()
+def fix_video_with_ffmpeg(path):
+    temp_path = path + ".fixed.mp4"
+    subprocess.run([
+        "ffmpeg", "-y", "-i", path,
+        "-vcodec", "libx264", "-pix_fmt", "yuv420p", temp_path
+    ])
+    os.replace(temp_path, path)