Update preprocessing.py

preprocessing.py

@@ -17,25 +17,51 @@ def read_video(video_base64, num_frames=24, target_size=(224, 224)):
         frames.append(frame.to_ndarray(format="rgb24").astype(np.uint8))
     
     sampled_frames = sample_frames(frames, num_frames)
-    processed_frames = preprocess_frames(sampled_frames, target_size)
+    processed_frames = pad_and_resize(sampled_frames, target_size)
+    processed_frames = processed_frames.permute(1, 0, 2, 3)  # (T, C, H, W) -> (C, T, H, W)
     return processed_frames
 
 def sample_frames(frames, num_frames):
     total_frames = len(frames)
+    sampled_frames = list(frames)
     if total_frames <= num_frames:
+        # sampled_frames = frames
         if total_frames < num_frames:
             padding = [np.zeros_like(frames[0]) for _ in range(num_frames - total_frames)]
-            frames.extend(padding)
+            sampled_frames.extend(padding)
     else:
         indices = np.linspace(0, total_frames - 1, num=num_frames, dtype=int)
-        frames = [frames[i] for i in indices]
-    return np.array(frames)
+        sampled_frames = [frames[i] for i in indices]
 
-def preprocess_frames(frames, target_size):
+    return np.array(sampled_frames)
+
+    
+    # total_frames = len(frames)
+    # if total_frames <= num_frames:
+    #     if total_frames < num_frames:
+    #         padding = [np.zeros_like(frames[0]) for _ in range(num_frames - total_frames)]
+    #         frames.extend(padding)
+    # else:
+    #     indices = np.linspace(0, total_frames - 1, num=num_frames, dtype=int)
+    #     frames = [frames[i] for i in indices]
+    # return np.array(frames)
+
+def pad_and_resize(frames, target_size):
     transform = transforms.Compose([
         transforms.ToPILImage(),
         transforms.Resize(target_size),
         transforms.ToTensor()
     ])
     processed_frames = [transform(frame) for frame in frames]
-    return torch.stack(processed_frames).permute(1, 0, 2, 3).numpy()  # (T, C, H, W) -> (C, T, H, W)
+
+    return torch.stack(processed_frames)
+
+
+# def pad_and_resize(frames, target_size):
+#     transform = transforms.Compose([
+#         transforms.ToPILImage(),
+#         transforms.Resize(target_size),
+#         transforms.ToTensor()
+#     ])
+#     processed_frames = [transform(frame) for frame in frames]
+#     return torch.stack(processed_frames).permute(1, 0, 2, 3)  # (T, C, H, W) -> (C, T, H, W)