Update app.py

app.py

@@ -1,18 +1,18 @@
 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
+from PIL import Image
 import gradio as gr
 from roboflow import Roboflow
 
 
-def parse_roboflow_url(url):
+def parse_roboflow_url(url: str):
     """Extract workspace, project name, and version from a Roboflow URL."""
     parsed = urlparse(url)
     parts = parsed.path.strip('/').split('/')
@@ -25,47 +25,43 @@ def parse_roboflow_url(url):
     return workspace, project, version
 
 
-def convert_seg_to_bbox(api_key, dataset_url):
-    # Initialize Roboflow client & download
+def convert_seg_to_bbox(api_key: str, dataset_url: str):
+    # 1) download the dataset
     rf = Roboflow(api_key=api_key)
-    workspace, project_name, version = parse_roboflow_url(dataset_url)
-    version_obj = rf.workspace(workspace).project(project_name).version(version)
+    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
+    root = dataset.location  # e.g. "/home/user/app/YourProject"
 
-    # Locate annotations file
+    # 2) find the train JSON (anything with "train" in the name)
     ann_file = None
-    for sub in ("coco-annotations", "annotations"):
-        candidate = os.path.join(root, sub, "train.json")
-        if os.path.isfile(candidate):
-            ann_file = candidate
-            break
-    if ann_file is None:
-        # fallback: walk
-        for dp, _, files in os.walk(root):
-            if "train.json" in files:
-                ann_file = os.path.join(dp, "train.json")
+    for dp, _, files in os.walk(root):
+        for fname in files:
+            if fname.lower().endswith(".json") and "train" in fname.lower():
+                ann_file = os.path.join(dp, fname)
                 break
+        if ann_file:
+            break
     if ann_file is None:
-        raise FileNotFoundError(f"Could not find train.json under {root}")
+        raise FileNotFoundError(f"Could not locate any '*train*.json' under {root}")
 
-    # Load COCO annotations
+    # 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']}
 
-    # Map category IDs → YOLO class indices
+    # 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)}
 
-    # Prepare YOLO output dirs
+    # 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)
 
-    # Build YOLO annotations per image
+    # 6) convert each segmentation annotation to a YOLO bbox line
     annos = {}
     for anno in coco['annotations']:
         img_id = anno['image_id']
@@ -75,79 +71,73 @@ def convert_seg_to_bbox(api_key, dataset_url):
         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]
-        img_w, img_h = info['width'], info['height']
-        line = f"{id_to_index[anno['category_id']]} " \
-               f"{cx/img_w:.6f} {cy/img_h:.6f} {w/img_w:.6f} {h/img_h:.6f}"
+        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)
 
-    # Locate train images directory
-    possible = [
-        os.path.join(root, "train", "images"),
-        os.path.join(root, "train"),
-        os.path.join(root, "images")
-    ]
+    # 7) locate your train-images folder (first folder under root with any .jpg/.png)
     train_img_dir = None
-    for p in possible:
-        if os.path.isdir(p):
-            train_img_dir = p
+    for dp, _, files in os.walk(root):
+        if any(f.lower().endswith((".jpg", ".png", ".jpeg")) for f in files):
+            train_img_dir = dp
             break
     if train_img_dir is None:
-        for dp, _, files in os.walk(root):
-            if any(f.lower().endswith(('.jpg','.png','.jpeg')) for f in files):
-                train_img_dir = dp
-                break
-    if train_img_dir is None:
-        raise FileNotFoundError(f"No image directory found under {root}")
+        raise FileNotFoundError(f"No image files found under {root}")
 
-    # Map filenames → IDs, copy images & write labels
+    # 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))
-        with open(os.path.join(lbl_out, fname.rsplit('.',1)[0] + '.txt'), 'w') as lf:
-            lf.write('\n'.join(annos.get(img_id, [])))
+        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, [])))
 
-    # Build example galleries
-    before_imgs, after_imgs = [], []
+    # 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)
 
-        # segmentation overlay
+        # draw polys
         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)
+            pts = np.array(anno['segmentation'][0], dtype=np.int32).reshape(-1, 2)
+            cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)
 
-        # bounding-box overlay
+        # draw boxes
         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)
+            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_imgs.append(Image.fromarray(seg_vis))
-        after_imgs.append(Image.fromarray(box_vis))
+        before.append(Image.fromarray(seg_vis))
+        after.append(Image.fromarray(box_vis))
 
-    return before_imgs, after_imgs
+    return before, after
 
 
-# Gradio UI
+# --- 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="Dataset URL (Segmentation)")
-    run_btn = gr.Button("Convert")
+    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")
@@ -159,5 +149,5 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
-    # if you need a public link change to share=True:
+    # set share=True if you need a public link
     app.launch()