Added exception handling for video read bucket_sampler.py

easyanimate/data/bucket_sampler.py

@@ -243,6 +243,9 @@ class AspectRatioBatchSampler(BatchSampler):
                         videoid, name, page_dir = video_dict['videoid'], video_dict['name'], video_dict['page_dir']
                         video_dir = os.path.join(self.video_folder, f"{videoid}.mp4")
                     cap = cv2.VideoCapture(video_dir)
+                    if not cap.isOpened():
+                        print(f"Open video {video_dir} is error! ")
+                        continue
 
                     # 获取视频尺寸
                     width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))   # 浮点数转换为整数
@@ -354,6 +357,9 @@ class AspectRatioBatchImageVideoSampler(BatchSampler):
                         else:
                             video_dir = os.path.join(self.train_folder, video_id)
                         cap = cv2.VideoCapture(video_dir)
+                        if not cap.isOpened():
+                            print(f"Open video {video_dir} is error! ")
+                            continue
 
                         # 获取视频尺寸
                         width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))   # 浮点数转换为整数
@@ -376,4 +382,4 @@ class AspectRatioBatchImageVideoSampler(BatchSampler):
                 # yield a batch of indices in the same aspect ratio group
                 if len(bucket) == self.batch_size:
                     yield bucket[:]
-                    del bucket[:]
+                    del bucket[:]