src/metadata.py

import os
import json
import time


from PIL import Image


from qwen_vl import (
    analyze_image_with_qwen
)


from ocr import (
    extract_text_with_paddleocr_vl_spotting
)


from grounding import (
    detect_objects_with_ram_dino
)


def build_metadata(
    dataset_path,
    output_path,
    checkpoint_path=None
):

    # ========================================================
    # Load existing metadata if checkpoint exists
    # ========================================================

    metadata = []


    if (
        checkpoint_path
        and os.path.exists(
            checkpoint_path
        )
    ):

        with open(
            checkpoint_path,
            "r",
            encoding="utf-8"
        ) as file:

            metadata = json.load(
                file
            )


    # ========================================================
    # Create image list
    # ========================================================

    images = sorted([
        file_name
        for file_name
        in os.listdir(dataset_path)
        if file_name.lower().endswith(
            (
                ".jpg",
                ".jpeg",
                ".png",
                ".webp"
            )
        )
    ])


    # ========================================================
    # Create metadata entries
    # ========================================================

    existing_images = {
        item["image"]
        for item in metadata
        if "image" in item
    }


    for image_name in images:

        if image_name not in existing_images:

            metadata.append(
                {
                    "image": image_name,
                    "qwen_processed": False
                }
            )


    start_time = time.time()


    # ========================================================
    # Process images
    # ========================================================

    for index, item in enumerate(
        metadata,
        start=1
    ):

        if item.get(
            "pipeline_processed"
        ) is True:

            continue


        image_path = os.path.join(
            dataset_path,
            item["image"]
        )


        try:

            if not os.path.exists(
                image_path
            ):

                raise FileNotFoundError(
                    f"Image not found: {image_path}"
                )


            image = Image.open(
                image_path
            ).convert("RGB")


            # ==================================================
            # 1. Qwen
            # ==================================================

            qwen_result = (
                analyze_image_with_qwen(
                    image_path
                )
            )


            item["caption"] = (
                qwen_result.get(
                    "description",
                    ""
                )
            )


            item["vlm_objects"] = (
                qwen_result.get(
                    "objects",
                    []
                )
            )


            item["visible_text"] = (
                qwen_result.get(
                    "visible_text",
                    []
                )
            )


            item["attributes"] = (
                qwen_result.get(
                    "attributes",
                    []
                )
            )


            item["relations"] = (
                qwen_result.get(
                    "relations",
                    []
                )
            )


            # ==================================================
            # 2. OCR
            # ==================================================

            ocr_result = (
                extract_text_with_paddleocr_vl_spotting(
                    image
                )
            )


            item["ocr_text"] = (
                ocr_result
            )


            # ==================================================
            # 3. RAM++ + Grounding DINO
            # ==================================================

            grounding_result = (
                detect_objects_with_ram_dino(
                    image_path
                )
            )


            item["objects"] = (
                grounding_result.get(
                    "objects",
                    []
                )
            )


            item["ram_tags"] = (
                grounding_result.get(
                    "ram_tags",
                    []
                )
            )


            item["detections"] = (
                grounding_result.get(
                    "detections",
                    []
                )
            )


            # ==================================================
            # Pipeline status
            # ==================================================

            item["pipeline_processed"] = True
            item["pipeline_error"] = None


        except Exception as error:

            item["pipeline_processed"] = False

            item["pipeline_error"] = str(
                error
            )


        print(
            f"{index}/{len(metadata)} "
            f"{item['image']}"
        )


        # ======================================================
        # Checkpoint
        # ======================================================

        if (
            checkpoint_path
            and index % 5 == 0
        ):

            with open(
                checkpoint_path,
                "w",
                encoding="utf-8"
            ) as file:

                json.dump(
                    metadata,
                    file,
                    indent=4,
                    ensure_ascii=False
                )


    # ========================================================
    # Final metadata
    # ========================================================

    with open(
        output_path,
        "w",
        encoding="utf-8"
    ) as file:

        json.dump(
            metadata,
            file,
            indent=4,
            ensure_ascii=False
        )


    elapsed = (
        time.time()
        - start_time
    ) / 60


    return {
        "metadata": metadata,

        "failed_images": [
            item
            for item in metadata
            if item.get(
                "pipeline_processed"
            ) is False
        ],

        "time_minutes": elapsed
    }