facebookresearch · ShuoShenDe · Aug 8, 2024 · Aug 8, 2024 · Aug 8, 2024 · Aug 9, 2024
diff --git a/notebooks/video_predictor_example_with_png.ipynb b/notebooks/video_predictor_example_with_png.ipynb
diff --git a/notebooks/videos/car/0001.png b/notebooks/videos/car/0001.png
diff --git a/notebooks/videos/car/0002.png b/notebooks/videos/car/0002.png
diff --git a/notebooks/videos/car/0003.png b/notebooks/videos/car/0003.png
diff --git a/notebooks/videos/car/0004.png b/notebooks/videos/car/0004.png
diff --git a/notebooks/videos/car/0005.png b/notebooks/videos/car/0005.png
diff --git a/notebooks/videos/car/0006.png b/notebooks/videos/car/0006.png
diff --git a/notebooks/videos/car/0007.png b/notebooks/videos/car/0007.png
diff --git a/notebooks/videos/car/0008.png b/notebooks/videos/car/0008.png
diff --git a/notebooks/videos/car/0009.png b/notebooks/videos/car/0009.png
diff --git a/notebooks/videos/car/0010.png b/notebooks/videos/car/0010.png
diff --git a/notebooks/videos/car/0011.png b/notebooks/videos/car/0011.png
diff --git a/notebooks/videos/car/0012.png b/notebooks/videos/car/0012.png
diff --git a/notebooks/videos/car/0013.png b/notebooks/videos/car/0013.png
diff --git a/notebooks/videos/car/0014.png b/notebooks/videos/car/0014.png
diff --git a/notebooks/videos/car/0015.png b/notebooks/videos/car/0015.png
diff --git a/notebooks/videos/car/0016.png b/notebooks/videos/car/0016.png
diff --git a/notebooks/videos/car/0017.png b/notebooks/videos/car/0017.png
diff --git a/notebooks/videos/car/0018.png b/notebooks/videos/car/0018.png
diff --git a/notebooks/videos/car/0019.png b/notebooks/videos/car/0019.png
diff --git a/notebooks/videos/car/0020.png b/notebooks/videos/car/0020.png
diff --git a/notebooks/videos/car/0021.png b/notebooks/videos/car/0021.png
diff --git a/sam2/utils/misc.py b/sam2/utils/misc.py
@@ -179,7 +179,7 @@ def load_video_frames(
     compute_device=torch.device("cuda"),
 ):
     """
-    Load the video frames from a directory of JPEG files ("<frame_index>.jpg" format).
+    Load the video frames from a directory of JPEG or PNG files ("<frame_index>.jpg" "<frame_index>.png" format).
 
     The frames are resized to image_size x image_size and are loaded to GPU if
     `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`.
@@ -202,7 +202,7 @@ def load_video_frames(
     frame_names = [
         p
         for p in os.listdir(jpg_folder)
-        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png", ".PNG"]
     ]
     frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))
     num_frames = len(frame_names)
@@ -224,7 +224,7 @@ def load_video_frames(
         return lazy_images, lazy_images.video_height, lazy_images.video_width
 
     images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)
-    for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")):
+    for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG, PNG)")):
         images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)
     if not offload_video_to_cpu:
         images = images.to(compute_device)