Create app.py

app.py

@@ -0,0 +1,149 @@
+import os
+import json
+import math
+import random
+import shutil
+import tempfile
+from urllib.parse import urlparse
+from PIL import Image
+import cv2
+import numpy as np
+import gradio as gr
+from roboflow import Roboflow
+
+
+def parse_roboflow_url(url):
+    """Extract workspace, project name, and version from a Roboflow URL."""
+    parsed = urlparse(url)
+    parts = parsed.path.strip('/').split('/')
+    # Expect at least [workspace, project, ..., version]
+    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, dataset_url):
+    # Initialize Roboflow client
+    rf = Roboflow(api_key=api_key)
+    workspace, project_name, version = parse_roboflow_url(dataset_url)
+    project = rf.workspace(workspace).project(project_name)
+    version_obj = project.version(version)
+
+    # Download the segmentation dataset in COCO format
+    dataset = version_obj.download("coco-segmentation")
+    root = dataset.location  # root of downloaded dataset
+
+    # Load COCO train annotations
+    ann_dir = os.path.join(root, "coco-annotations")
+    ann_file = os.path.join(ann_dir, "train.json")
+    with open(ann_file, 'r') as f:
+        coco = json.load(f)
+    images_info = {img['id']: img for img in coco['images']}
+
+    # Map original category IDs to contiguous YOLO class indices
+    categories = coco.get('categories', [])
+    cat_ids = sorted(cat['id'] for cat in categories)
+    id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}
+
+    # Prepare output directories for YOLOv8 dataset
+    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)
+
+    # Build YOLO annotation strings grouped by image
+    annos = {}
+    for anno in coco['annotations']:
+        img_id = anno['image_id']
+        poly = anno['segmentation'][0]
+        xs = poly[0::2]
+        ys = poly[1::2]
+        x_min, x_max = min(xs), max(xs)
+        y_min, y_max = min(ys), max(ys)
+        width = x_max - x_min
+        height = y_max - y_min
+        cx = x_min + width / 2
+        cy = y_min + height / 2
+        info = images_info[img_id]
+        img_w, img_h = info['width'], info['height']
+        cxn = cx / img_w
+        cyn = cy / img_h
+        wnorm = width / img_w
+        hnorm = height / img_h
+        cls_idx = id_to_index[anno['category_id']]
+        line = f"{cls_idx} {cxn:.6f} {cyn:.6f} {wnorm:.6f} {hnorm:.6f}"
+        annos.setdefault(img_id, []).append(line)
+
+    # Determine train images directory
+    train_img_dir = os.path.join(root, "train", "images")
+    if not os.path.isdir(train_img_dir):
+        train_img_dir = os.path.join(root, "train")
+
+    # Map filenames to image IDs
+    name_to_id = {img['file_name']: img['id'] for img in coco['images']}
+
+    # Copy images and write YOLO label files
+    for fname, img_id in name_to_id.items():
+        src_img = os.path.join(train_img_dir, fname)
+        if not os.path.isfile(src_img):
+            continue
+        dst_img = os.path.join(img_out, fname)
+        shutil.copy(src_img, dst_img)
+        lbl_path = os.path.join(lbl_out, os.path.splitext(fname)[0] + ".txt")
+        with open(lbl_path, 'w') as lf:
+            for line in annos.get(img_id, []):
+                lf.write(line + '\n')
+
+    # Prepare before/after example galleries
+    before_imgs, after_imgs = [], []
+    example_files = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))
+
+    for fname in example_files:
+        src_img = os.path.join(train_img_dir, fname)
+        img = cv2.cvtColor(cv2.imread(src_img), cv2.COLOR_BGR2RGB)
+
+        # Overlay segmentation polygons
+        seg_vis = img.copy()
+        img_id = name_to_id[fname]
+        for anno in coco['annotations']:
+            if anno['image_id'] != img_id:
+                continue
+            poly = anno['segmentation'][0]
+            pts = np.array(poly, dtype=np.int32).reshape(-1, 2)
+            cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)
+
+        # Overlay bounding boxes
+        box_vis = img.copy()
+        for line in annos.get(img_id, []):
+            _, cxn, cyn, wnorm, hnorm = line.split()
+            cxn, cyn, wnorm, hnorm = map(float, (cxn, cyn, wnorm, hnorm))
+            iw, ih = images_info[img_id]['width'], images_info[img_id]['height']
+            w0 = int(wnorm * iw)
+            h0 = 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_imgs.append(Image.fromarray(seg_vis))
+        after_imgs.append(Image.fromarray(box_vis))
+
+    return before_imgs, after_imgs
+
+
+# Build Gradio interface
+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)").style(grid=[5], height="auto")
+    after_gallery = gr.Gallery(label="After (Bounding Boxes)").style(grid=[5], height="auto")
+    run_btn.click(convert_seg_to_bbox, inputs=[api_input, url_input], outputs=[before_gallery, after_gallery])
+
+if __name__ == "__main__":
+    app.launch()