pipeline/prompt_mode.py

import uuid
import os
from models.utils import save_mask_as_png


def run_prompt_mode(image, prompt_text, sam_wrapper, dino_wrapper, save_dir="outputs"):
    """
    Process image using DINO + SAM based on user prompt.

    :param image: Input image (BGR format)
    :param prompt_text: Text prompt from user
    :param sam_wrapper: Initialized SamWrapper
    :param dino_wrapper: Initialized DinoWrapper
    :param clip_wrapper: Optional CLIP for verification
    :param save_dir: Directory to save results
    :return: List of saved file paths
    """
    saved_paths = []
    os.makedirs(save_dir, exist_ok=True)

    # 1. Detect boxes with DINO
    boxes = dino_wrapper.detect(image, prompt_text)

    # 2. Segment with SAM for each box
    for i, box in enumerate(boxes):
        mask = sam_wrapper.predict_with_box(image, box)
        if mask is None:
            continue  # safety check

        filename = f"{save_dir}/{uuid.uuid4().hex[:8]}_{i}_{prompt_text.replace(' ', '_')}.png"
        save_mask_as_png(image, mask, filename)
        saved_paths.append(filename)

    return saved_paths