Update app.py

app.py

@@ -9,28 +9,31 @@ import cv2
 import numpy as np
 from PIL import Image
 import gradio as gr
-from roboflow import Roboflow
+from roboflow import Roboflow, RoboflowError
 
 
 def parse_roboflow_url(url: str):
+    """
+    Extract (workspace, project slug, version) from any Roboflow URL.
+    """
     parsed = urlparse(url)
     parts = parsed.path.strip('/').split('/')
-    workspace = parts[0]
-    project   = parts[1]
+    ws       = parts[0]
+    proj     = parts[1]
     try:
-        version = int(parts[-1])
+        ver = int(parts[-1])
     except ValueError:
-        version = int(parts[-2])
-    return workspace, project, version
+        ver = int(parts[-2])
+    return ws, proj, ver
 
 
 def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1, 0.1)):
     # --- download segmentation export
-    rf = Roboflow(api_key=api_key)
-    ws, proj_name, ver = parse_roboflow_url(dataset_url)
-    version_obj = rf.workspace(ws).project(proj_name).version(ver)
-    dataset     = version_obj.download("coco-segmentation")
-    root        = dataset.location
+    rf           = Roboflow(api_key=api_key)
+    workspace, proj_name, ver = parse_roboflow_url(dataset_url)
+    version_obj  = rf.workspace(workspace).project(proj_name).version(ver)
+    dataset      = version_obj.download("coco-segmentation")
+    root         = dataset.location
 
     # --- find the COCO JSON
     ann_file = None
@@ -42,39 +45,38 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         if ann_file:
             break
     if not ann_file:
-        raise FileNotFoundError(f"No JSON annotations under {root}")
+        raise FileNotFoundError(f"No JSON annotations found under {root}")
 
     coco        = json.load(open(ann_file, 'r'))
     images_info = {img['id']: img for img in coco['images']}
     cat_ids     = sorted(c['id'] for c in coco.get('categories', []))
     id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}
 
-    # --- prepare flat_images + flat_labels
+    # --- flatten + convert to YOLO bboxes
     out_root = tempfile.mkdtemp(prefix="yolov8_")
     flat_img = os.path.join(out_root, "flat_images")
     flat_lbl = os.path.join(out_root, "flat_labels")
     os.makedirs(flat_img, exist_ok=True)
     os.makedirs(flat_lbl, exist_ok=True)
 
-    # --- convert each segmentation → 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
+        xmin, xmax = min(xs), max(xs)
+        ymin, ymax = min(ys), max(ys)
+        w, h     = xmax - xmin, ymax - ymin
+        cx, cy   = xmin + w/2, ymin + h/2
 
-        iw, ih = images_info[img_id]['width'], images_info[img_id]['height']
-        line = (
+        iw, ih   = images_info[img_id]['width'], images_info[img_id]['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)
 
-    # --- map each file_name to its actual path on disk
+    # --- map filenames to their disk paths
     name_to_id = {img['file_name']: img['id'] for img in coco['images']}
     file_paths = {}
     for dp, _, files in os.walk(root):
@@ -82,26 +84,22 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
             if f in name_to_id:
                 file_paths[f] = os.path.join(dp, f)
 
-    # --- copy images + write flat_labels
+    # --- copy images and write YOLO .txt labels
     for fname, img_id in name_to_id.items():
-        src_path = file_paths.get(fname)
-        if not src_path:
-            # skip if we couldn't find this image under root
+        src = file_paths.get(fname)
+        if not src:
             continue
-        shutil.copy(src_path, os.path.join(flat_img, fname))
+        shutil.copy(src, os.path.join(flat_img, fname))
         with open(os.path.join(flat_lbl, fname.rsplit('.',1)[0] + ".txt"), 'w') as lf:
             lf.write("\n".join(annos.get(img_id, [])))
 
-    # --- split into train/valid/test
-    all_files = sorted(
-        f for f in os.listdir(flat_img)
-        if f.lower().endswith(('.jpg','.png','.jpeg'))
-    )
+    # --- split into train/val/test
+    all_files = sorted(f for f in os.listdir(flat_img)
+                       if f.lower().endswith(('.jpg','.png','.jpeg')))
     random.shuffle(all_files)
     n = len(all_files)
     n_train = max(1, int(n * split_ratios[0]))
     n_valid = max(1, int(n * split_ratios[1]))
-    # ensure at least 1 left for test
     n_valid = min(n_valid, n - n_train - 1)
 
     splits = {
@@ -110,28 +108,23 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         "test":  all_files[n_train+n_valid:]
     }
 
-    # --- build Roboflow‐friendly folder structure
+    # --- arrange into Roboflow‑friendly folder tree
     for split, files in splits.items():
-        out_img_dir = os.path.join(out_root, "images", split)
-        out_lbl_dir = os.path.join(out_root, "labels", split)
-        os.makedirs(out_img_dir, exist_ok=True)
-        os.makedirs(out_lbl_dir, exist_ok=True)
+        idir = os.path.join(out_root, "images", split)
+        ldir = os.path.join(out_root, "labels", split)
+        os.makedirs(idir, exist_ok=True)
+        os.makedirs(ldir, exist_ok=True)
         for fn in files:
-            shutil.move(
-                os.path.join(flat_img, fn),
-                os.path.join(out_img_dir, fn)
-            )
-            lbl_fn = fn.rsplit('.',1)[0] + ".txt"
-            shutil.move(
-                os.path.join(flat_lbl, lbl_fn),
-                os.path.join(out_lbl_dir, lbl_fn)
-            )
+            shutil.move(os.path.join(flat_img, fn),
+                        os.path.join(idir, fn))
+            lbl = fn.rsplit('.',1)[0] + ".txt"
+            shutil.move(os.path.join(flat_lbl, lbl),
+                        os.path.join(ldir, lbl))
 
-    # --- clean up the flat dirs
     shutil.rmtree(flat_img)
     shutil.rmtree(flat_lbl)
 
-    # --- prepare a few before/after visuals
+    # --- make a few before/after visual samples
     before, after = [], []
     sample = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))
     for fname in sample:
