From 7a2a9908a2931843d235cae9fbf110b0bce72ad2 Mon Sep 17 00:00:00 2001 From: Mandlin Sarah Date: Fri, 30 Aug 2024 21:15:17 -0700 Subject: [PATCH] Improved Image Loading Error Handling --- captioner/fast_captioner_lmdeploy.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/captioner/fast_captioner_lmdeploy.py b/captioner/fast_captioner_lmdeploy.py index e76a2d6..8875b1f 100644 --- a/captioner/fast_captioner_lmdeploy.py +++ b/captioner/fast_captioner_lmdeploy.py @@ -51,9 +51,13 @@ def __len__(self): def __getitem__(self, idx): video_img_list = self.video_dir_list[idx] - img_path_list = os.listdir(video_img_list) - img_path_list.sort() - images = [Image.open(os.path.join(video_img_list, img_path)).convert('RGB') for img_path in img_path_list] + try: + img_path_list = os.listdir(video_img_list) + img_path_list.sort() + images = [Image.open(os.path.join(video_img_list, img_path)).convert('RGB') for img_path in img_path_list] + except Exception as e: + print(f"Error loading image: {e}") + return None img_width, img_height = images[0].size grid_image = Image.new('RGB', (self.img_grid_w * img_width, self.img_grid_h * img_height)) @@ -72,8 +76,9 @@ def custom_collate_fn(batch): input_list = [] filename_list = [] for item in batch: - input_list.append(item[0]) - filename_list.append(item[1]) + if item: + input_list.append(item[0]) + filename_list.append(item[1]) return input_list, filename_list def parse_args(): @@ -111,4 +116,4 @@ def parse_args(): } filename = filename.split('/')[-1] with open(os.path.join(args.save_path, filename+'.json'), 'w') as f: - f.write(json.dumps(finish_data, indent=2, ensure_ascii=False)) \ No newline at end of file + f.write(json.dumps(finish_data, indent=2, ensure_ascii=False))