vlm-dataset-captioner 0.0.1__py3-none-any.whl

vlm_dataset_captioner/vlm_caption.py

@@ -0,0 +1,219 @@
+import os
+import re
+
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+DEFAULT_MODEL = "Qwen/Qwen2.5-VL-32B-Instruct"
+IMAGE_FILE_EXTENSIONS = (
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".gif",
+    ".bmp",
+    ".tif",
+)
+
+
+def init_model(model_name=None):
+    if model_name is None:
+        model_name = DEFAULT_MODEL
+        print(f"INFO: No model name provided. Initializing default model {model_name}.")
+
+    print(f"INFO: Initializing model {model_name}.", flush=True)
+
+    # Load the model on the available device(s)
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        model_name, torch_dtype="auto", device_map="auto"
+    )
+
+    print(f"INFO: Model {model_name} loaded successfully.", flush=True)
+    print(f"INFO: Loading processor for model {model_name}.", flush=True)
+
+    # Load the default processor
+    processor = AutoProcessor.from_pretrained(model_name)
+
+    print(f"INFO: Processor for model {model_name} loaded successfully.", flush=True)
+
+    return model, processor
+
+
+def get_prompt_for_directory(directory_path):
+    prompt = ""
+    prompt_file_path = os.path.join(directory_path, "prompt.txt")
+    try:
+        with open(prompt_file_path) as prompt_file:
+            prompt = prompt_file.read()
+    except FileNotFoundError:
+        print(
+            f"WARN: Prompt file not found for directory {prompt_file_path}. Using default prompt.",
+            flush=True,
+        )
+        prompt = "Describe the image in detail."
+
+    print(f"INFO: Using prompt: '{prompt}'", flush=True)
+
+    return prompt
+
+
+def is_image_file(filename):
+    return filename.lower().endswith(IMAGE_FILE_EXTENSIONS)
+
+
+def is_image_directory(directory_path):
+    return any(is_image_file(filename) for filename in os.listdir(directory_path))
+
+
+def get_messages(prompt, image):
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": prompt},
+            ],
+        }
+    ]
+
+
+def contains_chinese(text_string):
+    # The Unicode range for common CJK Unified Ideographs (Han characters)
+    # is typically from U+4E00 to U+9FFF.
+    chinese_char_pattern = re.compile(r"[\u4e00-\u9fff]")
+    return bool(chinese_char_pattern.search(text_string))
+
+
+def caption_image(prompt, image, model, processor, max_new_tokens=None):
+    messages = get_messages(prompt, image)
+
+    # Prepare inputs for the model
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to("cuda")
+
+    # Generate caption
+    generated_ids = model.generate(
+        **inputs,
+        max_new_tokens=128,
+        do_sample=True,
+        top_p=1.0,
+        temperature=0.7,
+        top_k=50,
+    )
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :]
+        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+
+    # Truncate caption if it exceeds max_new_tokens
+    if max_new_tokens is not None and len(generated_ids_trimmed[0]) > max_new_tokens:
+        print(
+            f"WARN: Generated tokens for {image} exceed max_new_tokens={max_new_tokens}. Truncating caption.",
+            flush=True,
+        )
+        generated_ids_trimmed[0] = generated_ids_trimmed[0][:max_new_tokens]
+
+    # Decode the generated tokens to text
+    output_text = processor.batch_decode(
+        generated_ids_trimmed,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=False,
+    )
+
+    return output_text[0]
+
+
+def write_caption_to_file(image_file, caption, output_directory):
+    if not os.path.exists(output_directory):
+        os.makedirs(output_directory)
+
+    caption_file = os.path.join(
+        output_directory, f"{os.path.splitext(image_file)[0]}.txt"
+    )
+    with open(caption_file, "w") as f:
+        f.write(caption)
+
+
+def ignore_file(filename, ignore_substring):
+    if ignore_substring is None:
+        return False
+    return ignore_substring in filename
+
+
+def requires_caption(image_file, output_directory, overwrite):
+    caption_file = os.path.join(
+        output_directory, f"{os.path.splitext(image_file)[0]}.txt"
+    )
+    return overwrite or not os.path.exists(caption_file)
+
+
+def caption_entire_directory(
+    directory_path,
+    output_directory,
+    model,
+    processor,
+    max_new_tokens=None,
+    ignore_substring=None,
+    num_captions=None,
+    overwrite=False,
+):
+    print(
+        f"INFO: Processing directory {directory_path} for image captions.", flush=True
+    )
+
+    if not is_image_directory(directory_path):
+        for subdir in os.listdir(directory_path):
+            if not ignore_file(subdir, ignore_substring):
+                subdir_path = os.path.join(directory_path, subdir)
+                if os.path.isdir(subdir_path):
+                    caption_entire_directory(
+                        subdir_path,
+                        os.path.join(output_directory, subdir),
+                        model,
+                        processor,
+                        max_new_tokens,
+                        ignore_substring,
+                        num_captions,
+                        overwrite,
+                    )
+    else:
+        prompt = get_prompt_for_directory(directory_path)
+        for image_file in os.listdir(directory_path):
+            if (
+                not ignore_file(image_file, ignore_substring)
+                and requires_caption(image_file, output_directory, overwrite)
+                and image_file.lower().endswith(
+                    (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tif")
+                )
+            ):
+                try:
+                    caption = ""
+                    for i in range(int(num_captions) if num_captions else 1):
+                        if i != 0:
+                            caption += "\n"
+
+                        while True:
+                            caption += caption_image(
+                                prompt,
+                                os.path.join(directory_path, image_file),
+                                model,
+                                processor,
+                                max_new_tokens,
+                            )
+                            if not contains_chinese(caption):
+                                break
+                    write_caption_to_file(image_file, caption, output_directory)
+                except Exception as e:
+                    print(
+                        f"WARN: Error processing image {image_file} in {directory_path}: {e}",
+                        flush=True,
+                    )

vlm_dataset_captioner/vlm_caption_cli.py

@@ -0,0 +1,83 @@
+import argparse
+
+from vlm_caption import caption_entire_directory, init_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Caption images from a dataset using a VLM."
+    )
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        required=True,
+        help="The path of the input directory containing images to be captioned.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="The HuggingFace model used to generate captions.",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=None,
+        help="The maximum number of tokens to be generated in any given caption.",
+    )
+    parser.add_argument(
+        "--ignore_substring",
+        type=str,
+        default=None,
+        help="Ignore files and subdirectories that contain this substring in their names.",
+    )
+    parser.add_argument(
+        "--num_captions",
+        type=str,
+        default=None,
+        help="Number of captions to be generated.",
+    )
+    parser.add_argument(
+        "--overwrite",
+        type=bool,
+        default=False,
+        help="If true, overwrites existing captions.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The directory to act as the root of the caption file structure. Defaults to `<input_dir>_caption`.",
+    )
+    return parser.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"
+
+    if args.model is not None:
+        print(f"INFO: Using model {args.model} for captioning.", flush=True)
+    if args.max_length is not None:
+        print(f"INFO: Setting max length to {args.max_length} tokens.", flush=True)
+    if args.ignore_substring is not None:
+        print(
+            f"INFO: Ignoring files/directories containing substring '{args.ignore_substring}'.",
+            flush=True,
+        )
+
+    caption_entire_directory(
+        args.input_dir,
+        output_dir,
+        model,
+        processor,
+        max_new_tokens=args.max_length,
+        ignore_substring=args.ignore_substring,
+        num_captions=args.num_captions,
+        overwrite=args.overwrite,
+    )
+
+
+main()