@@ -159,32 +152,37 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         before.append(Image.fromarray(seg_vis))
         after.append(Image.fromarray(box_vis))
 
-    project_slug = f"{proj_name}-detection"
-    return before, after, out_root, project_slug
+    # return samples + local folder + the two slugs we need downstream
+    return before, after, out_root, proj_name + "-detection", workspace
 
 
 def upload_and_train_detection(
     api_key: str,
+    workspace: str,
     project_slug: str,
     dataset_path: str,
     project_license: str = "MIT",
     project_type: str = "object-detection"
 ):
     rf = Roboflow(api_key=api_key)
-    ws = rf.workspace()
+    ws = rf.workspace(workspace)
 
-    # get‐or‐create your detection project
+    # --- get‑or‑create project
     try:
         proj = ws.project(project_slug)
-    except:
-        proj = ws.create_project(
-            project_slug,
-            annotation=project_type,
-            project_type=project_type,
-            project_license=project_license
-        )
+    except RoboflowError as e:
+        # only create if truly “not found”
+        if "does not exist" in str(e):
+            proj = ws.create_project(
+                project_slug,
+                annotation=project_type,
+                project_type=project_type,
+                project_license=project_license
+            )
+        else:
+            raise
 
-    # upload the folder with proper train/valid/test
+    # --- upload the new train/val/test
     ws.upload_dataset(
         dataset_path,
         project_slug,
@@ -192,34 +190,36 @@ def upload_and_train_detection(
         project_type=project_type
     )
 
-    # create a new version & queue training
+    # --- spin up a new version and start training
     version_num = proj.generate_version(settings={
         "augmentation": {},
         "preprocessing": {},
     })
     proj.version(str(version_num)).train()
 
-    # return the hosted endpoint URL
     m = proj.version(str(version_num)).model
     return f"{m['base_url']}{m['id']}?api_key={api_key}"
 
 
 # --- Gradio UI ---
 with gr.Blocks() as app:
-    gr.Markdown("## 🔄 Seg→BBox + Auto‐Upload/Train")
+    gr.Markdown("## 🔄 Seg→BBox + Auto‑Upload/Train")
 
     api_input  = gr.Textbox(label="Roboflow API Key", type="password")
     url_input  = gr.Textbox(label="Segmentation Dataset URL")
     run_btn    = gr.Button("Convert to BBoxes")
     before_g   = gr.Gallery(columns=5, label="Before")
     after_g    = gr.Gallery(columns=5, label="After")
-    ds_state   = gr.Textbox(visible=False)
-    slug_state = gr.Textbox(visible=False)
+
+    # hidden states
+    ds_state   = gr.Textbox(visible=False)  # local dataset folder
+    slug_state = gr.Textbox(visible=False)  # project‑slug e.g. "myproj-detection"
+    ws_state   = gr.Textbox(visible=False)  # workspace name
 
     run_btn.click(
         convert_seg_to_bbox,
         inputs=[api_input, url_input],
-        outputs=[before_g, after_g, ds_state, slug_state]
+        outputs=[before_g, after_g, ds_state, slug_state, ws_state]
     )
 
     gr.Markdown("## 🚀 Upload & Train Detection Model")
@@ -228,7 +228,7 @@ with gr.Blocks() as app:
 
     train_btn.click(
         upload_and_train_detection,
-        inputs=[api_input, slug_state, ds_state],
+        inputs=[api_input, ws_state, slug_state, ds_state],
         outputs=[url_out]
     )