update src

app.py

@@ -52,7 +52,8 @@ def make_video(video_path, outdir='./summarized_video', algorithm='Kmeans', mode
         clip_sample_rate = 1
         # F
         num_frames = 8
-
+        
+        original_frames = []
         frames = []
         features = []
 
@@ -64,7 +65,10 @@ def make_video(video_path, outdir='./summarized_video', algorithm='Kmeans', mode
             if not ret:
                 break
             
+            # use the original frames to write the output video
+            original_frames.append(raw_frame)
             raw_frame = cv2.resize(raw_frame, resize_to)
+            # use the resized frames to extract features
             frames.append(raw_frame)
 
         # Find key frames by selecting frames with clip_sample_rate
@@ -109,7 +113,7 @@ def make_video(video_path, outdir='./summarized_video', algorithm='Kmeans', mode
         
         video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, (frames[0].shape[1], frames[0].shape[0]))
         for idx in selected_frames:
-            video_writer.write(frames[idx])
+            video_writer.write(original_frames[idx])
         
         raw_video.release()
         video_writer.release()

utils.py

@@ -55,7 +55,7 @@ def to_video(selected_frames, frames, output_path, video_fps):
 def load_model(model_version):
     try:
         DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-        model = TimesformerModel.from_pretrained(f"facebook/timesformer-base-finetuned-{model_version}").to(DEVICE).eval()
+        model = TimesformerModel.from_pretrained(f"facebook/timesformer-base-finetuned-{model_version}")
         processor=VideoMAEImageProcessor.from_pretrained("MCG-NJU/videomae-base")
         return model, processor, DEVICE