vlm_dataset_captioner-0.0.1.dist-info/METADATA

@@ -0,0 +1,77 @@
+Metadata-Version: 2.4
+Name: vlm-dataset-captioner
+Version: 0.0.1
+Summary: Uses a VLM to caption images from a dataset.
+Author: Alex Senden
+Maintainer: Alex Senden
+License: MIT
+Keywords: computer-vision,image-captioning,machine-learning,vision-language-model
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Requires-Python: >=3.8
+Requires-Dist: accelerate
+Requires-Dist: huggingface-hub
+Requires-Dist: qwen-vl-utils
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: transformers
+Description-Content-Type: text/markdown
+
+# VLM Captioner
+
+Uses a VLM (initially configured to use `Qwen2.5-VL-32B-Instruct`) to caption images from a dataset.
+
+### Dataset Structure
+
+One VLM prompt will be used for each entire image directory.
+For each image directory, a mirror file structure is created with the suffix `_caption`. This structure contains individual `.txt` caption files with filenames matching that of their image counterparts.
+
+```
+dataset/
+└── top_level_folder_1/
+    ├── image_folder_1 (contains prompt for entire folder)/
+    │   ├── prompt.txt
+    │   ├── image_1.png
+    │   ├── image_2.png
+    │   └── ...
+    └── ...
+```
+
+### Running
+
+First, install the required packages:
+
+```
+pip install -r requirements.txt
+```
+
+Then, run the script:
+
+```
+python vlm_caption_cli.py --input_dir=<input_dir> [--model=<vlm_model>]
+```
+
+### Command Line Args
+
+##### Required Args:
+
+```
+--input_dir=<input_dir> || The path of the input directory containing images to be captioned.
+```
+
+##### Optional Args:
+
+```
+--model=<vlm_model> || VLM to use to generate captions
+--max_length=<max_new_tokens> || Maximum number of new tokens before truncation
+--ignore_substring=<ignore_substring> || Ignore files/directories containing this substring
+--num_captions=<number_of_captions> || Number of captions to generate per image
+--overwrite=<True/False> || If true, overwrites captions that already exist
+--output_dir=<output_dir> || The directory to act as the root of the caption file structure. Defaults to `<input_dir>_caption`.
+```

vlm_dataset_captioner-0.0.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+vlm_dataset_captioner/vlm_caption.py,sha256=k711kghgWmWXZIYva8t7v2ew519BjwcchZt3vwzmfZc,6854
+vlm_dataset_captioner/vlm_caption_cli.py,sha256=i1SS43ga2hpxCAQ2XtOkzNFBfI0zKZ5y-aKWI6djt4M,2341
+vlm_dataset_captioner-0.0.1.dist-info/METADATA,sha256=b7B8SwZAIIs2DPsfoY0nUb3ZomPujDl0-LEYDGco-x8,2430
+vlm_dataset_captioner-0.0.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+vlm_dataset_captioner-0.0.1.dist-info/entry_points.txt,sha256=k-zH3SWvcplaeDuGV4J6OyHKLr9GieWcOhRB5sF2pEI,75
+vlm_dataset_captioner-0.0.1.dist-info/RECORD,,

vlm_dataset_captioner-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

vlm_dataset_captioner-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+vlm-caption = vlm_dataset_captioner.vlm_caption_cli:main