From 4f69a9cfddfb989ada769dbfc7e5f1c75da2ed86 Mon Sep 17 00:00:00 2001
From: ShowLo <945519267@qq.com>
Date: Thu, 19 Sep 2024 17:09:35 +0800
Subject: [PATCH] Fix the bug that causes an infinite loop when the total
 number of frames in the video does not exceed 11.

eg, the video has 11 frames, when select the NO.6 frame, `while abs(random_element - img_idx) <= 5:` will result in an infinite loop
---
 train_codes/DataLoader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train_codes/DataLoader.py b/train_codes/DataLoader.py
index eb67842..babfe39 100644
--- a/train_codes/DataLoader.py
+++ b/train_codes/DataLoader.py
@@ -135,7 +135,7 @@ class Dataset(object):
             #随机选择某个video里
             vidname = self.all_videos[idx].split('/')[-1]
             video_imgs = self.all_img_names[idx]
-            if len(video_imgs) == 0:
+            if len(video_imgs) <= 11:
                 continue
             img_name = random.choice(video_imgs)
             img_idx = int(basename(img_name).split(".")[0])