Update app.py

app.py

@@ -16,7 +16,7 @@ def parse_roboflow_url(url: str):
     parsed = urlparse(url)
     parts = parsed.path.strip('/').split('/')
     workspace = parts[0]
-    project = parts[1]
+    project   = parts[1]
     try:
         version = int(parts[-1])
     except ValueError:
@@ -29,8 +29,8 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
     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
+    dataset     = version_obj.download("coco-segmentation")
+    root        = dataset.location
 
     # --- find the COCO JSON
     ann_file = None
@@ -44,23 +44,23 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
     if not ann_file:
         raise FileNotFoundError(f"No JSON annotations under {root}")
 
-    coco = json.load(open(ann_file, 'r'))
+    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', []))
+    cat_ids     = sorted(c['id'] for c in coco.get('categories', []))
     id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}
 
-    # --- make a flat YOLO folder
+    # --- prepare flat_images + flat_labels
     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 to a YOLO bbox line
+    # --- convert each segmentation → YOLO bbox line
     annos = {}
     for anno in coco['annotations']:
         img_id = anno['image_id']
-        poly = anno['segmentation'][0]
+        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)
@@ -74,16 +74,16 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         )
         annos.setdefault(img_id, []).append(line)
 
-    # --- locate the single images folder
+    # --- locate the single folder of raw images
     img_src = None
     for dp, _, files in os.walk(root):
         if any(f.lower().endswith(('.jpg','.png','.jpeg')) for f in files):
             img_src = dp
             break
     if not img_src:
-        raise FileNotFoundError(f"No images folder in {root}")
+        raise FileNotFoundError(f"No images under {root}")
 
-    # --- copy images + write flat labels
+    # --- copy images + write flat label files
     name_to_id = {img['file_name']: img['id'] for img in coco['images']}
     for fname, img_id in name_to_id.items():
         src_path = os.path.join(img_src, fname)
@@ -93,38 +93,48 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         with open(os.path.join(flat_lbl, fname.rsplit('.',1)[0] + ".txt"), 'w') as lf:
             lf.write("\n".join(annos.get(img_id, [])))
 
-    # --- split filenames into train/valid/test lists
-    all_files = sorted([f for f in os.listdir(flat_img) if f.lower().endswith(('.jpg','.png','.jpeg'))])
+    # --- split into train/valid/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 we don’t overshoot
+    # ensure at least 1 for each split
     n_valid = min(n_valid, n - n_train - 1)
+
     splits = {
         "train": all_files[:n_train],
         "valid": all_files[n_train:n_train+n_valid],
         "test":  all_files[n_train+n_valid:]
     }
 
-    # --- create Roboflow‑friendly structure:
-    # out_root/images/{train,valid,test}
-    # out_root/labels/{train,valid,test}
+    # --- build Roboflow‐friendly folder structure
     for split, files in splits.items():
-        img_dir = os.path.join(out_root, "images", split)
-        lbl_dir = os.path.join(out_root, "labels", split)
-        os.makedirs(img_dir, exist_ok=True)
-        os.makedirs(lbl_dir, exist_ok=True)
+        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)
         for fn in files:
-            shutil.move(os.path.join(flat_img, fn),  os.path.join(img_dir, fn))
-            shutil.move(os.path.join(flat_lbl, fn.rsplit('.',1)[0] + ".txt"),
-                        os.path.join(lbl_dir, fn.rsplit('.',1)[0] + ".txt"))
-
-    # --- clean up flats
+            # move image
+            shutil.move(
+                os.path.join(flat_img, fn),
+                os.path.join(out_img_dir, fn)
+            )
+            # move corresponding .txt label
+            lbl_fn = fn.rsplit('.',1)[0] + ".txt"
+            shutil.move(
+                os.path.join(flat_lbl, lbl_fn),
+                os.path.join(out_lbl_dir, lbl_fn)
+            )
+
+    # --- clean up the flat dirs
     shutil.rmtree(flat_img)
     shutil.rmtree(flat_lbl)
 
-    # --- build a few before/after previews
+    # --- prepare a few before/after visuals
     before, after = [], []
     sample = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))
     for fname in sample:
@@ -135,17 +145,17 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         for anno in coco['annotations']:
             if anno['image_id'] != name_to_id[fname]:
                 continue
-            pts = np.array(anno['segmentation'][0], np.int32).reshape(-1, 2)
-            cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)
+            pts = np.array(anno['segmentation'][0], np.int32).reshape(-1,2)
+            cv2.polylines(seg_vis, [pts], True, (255,0,0), 2)
 
         box_vis = img.copy()
         for line in annos.get(name_to_id[fname], []):
             _, cxn, cyn, wnorm, hnorm = map(float, line.split())
             iw, ih = images_info[name_to_id[fname]]['width'], images_info[name_to_id[fname]]['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.append(Image.fromarray(seg_vis))
         after.append(Image.fromarray(box_vis))
@@ -164,10 +174,10 @@ def upload_and_train_detection(
     rf = Roboflow(api_key=api_key)
     ws = rf.workspace()
 
-    # get-or-create your detection project
+    # get-or-create detection project
     try:
         proj = ws.project(project_slug)
-    except Exception:
+    except:
         proj = ws.create_project(
             project_slug,
             annotation=project_type,
@@ -175,7 +185,7 @@ def upload_and_train_detection(
             project_license=project_license
         )
 
-    # upload the properly‑split folder
+    # upload the folder that now has train/valid/test
     ws.upload_dataset(
         dataset_path,
         project_slug,
@@ -183,13 +193,11 @@ def upload_and_train_detection(
         project_type=project_type
     )
 
-    # create a new version
+    # create a new version & queue training
     version_num = proj.generate_version(settings={
         "augmentation": {},
         "preprocessing": {},
     })
-
-    # enqueue training (now finds train/valid/test)
     proj.version(str(version_num)).train()
 
     # return the hosted endpoint URL
@@ -201,13 +209,13 @@ def upload_and_train_detection(
 with gr.Blocks() as app:
     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)
+    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)
 
     run_btn.click(
         convert_seg_to_bbox,