vlm-dataset-captioner 0.0.2__tar.gz → 0.0.4__tar.gz

{vlm_dataset_captioner-0.0.2 → vlm_dataset_captioner-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vlm-dataset-captioner
-Version: 0.0.2
+Version: 0.0.4
 Summary: Uses a VLM to caption images from a dataset.
 Author: Alex Senden
 Maintainer: Alex Senden

{vlm_dataset_captioner-0.0.2 → vlm_dataset_captioner-0.0.4}/vlm_dataset_captioner/vlm_caption.py

@@ -49,7 +49,7 @@ def get_prompt_for_directory(directory_path):
             f"WARN: Prompt file not found for directory {prompt_file_path}. Using default prompt.",
             flush=True,
         )
-        prompt = "Describe the image in detail."
+        prompt = "In one short sentence. The caption will be used for image indexing and search, so include relevant details. 1 sentence only."
 
     print(f"INFO: Using prompt: '{prompt}'", flush=True)
 
@@ -86,6 +86,8 @@ def contains_chinese(text_string):
 def caption_image(prompt, image, model, processor, max_new_tokens=None):
     messages = get_messages(prompt, image)
 
+    print(f"INFO: Generating caption for image: {image}.", flush=True)
+
     # Prepare inputs for the model
     text = processor.apply_chat_template(
         messages, tokenize=False, add_generation_prompt=True
@@ -129,6 +131,8 @@ def caption_image(prompt, image, model, processor, max_new_tokens=None):
         clean_up_tokenization_spaces=False,
     )
 
+    print(f"INFO: Caption generated for image: {image}.", flush=True)
+
     return output_text[0]
 
 
@@ -206,15 +210,20 @@ def caption_entire_directory(
                             caption += "\n"
 
                         while True:
-                            caption += caption_image(
+                            individual_caption = caption_image(
                                 prompt,
                                 os.path.join(directory_path, image_file),
                                 model,
                                 processor,
                                 max_new_tokens,
                             )
-                            if not contains_chinese(caption):
+                            if not contains_chinese(individual_caption):
+                                caption += individual_caption
                                 break
+                            print(
+                                "WARN: Detected Chinese characters in caption. Regenerating...",
+                                flush=True,
+                            )
                     write_caption_to_file(image_file, caption, output_directory)
                 except Exception as e:
                     print(

{vlm_dataset_captioner-0.0.2 → vlm_dataset_captioner-0.0.4}/vlm_dataset_captioner/vlm_caption_cli.py

@@ -16,7 +16,7 @@ def parse_args():
     parser.add_argument(
         "--model",
         type=str,
-        default=None,
+        default="Qwen/Qwen2.5-VL-32B-Instruct",
         help="The HuggingFace model used to generate captions.",
     )
     parser.add_argument(
@@ -55,8 +55,10 @@ def parse_args():
 def main():
     args = parse_args()
     model, processor = init_model(args.model)
-    
-    output_dir = args.output_dir if args.output_dir is not None else f"{args.input_dir}_caption"
+
+    output_dir = (
+        args.output_dir if args.output_dir is not None else f"{args.input_dir}_caption"
+    )
 
     if args.model is not None:
         print(f"INFO: Using model {args.model} for captioning.", flush=True)
@@ -71,8 +73,8 @@ def main():
     caption_entire_directory(
         args.input_dir,
         output_dir,
-        model,
-        processor,
+        model=model,
+        processor=processor,
         max_new_tokens=args.max_length,
         ignore_substring=args.ignore_substring,
         num_captions=args.num_captions,