Update src/utils.py

src/utils.py

@@ -1,50 +1,71 @@
-# Importing the requirements
-from PIL import Image
-from decord import VideoReader, cpu
-
-
-# Maximum number of frames to use
-MAX_NUM_FRAMES = 64  # If CUDA OOM, set a smaller number
-
-
-def encode_video(video_path):
-    """
-    Encodes a video file into a list of frames.
-
-    Args:
-        video_path (str): The path to the video file.
-
-    Returns:
-        list: A list of frames, where each frame is represented as an Image object.
-    """
-
-    def uniform_sample(l, n):
-        """
-        Uniformly samples elements from a list.
-
-        Args:
-            - l (list): The input list.
-            - n (int): The number of elements to sample.
-
-        Returns:
-            list: A list of sampled elements.
-        """
-        gap = len(l) / n
-        idxs = [int(i * gap + gap / 2) for i in range(n)]
-        return [l[i] for i in idxs]
-
-    # Read the video file and sample frames
-    vr = VideoReader(video_path, ctx=cpu(0))
-    sample_fps = round(vr.get_avg_fps() / 1)  # FPS
-    frame_idx = [i for i in range(0, len(vr), sample_fps)]
-
-    # Uniformly sample frames if the number of frames is too large
-    if len(frame_idx) > MAX_NUM_FRAMES:
-        frame_idx = uniform_sample(frame_idx, MAX_NUM_FRAMES)
-
-    # Extract frames from the video
-    frames = vr.get_batch(frame_idx).asnumpy()
-    frames = [Image.fromarray(v.astype("uint8")) for v in frames]
-
-    # Return video frames
-    return frames
+# Importing the requirements
+from PIL import Image
+from decord import VideoReader, cpu
+
+
+# Maximum number of frames to use
+MAX_NUM_FRAMES = 64  # If CUDA OOM, set a smaller number
+
+
+def encode_image(image):
+    if not isinstance(image, Image.Image):
+        if hasattr(image, 'path'):
+            image = Image.open(image.path).convert("RGB")
+        else:
+            image = Image.open(image.file.path).convert("RGB")
+    # resize to max_size
+    max_size = 448*16 
+    if max(image.size) > max_size:
+        w,h = image.size
+        if w > h:
+            new_w = max_size
+            new_h = int(h * max_size / w)
+        else:
+            new_h = max_size
+            new_w = int(w * max_size / h)
+        image = image.resize((new_w, new_h), resample=Image.BICUBIC)
+    return image
+
+
+def encode_video(video_path):
+    """
+    Encodes a video file into a list of frames.
+
+    Args:
+        video_path (str): The path to the video file.
+
+    Returns:
+        list: A list of frames, where each frame is represented as an Image object.
+    """
+
+    def uniform_sample(l, n):
+        """
+        Uniformly samples elements from a list.
+
+        Args:
+            - l (list): The input list.
+            - n (int): The number of elements to sample.
+
+        Returns:
+            list: A list of sampled elements.
+        """
+        gap = len(l) / n
+        idxs = [int(i * gap + gap / 2) for i in range(n)]
+        return [l[i] for i in idxs]
+
+    # Read the video file and sample frames
+    vr = VideoReader(video_path, ctx=cpu(0))
+    sample_fps = round(vr.get_avg_fps() / 1)  # FPS
+    frame_idx = [i for i in range(0, len(vr), sample_fps)]
+
+    # Uniformly sample frames if the number of frames is too large
+    if len(frame_idx) > MAX_NUM_FRAMES:
+        frame_idx = uniform_sample(frame_idx, MAX_NUM_FRAMES)
+
+    # Extract frames from the video
+    frames = vr.get_batch(frame_idx).asnumpy()
+    frames = [Image.fromarray(v.astype("uint8")) for v in frames]
+    frames = [encode_image(v) for v in frames]
+
+    # Return video frames
+    return frames