run_local.py

import cv2
import os
from models.sam import SamWrapper
from models.dino import DinoWrapper
from pipeline.prompt_mode import run_prompt_mode


def main_prompt():
    """
    Run prompt-based segmentation using DINO + SAM.
    """
    # --- Ensure necessary folders exist
    os.makedirs("assets", exist_ok=True)
    os.makedirs("outputs", exist_ok=True)

    image_path = "assets/fir_tree.jpg"
    output_dir = "outputs"
    # prompt = "a vegetable with purple leaves"  # You can change this prompt
    # prompt = "a long light green vegetable"
    # prompt = "a round fruit with a red rough surface"
    # prompt = "a long green cucumber"
    # prompt = "a long green cucumber"
    # prompt = "a yellow car indoor"
    prompt = "a high fir tree"
    # prompt = "a car top view"

    image = cv2.imread(image_path)
    if image is None:
        print("Failed to load image.")
        return

    sam = SamWrapper(
        model_type="vit_b",
        checkpoint_path="weights/sam_vit_b_01ec64.pth"
    )
    dino = DinoWrapper()

    results = run_prompt_mode(image, prompt, sam, dino, save_dir=output_dir)


    print(f"\nProcessed {len(results)} objects (prompt mode):")
    for path in results:
        print(f"→ Saved at {path}")


if __name__ == "__main__":
    main_prompt()