script.py

import cv2
import os
import subprocess
from PIL import Image
import easyocr
from spellchecker import SpellChecker
import numpy as np
import webcolors
from collections import Counter
import torch
from transformers import AutoProcessor, Blip2ForConditionalGeneration
import tensorflow as tf
import argparse
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
from utils.json_helpers import NoIndent, CustomEncoder


# constants
BARRIER = "********\n"

# Check if a model is in the cache
def is_model_downloaded(model_name, cache_directory):
    model_path = os.path.join(cache_directory, model_name.replace('/', '_'))
    return os.path.exists(model_path)

# Convert color to the closest name
def closest_colour(requested_colour):
    min_colours = {}
    css3_names = webcolors.names("css3")
    for name in css3_names:
        hex_value = webcolors.name_to_hex(name, spec='css3')
        r_c, g_c, b_c = webcolors.hex_to_rgb(hex_value)
        rd = (r_c - requested_colour[0]) ** 2
        gd = (g_c - requested_colour[1]) ** 2
        bd = (b_c - requested_colour[2]) ** 2
        distance = rd + gd + bd
        min_colours[distance] = name
    return min_colours[min(min_colours.keys())]

def get_colour_name(requested_colour):
    """
    Returns a tuple: (exact_name, closest_name).
    If an exact match fails, 'exact_name' is None, use the 'closest_name' fallback.
    """
    try:
        actual_name = webcolors.rgb_to_name(requested_colour, spec='css3')
        closest_name = actual_name
    except ValueError:
        closest_name = closest_colour(requested_colour)
        actual_name = None
    return actual_name, closest_name

def get_most_frequent_color(pixels, bin_size=10):
    """
    Returns the most frequent color among the given pixels,
    using a binning approach (default bin size=10).
    """
    bins = np.arange(0, 257, bin_size)
    r_bins = np.digitize(pixels[:, 0], bins) - 1
    g_bins = np.digitize(pixels[:, 1], bins) - 1
    b_bins = np.digitize(pixels[:, 2], bins) - 1
    combined_bins = r_bins * 10000 + g_bins * 100 + b_bins
    bin_counts = Counter(combined_bins)
    most_common_bin = bin_counts.most_common(1)[0][0]

    r_bin = most_common_bin // 10000
    g_bin = (most_common_bin % 10000) // 100
    b_bin = most_common_bin % 100
    r_value = bins[r_bin] + bin_size // 2
    g_value = bins[g_bin] + bin_size // 2
    b_value = bins[b_bin] + bin_size // 2

    return (r_value, g_value, b_value)

def get_most_frequent_alpha(alphas, bin_size=10):
    bins = np.arange(0, 257, bin_size)
    alpha_bins = np.digitize(alphas, bins) - 1
    bin_counts = Counter(alpha_bins)
    most_common_bin = bin_counts.most_common(1)[0][0]
    alpha_value = bins[most_common_bin] + bin_size // 2
    return alpha_value

