From d57c779217158e570219e4f0bc2906fd5d555fd4 Mon Sep 17 00:00:00 2001
From: Wang Qiang <37444407+wangqiang9@users.noreply.github.com>
Date: Wed, 10 Jul 2024 15:42:30 +0800
Subject: [PATCH] Added exception handling for video read bucket_sampler.py

---
 easyanimate/data/bucket_sampler.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/easyanimate/data/bucket_sampler.py b/easyanimate/data/bucket_sampler.py
index 2c5fded..78eb683 100644
--- a/easyanimate/data/bucket_sampler.py
+++ b/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[:]
\ No newline at end of file
+                    del bucket[:]