app.py

import os
import json
import random
import shutil
import tempfile
from urllib.parse import urlparse

import cv2
import numpy as np
from PIL import Image
import gradio as gr
from roboflow import Roboflow


def parse_roboflow_url(url: str):
    """Extract workspace, project name, and version from a Roboflow URL."""
    parsed = urlparse(url)
    parts = parsed.path.strip('/').split('/')
    workspace = parts[0]
    project = parts[1]
    try:
        version = int(parts[-1])
    except ValueError:
        version = int(parts[-2])
    return workspace, project, version


def convert_seg_to_bbox(api_key: str, dataset_url: str):
    # 1) download the dataset
    rf = Roboflow(api_key=api_key)
    ws, proj, ver = parse_roboflow_url(dataset_url)
    version_obj = rf.workspace(ws).project(proj).version(ver)
    dataset = version_obj.download("coco-segmentation")
    root = dataset.location  # e.g. "/home/user/app/ds-2"

    # --- DEBUG: print out the downloaded directory tree ---
    print(f"\n=== Downloaded dataset root: {root} ===")
    for dirpath, dirnames, filenames in os.walk(root):
        print(f"\nDirectory: {dirpath}")
        for d in dirnames:
            print(f"  [DIR ] {d}")
        for f in filenames:
            print(f"  [FILE] {f}")
    print("=== end tree dump ===\n")

    # 2) search for any JSON file with "train" in its name
    ann_file = None
    for dirpath, _, filenames in os.walk(root):
        for fname in filenames:
            if 'train' in fname.lower() and fname.lower().endswith('.json'):
                ann_file = os.path.join(dirpath, fname)
                print(f"Found TRAIN annotation file: {ann_file}")
                break
        if ann_file:
            break

    # 2b) fallback: first .json anywhere
    if ann_file is None:
        for dirpath, _, filenames in os.walk(root):
            for fname in filenames:
                if fname.lower().endswith('.json'):
                    ann_file = os.path.join(dirpath, fname)
                    print(f"No TRAIN file—falling back to first JSON: {ann_file}")
                    break
            if ann_file:
                break

    if ann_file is None:
        raise FileNotFoundError(f"No JSON annotations found under {root}")

    # 3) load COCO annotations
    with open(ann_file, 'r') as f:
        coco = json.load(f)
    images_info = {img['id']: img for img in coco['images']}

    # 4) build category→index map
    cat_ids = sorted(c['id'] for c in coco.get('categories', []))
    id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}

    # 5) prepare YOLO output directories
    out_root = tempfile.mkdtemp(prefix="yolov8_")
    img_out = os.path.join(out_root, "images")
    lbl_out = os.path.join(out_root, "labels")
    os.makedirs(img_out, exist_ok=True)
    os.makedirs(lbl_out, exist_ok=True)
    print(f"Preparing YOLOv8 output in: {out_root}")

    # 6) convert each segmentation annotation to a YOLO bbox line
    annos = {}
    for anno in coco['annotations']:
        img_id = anno['image_id']
        poly = anno['segmentation'][0]
        xs, ys = poly[0::2], poly[1::2]
        x_min, x_max = min(xs), max(xs)
        y_min, y_max = min(ys), max(ys)
        w, h = x_max - x_min, y_max - y_min
        cx, cy = x_min + w/2, y_min + h/2

        info = images_info[img_id]
        iw, ih = info['width'], info['height']
        line = (
            f"{id_to_index[anno['category_id']]} "
            f"{cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
        )
        annos.setdefault(img_id, []).append(line)

    # 7) locate your image folder (first with any .jpg/.png)
    train_img_dir = None
    for dirpath, _, files in os.walk(root):
        if any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in files):
            train_img_dir = dirpath
            print(f"Found image directory: {train_img_dir}")
            break

    if train_img_dir is None:
        raise FileNotFoundError(f"No image files found under {root}")

    # 8) copy images + write labels
    name_to_id = {img['file_name']: img['id'] for img in coco['images']}
    for fname, img_id in name_to_id.items():
        src = os.path.join(train_img_dir, fname)
        if not os.path.isfile(src):
            continue
        shutil.copy(src, os.path.join(img_out, fname))
        lbl_path = os.path.join(lbl_out, fname.rsplit('.', 1)[0] + ".txt")
        with open(lbl_path, 'w') as lf:
            lf.write("\n".join(annos.get(img_id, [])))

    # 9) build before/after galleries
    before, after = [], []
    sample = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))
    for fname in sample:
        src = os.path.join(train_img_dir, fname)
        img = cv2.cvtColor(cv2.imread(src), cv2.COLOR_BGR2RGB)

        seg_vis = img.copy()
        img_id = name_to_id[fname]
        for anno in coco['annotations']:
            if anno['image_id'] != img_id: continue
            pts = np.array(anno['segmentation'][0], dtype=np.int32).reshape(-1, 2)
            cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)

        box_vis = img.copy()
        for line in annos.get(img_id, []):
            _, cxn, cyn, wnorm, hnorm = map(float, line.split())
            iw, ih = images_info[img_id]['width'], images_info[img_id]['height']
            w0, h0 = int(wnorm * iw), int(hnorm * ih)
            x0 = int(cxn * iw - w0/2)
            y0 = int(cyn * ih - h0/2)
            cv2.rectangle(box_vis, (x0, y0), (x0 + w0, y0 + h0), (0, 255, 0), 2)

        before.append(Image.fromarray(seg_vis))
        after.append(Image.fromarray(box_vis))

    return before, after


# --- Gradio app ---
with gr.Blocks() as app:
    gr.Markdown("# Segmentation → YOLOv8 Converter")
    api_input      = gr.Textbox(label="Roboflow API Key", type="password")
    url_input      = gr.Textbox(label="Roboflow Dataset URL (Segmentation)")
    run_btn        = gr.Button("Convert")
    before_gallery = gr.Gallery(label="Before (Segmentation)", columns=5, height="auto")
    after_gallery  = gr.Gallery(label="After (Bounding Boxes)", columns=5, height="auto")

    run_btn.click(
        fn=convert_seg_to_bbox,
        inputs=[api_input, url_input],
        outputs=[before_gallery, after_gallery]
    )

if __name__ == "__main__":
    app.launch()