# downscale images for OCR. TODO change dim to a suitable one
def downscale_for_ocr(image_cv, max_dim=600):
    """
    If either dimension of `image_cv` is bigger than `max_dim`,
    scale it down proportionally. This speeds up EasyOCR on large images.
    """
    h, w = image_cv.shape[:2]
    if w <= max_dim and h <= max_dim:
        return image_cv  # No downscale needed

    scale = min(max_dim / float(w), max_dim / float(h))
    new_w = int(w * scale)
    new_h = int(h * scale)
    image_resized = cv2.resize(image_cv, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return image_resized

# Worker function to process a single bounding box
def process_single_region(
    idx, bounding_box, image, sr, reader, spell, icon_model,
    processor, model, device, no_captioning, output_json, json_mini,
    cropped_imageview_images_dir, base_name, save_images,
    model_to_use, log_prefix="",
    skip_ocr=False,
    skip_spell=False
):
    """
    Processes one bounding box (region)
    Returns a dict with:
      * "region_dict" (for JSON)
      * "text_log" (file/captions output)
    """
    (x_min, y_min, x_max, y_max, class_id) = bounding_box
    class_names = {0: 'View', 1: 'ImageView', 2: 'Text', 3: 'Line'}
    class_name = class_names.get(class_id, f'Unknown Class {class_id}')
    region_idx = idx + 1
    logs = []

    x_center = (x_min + x_max) // 2
    y_center = (y_min + y_max) // 2
    width = x_max - x_min
    height = y_max - y_min

    def open_and_upscale_image(img_path, cid):

        if cid == 2: # Text 
            MAX_WIDTH, MAX_HEIGHT = 30, 30
        else:
            MAX_WIDTH, MAX_HEIGHT = 10, 10

        def is_small(w, h):
            return w <= MAX_WIDTH and h <= MAX_HEIGHT

        if cid == 0:  # "View" - use PIL to preserve alpha
            pil_image = Image.open(img_path).convert("RGBA")
            w, h = pil_image.size
            if not is_small(w, h):
                logs.append(f"{log_prefix}Skipping upscale for large View (size={w}×{h}).")
                return pil_image

            # If super-resolution is provided, use it
            if sr:
                image_cv = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGBA2BGR)
                upscaled = sr.upsample(image_cv)
                return Image.fromarray(cv2.cvtColor(upscaled, cv2.COLOR_BGR2RGBA))
            else:
                return pil_image.resize((w * 4, h * 4), resample=Image.BICUBIC)
        else:
            # For other classes, load the image with OpenCV (BGR)
            cv_image = cv2.imread(img_path)
            if cv_image is None or cv_image.size == 0:
                logs.append(f"{log_prefix}Empty image at {img_path}, skipping.")
                return None

            h, w = cv_image.shape[:2]
            if not is_small(w, h):
                logs.append(f"{log_prefix}Skipping upscale for large region (size={w}×{h}).")
                return cv_image

            if sr:
                return sr.upsample(cv_image)
            else:
                return cv2.resize(cv_image, (w * 2, h * 2), interpolation=cv2.INTER_CUBIC)

    if json_mini:
        simplified_class_name = class_name.lower().replace('imageview', 'image')
        new_id = f"{simplified_class_name}_{region_idx}"
        mini_region_dict = {
            "id": new_id,
            "bbox": NoIndent([x_center, y_center, width, height])
        }
        
        # only need to process text for the mini format
        if class_name == 'Text' and not skip_ocr:
            cropped_image_region = image[y_min:y_max, x_min:x_max]
            if cropped_image_region.size > 0:
                # Save the cropped image so open_and_upscale_image can use it
                cropped_path = os.path.join(cropped_imageview_images_dir, f"region_{region_idx}_class_{class_id}.jpg")
                cv2.imwrite(cropped_path, cropped_image_region)

                upscaled = open_and_upscale_image(cropped_path, class_id)
                if upscaled is not None:
                    if isinstance(upscaled, Image.Image):
                        upscaled_cv = cv2.cvtColor(np.array(upscaled), cv2.COLOR_RGBA2BGR)
                    else:
                        upscaled_cv = upscaled
                    
                    gray = cv2.cvtColor(downscale_for_ocr(upscaled_cv), cv2.COLOR_BGR2GRAY)
                    text = ' '.join(reader.readtext(gray, detail=0, batch_size=8)).strip()

                    if text:
                        if not skip_spell and spell:
                            corrected_words  = []
                            for w in text.split():
                                corrected_words.append(spell.correction(w) or w)
                            mini_region_dict["text"] = " ".join(corrected_words)
                        else:
                             mini_region_dict["text"] = text
                
                # Clean up the temporary cropped image
                if os.path.exists(cropped_path) and not save_images:
                    os.remove(cropped_path)

        return {"mini_region_dict": mini_region_dict, "text_log": ""}

    logs.append(f"\n{log_prefix}Region {region_idx} - Class ID: {class_id} ({class_name})")
    x_center = (x_min + x_max) // 2
    y_center = (y_min + y_max) // 2
    logs.append(f"{log_prefix}Coordinates: x_center={x_center}, y_center={y_center}")
    width = x_max - x_min
    height = y_max - y_min
    logs.append(f"{log_prefix}Size: width={width}, height={height}")

    region_dict = {
        "id": f"region_{region_idx}_class_{class_name}",
        "x_coordinates_center": x_center,
        "y_coordinates_center": y_center,
        "width": width,
        "height": height
    }

    # Crop region
    cropped_image_region = image[y_min:y_max, x_min:x_max]
    if cropped_image_region.size == 0:
        logs.append(f"{log_prefix}Empty crop for Region {region_idx}, skipping...")
        return {"region_dict": region_dict, "text_log": "\n".join(logs)}

    # Save cropped region
    if class_id == 0:
        # Save as PNG if it's a View
        cropped_path = os.path.join(
            cropped_imageview_images_dir, f"region_{region_idx}_class_{class_id}.png"
        )
        cv2.imwrite(cropped_path, cropped_image_region)
    else:
        # Save as JPG
        cropped_path = os.path.join(
            cropped_imageview_images_dir, f"region_{region_idx}_class_{class_id}.jpg"
        )
        cv2.imwrite(cropped_path, cropped_image_region)

    # for LLaMA (ollama)
    def call_ollama(prompt_text, rid, task_type):
        model_name = "llama3.2-vision:11b"
        cmd = ["ollama", "run", model_name, prompt_text]
        try:
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if result.returncode != 0:
                logs.append(f"{log_prefix}Error generating {task_type} for Region {rid}: {result.stderr}")
                return None
            else:
                response = result.stdout.strip()
                logs.append(f"{log_prefix}Generated {task_type.capitalize()} for Region {rid}: {response}")
                return response
        except Exception as e:
            logs.append(f"{log_prefix}An error occurred while generating {task_type} for Region {rid}: {e}")
            return None

    # for BLIP-2
    def generate_caption_blip(img_path):
        pil_image = Image.open(img_path).convert('RGB')
        inputs = processor(images=pil_image, return_tensors="pt").to(device, torch.float16)
        gen_ids = model.generate(**inputs)
        return processor.batch_decode(gen_ids, skip_special_tokens=True)[0].strip()

    # Handle each class type
    if class_id == 1:  # ImageView
        if no_captioning:
            logs.append(f"{log_prefix}(Icon-image detection + captioning disabled by --no-captioning.)")
            if not output_json:
                block = (
                    f"Image: region_{region_idx}_class_{class_id} ({class_name})\n"
                    f"Coordinates: x_center={(x_min + x_max) // 2}, y_center={(y_min + y_max) // 2}\n"
                    f"Size: width={width}, height={height}\n"
                    f"{BARRIER}"
                )
                logs.append(block)
        else:
            upscaled = open_and_upscale_image(cropped_path, class_id)
            if upscaled is None:
                return {"region_dict": region_dict, "text_log": "\n".join(logs)}

            # Icon detection
            if icon_model:
                icon_input_size = (224, 224)
                if isinstance(upscaled, Image.Image):
                    upscaled_cv = cv2.cvtColor(np.array(upscaled), cv2.COLOR_RGBA2BGR)
                else:
                    upscaled_cv = upscaled
                resized = cv2.resize(upscaled_cv, icon_input_size)
                rgb_img = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) / 255.0
                rgb_img = np.expand_dims(rgb_img, axis=0)
                pred = icon_model.predict(rgb_img)
                logs.append(f"{log_prefix}Prediction output for Region {region_idx}: {pred}")
                if pred.shape == (1, 1):
                    probability = pred[0][0]
                    threshold = 0.5
                    predicted_class = 1 if probability >= threshold else 0
                    logs.append(f"{log_prefix}Probability of class 1: {probability}")
                elif pred.shape == (1, 2):
                    predicted_class = np.argmax(pred[0])
                    logs.append(f"{log_prefix}Class probabilities: {pred[0]}")
                else:
                    logs.append(f"{log_prefix}Unexpected prediction shape: {pred.shape}")
                    return {"region_dict": region_dict, "text_log": "\n".join(logs)}

                pred_text = "Icon/Mobile UI Element" if predicted_class == 1 else "Normal Image"
                region_dict["prediction"] = pred_text
                if predicted_class == 1:
                    prompt_text = "Describe the mobile UI element on this image. Keep it short."
                else:
                    prompt_text = "Describe what is in the image briefly. It's not an icon or typical UI element."
            else:
                logs.append(f"{log_prefix}Icon detection model not provided; skipping icon detection.")
                region_dict["prediction"] = "Icon detection skipped"
                prompt_text = "Describe what is in this image briefly."

            # Caption
            temp_image_path = os.path.abspath(
                os.path.join(cropped_imageview_images_dir, f"imageview_{region_idx}.jpg")
            )
            if isinstance(upscaled, Image.Image): # TODO check optimization
                upscaled.save(temp_image_path)
            else:
                cv2.imwrite(temp_image_path, upscaled)

            response = ""
            if model and processor and model_to_use == 'blip':
                response = generate_caption_blip(temp_image_path)
            else: # TODO check optimization
                resp = call_ollama(prompt_text + " " + temp_image_path, region_idx, "description")
                response = resp if resp else "Error generating description"

            region_dict["description"] = response

            if not output_json:
                block = (
                    f"Image: region_{region_idx}_class_{class_id} ({class_name})\n"
                    f"Coordinates: x_center={(x_min + x_max) // 2}, y_center={(y_min + y_max) // 2}\n"
                    f"Size: width={width}, height={height}\n"
                    f"Prediction: {region_dict['prediction']}\n"
                    f"{response}\n"
                    f"{BARRIER}"
                )
                logs.append(block)

            if os.path.exists(temp_image_path) and not save_images:
                os.remove(temp_image_path)

    elif class_id == 2:  # Text
        if skip_ocr or reader is None:
            logs.append(f"{log_prefix}OCR skipped for Region {region_idx}.")

            if not output_json:
                block = (
                    f"Text: region_{region_idx}_class_{class_id} ({class_name})\n"
                    f"Coordinates: x_center={(x_min + x_max) // 2}, "
                    f"y_center={(y_min + y_max) // 2}\n"
                    f"Size: width={width}, height={height}\n"
                    f"OCR + spell-check disabled\n"
                    f"{BARRIER}"
                )
                logs.append(block)

            return {"region_dict": region_dict, "text_log": "\n".join(logs)}

        upscaled = open_and_upscale_image(cropped_path, class_id)
        if upscaled is None:
            return {"region_dict": region_dict, "text_log": "\n".join(logs)}

        if isinstance(upscaled, Image.Image):
            upscaled_cv = cv2.cvtColor(np.array(upscaled), cv2.COLOR_RGBA2BGR)
        else:
            upscaled_cv = upscaled


        # TODO use other lib to improve the performance
        upscaled_cv = downscale_for_ocr(upscaled_cv, max_dim=600)
        gray = cv2.cvtColor(upscaled_cv, cv2.COLOR_BGR2GRAY)
        result_ocr = reader.readtext(gray, detail=0, batch_size=8)
        text = ' '.join(result_ocr).strip()

        # TODO use other lib to improve performance
        if skip_spell or spell is None:
            corrected_text = None
            logs.append(f"{log_prefix}Spell-check skipped for Region {region_idx}.")
        else:
            correction_cache = {}
            corrected_words  = []
            for w in text.split():
                if w not in correction_cache:
                    correction_cache[w] = spell.correction(w) or w
                corrected_words.append(correction_cache[w])
            corrected_text = " ".join(corrected_words)


        logs.append(f"{log_prefix}Extracted Text for Region {region_idx}: {text}")
        if corrected_text is not None:
            logs.append(f"{log_prefix}Corrected Text for Region {region_idx}: {corrected_text}")

        region_dict["extractedText"] = text
        if corrected_text is not None:
            region_dict["correctedText"] = corrected_text

        if not output_json:
            block = (
                f"Text: region_{region_idx}_class_{class_id} ({class_name})\n"
                f"Coordinates: x_center={(x_min + x_max) // 2}, y_center={(y_min + y_max) // 2}\n"
                f"Size: width={width}, height={height}\n"
                f"Extracted Text: {text}\n"
                + (f"Corrected Text: {corrected_text}\n" if corrected_text is not None else "")
                + f"{BARRIER}"
            )
            logs.append(block)

    elif class_id == 0:  # View
        upscaled = open_and_upscale_image(cropped_path, class_id)
        if upscaled is None:
            return {"region_dict": region_dict, "text_log": "\n".join(logs)}

        data = np.array(upscaled)
        if data.ndim == 2:
            data = cv2.cvtColor(data, cv2.COLOR_GRAY2BGRA)
        elif data.shape[-1] == 3:
            b, g, r = cv2.split(data)
            a = np.full_like(b, 255)
            data = cv2.merge((b, g, r, a))

        pixels = data.reshape((-1, 4))
        opaque_pixels = pixels[pixels[:, 3] > 0]

        if len(opaque_pixels) == 0:
            logs.append(f"{log_prefix}No opaque pixels found in Region {region_idx}, cannot determine background color.")
            color_name = "Unknown"
        else:
            dom_color = get_most_frequent_color(opaque_pixels[:, :3], bin_size=10)
            exact_name, closest_name = get_colour_name(dom_color)
            color_name = exact_name if exact_name else closest_name

        alphas = pixels[:, 3]
        dominant_alpha = get_most_frequent_alpha(alphas, bin_size=10)
        transparency = "opaque" if dominant_alpha >= 245 else "transparent"

        response = (
            f"1. The background color of the container is {color_name}.\n"
            f"2. The container is {transparency}."
        )
        logs.append(f"{log_prefix}{response}")
        region_dict["view_color"] = f"The background color of the container is {color_name}."
        region_dict["view_alpha"] = f"The container is {transparency}."

        if not output_json:
            block = (
                f"View: region_{region_idx}_class_{class_id} ({class_name})\n"
                f"Coordinates: x_center={(x_min + x_max) // 2}, y_center={(y_min + y_max) // 2}\n"
                f"Size: width={width}, height={height}\n"
                f"{response}\n"
                f"{BARRIER}"
            )
            logs.append(block)

    elif class_id == 3:  # Line
        logs.append(f"{log_prefix}Processing Line in Region {region_idx}")
        line_img = cv2.imread(cropped_path, cv2.IMREAD_UNCHANGED)
        if line_img is None:
            logs.append(f"{log_prefix}Failed to read image at {cropped_path}")
            return {"region_dict": region_dict, "text_log": "\n".join(logs)}

        hh, ww = line_img.shape[:2]
        logs.append(f"{log_prefix}Image dimensions: width={ww}, height={hh}")

        data = np.array(line_img)
        if data.ndim == 2:
            data = cv2.cvtColor(data, cv2.COLOR_GRAY2BGRA)
        elif data.shape[-1] == 3:
            b, g, r = cv2.split(data)
            a = np.full_like(b, 255)
            data = cv2.merge((b, g, r, a))

        pixels = data.reshape((-1, 4))
        opaque_pixels = pixels[pixels[:, 3] > 0]

        if len(opaque_pixels) == 0:
            logs.append(f"{log_prefix}No opaque pixels found in Region {region_idx}, cannot determine line color.")
            color_name = "Unknown"
        else:
            dom_color = get_most_frequent_color(opaque_pixels[:, :3], bin_size=10)
            exact_name, closest_name = get_colour_name(dom_color)
            color_name = exact_name if exact_name else closest_name

        alphas = pixels[:, 3]
        dom_alpha = get_most_frequent_alpha(alphas, bin_size=10)
        transparency = "opaque" if dom_alpha >= 245 else "transparent"

        response = (
            f"1. The color of the line is {color_name}.\n"
            f"2. The line is {transparency}."
        )
        logs.append(f"{log_prefix}{response}")
        region_dict["line_color"] = f"The color of the line is {color_name}."
        region_dict["line_alpha"] = f"The line is {transparency}."

        if not output_json:
            block = (
                f"Line: region_{region_idx}_class_{class_id} ({class_name})\n"
                f"Coordinates: x_center={(x_min + x_max) // 2}, y_center={(y_min + y_max) // 2}\n"
                f"Size: width={width}, height={height}\n"
                f"{response}\n"
                f"{BARRIER}"
            )
            logs.append(block)

    else:
        logs.append(f"{log_prefix}Class ID {class_id} not handled.")

    # Remove intermediate if not saving
    if os.path.exists(cropped_path) and not save_images:
        os.remove(cropped_path)

    return {
        "region_dict": region_dict,
        "text_log": "\n".join(logs),
    }


