Update run_caption_ds.py

run_caption_ds.py

@@ -123,22 +123,29 @@ def main():
     # Load dataset
     print(f"Loading dataset: {args.dataset_name}")
     dataset = load_dataset(args.dataset_name)
+    len_ = len(dataset["train"])
+    #len_ = 10
+    
+    # Initialize a list to store captions
+    captions = []
 
     # Generate captions for each image in the dataset
-    def add_caption(example):
+    print("Generating captions...")
+    for idx, example in enumerate(tqdm(dataset["train"].select(range(len_)), desc="Processing images")):  # 假设数据集是 "train" 拆分
         try:
             # Generate caption
             caption = generate_caption(example[args.image_column], clip_processor, clip_model, tokenizer, text_model, image_adapter)
-            # Add caption to the example
-            example[args.caption_column] = caption
+            captions.append(caption)
+            # Print the generated caption
+            print(f"Caption for image {idx + 1}: {caption}")
         except Exception as e:
-            print(f"Error processing image: {e}")
-            example[args.caption_column] = ""  # 如果出错，保存空字符串
-        return example
+            print(f"Error processing image {idx + 1}: {e}")
+            captions.append("")  # 如果出错，保存空字符串
+            print(f"Caption for image {idx + 1}: [Error]")
 
-    # Apply the function to the dataset
-    print("Generating captions...")
-    dataset = dataset.map(add_caption, desc="Generating captions")
+    # Add captions to the dataset
+    print("Adding captions to the dataset...")
+    dataset = dataset["train"].select(range(len_)).add_column(args.caption_column, captions)  # 将 captions 添加到数据集
 
     # Save the dataset with captions
     print(f"Saving dataset to {args.output_path}")