# Main function
def process_image(
    input_image_path,
    yolo_output_path,
    output_dir:str = '.',
    model_to_use='llama',     
    save_images=False,
    icon_model_path=None,
    cache_directory='./models_cache',
    huggingface_token='your_token', # for blip2
    no_captioning=False,
    output_json=False,
    json_mini=False,
    sr=None,
    reader=None,
    spell=None,
    skip_ocr=False,
    skip_spell=False
):
    if json_mini:
        json_output = {
            "image_size": None, # Will be populated later
            "bbox_format": "center_x, center_y, width, height",
            "elements": []
        }
    elif output_json:
        json_output = {
            "image": {"path": input_image_path, "size": {"width": None, "height": None}},
            "elements": []
        }
    else:
        json_output = None


    start_time = time.perf_counter()
    print("super-resolution initialization start (in script.py)")
    # Super-resolution initialization
    if sr is None:
        print("No sr reference passed; performing local init ...")
        model_path = 'EDSR_x4.pb'
        if hasattr(cv2, 'dnn_superres'):
            print("dnn_superres module is available.")
            import cv2.dnn_superres as dnn_superres
            try:
                sr = cv2.dnn_superres.DnnSuperResImpl_create()
                print("Using DnnSuperResImpl_create()")
            except AttributeError:
                sr = cv2.dnn_superres.DnnSuperResImpl()
                print("Using DnnSuperResImpl()")
            sr.readModel(model_path)
            sr.setModel('edsr', 4)
        else:
            print("dnn_superres module is NOT available; skipping super-resolution.")
    else:
        print("Using pre-initialized sr reference.")


    elapsed = time.perf_counter() - start_time
    print(f"super-resoulution init (in script.py) took {elapsed:.3f} seconds.")

    start_time = time.perf_counter()

    if skip_ocr:
        print("skip_ocr flag set - skipping EasyOCR and SpellChecker.")
        reader = None
        spell  = None

    else:
        print("OCR initialisation start (in script.py)")
        if reader is None:
            print("No EasyOCR reference passed; performing local init")
            reader = easyocr.Reader(['en'], gpu=True)
        else:
            print("Using pre-initialised EasyOCR object.")

        if skip_spell:
            print("skip_spell flag set - not initialising SpellChecker.")
            spell = None
        else:
            if spell is None:
                print("No SpellChecker reference passed; performing local init")
                spell = SpellChecker()
            else:
                print("Using pre-initialised SpellChecker object.")

    elapsed = time.perf_counter() - start_time
    print(f"OCR init (in script.py) took {elapsed:.3f} seconds.")


    start_time = time.perf_counter()
    print("icon-model init start (in script.py)")
    # Load icon detection model (if provided)
    if icon_model_path:
        icon_model = tf.keras.models.load_model(icon_model_path)
        print(f"Icon detection model loaded: {icon_model_path}")
    else:
        icon_model = None

    elapsed = time.perf_counter() - start_time
    print(f"icon-model init (in script.py) took {elapsed:.3f} seconds.")

    
    # Load the original image
    image = cv2.imread(input_image_path, cv2.IMREAD_UNCHANGED)
    if image is None:
        print(f"Image at {input_image_path} could not be loaded.")
        return

    image_height, image_width = image.shape[:2]


    # Read YOLO labels
    with open(yolo_output_path, 'r') as f:
        lines = f.readlines()

    # Check torch device
    if torch.backends.mps.is_available():
        device = torch.device("mps")
        print("Using MPS")
    elif torch.cuda.is_available():
        device = torch.device("cuda")
        print("Using CUDA")
    else:
        device = torch.device("cpu")
        print("Using CPU")

    # Conditionally load the captioning model
    processor, model = None, None
    if not no_captioning:
        if model_to_use == 'blip':
            print("Loading BLIP-2 model...")
            blip_model_name = "Salesforce/blip2-opt-2.7b"
            if not is_model_downloaded(blip_model_name, cache_directory):
                print("Model not found in cache. Downloading...")
            else:
                print("Model found in cache. Loading...")
            processor = AutoProcessor.from_pretrained(
                blip_model_name,
                use_auth_token=huggingface_token,
                cache_dir=cache_directory,
                resume_download=True
            )
            model = Blip2ForConditionalGeneration.from_pretrained(
                blip_model_name,
                device_map='auto',
                torch_dtype=torch.float16,
                use_auth_token=huggingface_token,
                cache_dir=cache_directory,
                resume_download=True
            ).to(device)
        else:
            print("Using LLaMA model via external call (ollama).")
    else:
        print("--no-captioning flag is set; skipping model loading.")

    # Prepare bounding boxes from YOLO
    bounding_boxes = []
    for line in lines:
        parts = line.strip().split()
        class_id = int(parts[0])
        x_center_norm, y_center_norm, width_norm, height_norm = map(float, parts[1:])
        x_center = x_center_norm * image_width
        y_center = y_center_norm * image_height
        box_width = width_norm * image_width
        box_height = height_norm * image_height
        x_min = int(x_center - box_width / 2)
        y_min = int(y_center - box_height / 2)
        x_max = int(x_center + box_width / 2)
        y_max = int(y_center + box_height / 2)
        x_min = max(0, x_min)
        y_min = max(0, y_min)
        x_max = min(image_width - 1, x_max)
        y_max = min(image_height - 1, y_max)
        bounding_boxes.append((x_min, y_min, x_max, y_max, class_id))


    # Create output dirs
    cropped_dir = os.path.join(output_dir, "cropped_imageview_images")
    os.makedirs(cropped_dir, exist_ok=True)
    result_dir = os.path.join(output_dir, "result")
    os.makedirs(result_dir, exist_ok=True)

    base_name = os.path.splitext(os.path.basename(input_image_path))[0]
    captions_file_path = None
    if json_mini:
        json_output["image_size"] = NoIndent([image_width, image_height])
    elif output_json:
        json_output["image"]["size"]["width"] = image_width
        json_output["image"]["size"]["height"] = image_height
    else: # Text output
        captions_filename = f"{base_name}_regions_captions.txt"
        captions_file_path = os.path.join(result_dir, captions_filename)
        with open(captions_file_path, 'w', encoding='utf-8') as f:
            f.write(f"Image path: {input_image_path}\n")
            f.write(f"Image Size: width={image_width}, height={image_height}\n")
            f.write(BARRIER)

    # Number of workers can be increased if hardware is suitable for it. But testing is needed
    start_time = time.perf_counter()
    print("Process single region start (in script.py)")

    with ThreadPoolExecutor(max_workers=1) as executor:
        futures = [
            executor.submit(
                process_single_region,
                idx, box, image, sr, reader, spell,
                icon_model, processor, model, (model and device),
                no_captioning, output_json, json_mini,
                cropped_dir, base_name, save_images,
                model_to_use, log_prefix="",
                skip_ocr=skip_ocr,
                skip_spell=skip_spell
            ) for idx, box in enumerate(bounding_boxes)
        ]

        for future in as_completed(futures):
            item = future.result()
            if json_mini:
                if item.get("mini_region_dict"):
                    json_output["elements"].append(item["mini_region_dict"])
            elif output_json:
                if item.get("region_dict"):
                    json_output["elements"].append(item["region_dict"])
            else: # Text output
                if item.get("text_log") and captions_file_path:
                    with open(captions_file_path, 'a', encoding='utf-8') as f:
                        f.write(item["text_log"])

    elapsed = time.perf_counter() - start_time
    print(f"Processing regions took {elapsed:.3f} seconds.")

    if json_mini or output_json:
        json_file = os.path.join(result_dir, f"{base_name}.json")
        with open(json_file, 'w', encoding='utf-8') as f:
            json.dump(json_output, f, indent=2, ensure_ascii=False, cls=CustomEncoder)
        
        output_type = "mini JSON" if json_mini else "JSON"
        print(f"{output_type} output written to {json_file}")
    else:
        print(f"Text output written to {captions_file_path}")
    

# CLI entry point
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process an image and its YOLO labels.')
    parser.add_argument('input_image', help='Path to the input YOLO image.')
    parser.add_argument('input_labels', help='Path to the input YOLO labels file.')
    parser.add_argument('--output_dir', default='.', 
                        help='Directory to save output files. Defaults to the current directory.')
    parser.add_argument('--model_to_use', choices=['llama', 'blip'], default='llama',
                        help='Model for captioning (llama or blip).')
    parser.add_argument('--save_images', action='store_true',
                        help='Flag to save intermediate images.')
    parser.add_argument('--icon_detection_path', help='Path to icon detection model.')
    parser.add_argument('--cache_directory', default='./models_cache',
                        help='Cache directory for Hugging Face models.')
    parser.add_argument('--huggingface_token', default='your_token',
                        help='Hugging Face token for model downloads.')
    parser.add_argument('--no-captioning', action='store_true',
                        help='Disable any image captioning.')
    parser.add_argument('--json', dest='output_json', action='store_true',
                        help='Output the image data in JSON format')
    parser.add_argument('--json-mini', action='store_true',
                        help='Output the image data in a condensed JSON format')
    args = parser.parse_args()

    process_image(
        input_image_path=args.input_image,
        yolo_output_path=args.input_labels,
        output_dir=args.output_dir,
        model_to_use=args.model_to_use,
        save_images=args.save_images,
        icon_model_path=args.icon_detection_path,
        cache_directory=args.cache_directory,
        huggingface_token=args.huggingface_token,
        no_captioning=args.no_captioning,
        output_json=args.output_json,
        json_mini=args.json_mini
    )