Update app.py

app.py

@@ -3,7 +3,7 @@ import json
 import random
 import shutil
 import tempfile
-import time              # ← added for polling
+import time
 from urllib.parse import urlparse
 
 import cv2
@@ -32,11 +32,11 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
     3) Split into train/valid/test
     4) Return before/after visuals plus (dataset_path, detection_slug)
     """
-    rf            = Roboflow(api_key=api_key)
+    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
+    version_obj = rf.workspace(ws).project(proj_name).version(ver)
+    dataset = version_obj.download("coco-segmentation")
+    root = dataset.location
 
     # find the COCO JSON
     ann_file = None
@@ -50,9 +50,9 @@ 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 found 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)}
 
     # build YOLO bboxes
@@ -71,13 +71,8 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         )
         annos.setdefault(img_id, []).append(line)
 
-    # copy and write out flat images + labels
+    # prepare temporary split folder
     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)
-
     name_to_id = {img['file_name']: img['id'] for img in coco['images']}
     file_paths = {
         f: os.path.join(dp, f)
@@ -86,17 +81,8 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         if f in name_to_id
     }
 
-    for fname, img_id in name_to_id.items():
-        src = file_paths.get(fname)
-        if not src:
-            continue
-        shutil.copy(src, os.path.join(flat_img, fname))
-        lbl_path = os.path.join(flat_lbl, fname.rsplit('.',1)[0] + ".txt")
-        with open(lbl_path, 'w') as lf:
-            lf.write("\n".join(annos.get(img_id, [])))
-
-    # split filenames
-    all_files = [f for f in os.listdir(flat_img) if f.lower().endswith(('.jpg','.png','.jpeg'))]
+    # determine splits
+    all_files = list(name_to_id.keys())
     random.shuffle(all_files)
     n = len(all_files)
     n_train = max(1, int(n * split_ratios[0]))
@@ -108,27 +94,26 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
         "test":  all_files[n_train+n_valid:]
     }
 
-    # move into final folder structure
+    # create Roboflow‐style folders
     for split, files in splits.items():
-        img_dir = os.path.join(out_root, "images", split)
-        lbl_dir = os.path.join(out_root, "labels", split)
+        img_dir = os.path.join(out_root, split, "images")
+        lbl_dir = os.path.join(out_root, split, "labels")
         os.makedirs(img_dir, exist_ok=True)
         os.makedirs(lbl_dir, exist_ok=True)
-        for fn in files:
-            shutil.move(os.path.join(flat_img, fn), os.path.join(img_dir, fn))
-            lbl = fn.rsplit('.',1)[0] + ".txt"
-            shutil.move(os.path.join(flat_lbl, lbl), os.path.join(lbl_dir, lbl))
-
-    shutil.rmtree(flat_img)
-    shutil.rmtree(flat_lbl)
-
-    # prepare a few before/after images for display
+        for fname in files:
+            # copy image
+            shutil.copy(file_paths[fname], os.path.join(img_dir, fname))
+            # write label
+            txt = "\n".join(annos.get(name_to_id[fname], []))
+            with open(os.path.join(lbl_dir, fname.rsplit('.',1)[0] + ".txt"), 'w') as f:
+                f.write(txt)
+
+    # prepare display examples
     before, after = [], []
-    sample = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))
+    sample = random.sample(all_files, min(5, len(all_files)))
     for fname in sample:
         img = cv2.cvtColor(cv2.imread(file_paths[fname]), cv2.COLOR_BGR2RGB)
 
-        # original segmentation overlay
         seg_vis = img.copy()
         for anno in coco['annotations']:
             if anno['image_id'] != name_to_id[fname]:
@@ -136,7 +121,6 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1,
             pts = np.array(anno['segmentation'][0], np.int32).reshape(-1,2)
             cv2.polylines(seg_vis, [pts], True, (255,0,0), 2)
 
-        # bbox overlay
         box_vis = img.copy()
         for line in annos.get(name_to_id[fname], []):
             _, cxn, cyn, wnorm, hnorm = map(float, line.split())
@@ -168,7 +152,7 @@ def upload_and_train_detection(
     rf = Roboflow(api_key=api_key)
     ws = rf.workspace()
 
-    # 1) get-or-create project
+    # get-or-create project
     try:
         proj = ws.project(detection_slug)
     except Exception as e:
@@ -182,7 +166,7 @@ def upload_and_train_detection(
         else:
             raise
 
-    # 2) extract real slug from proj.id and upload
+    # upload entire split folder
     _, real_slug = proj.id.rsplit("/", 1)
     ws.upload_dataset(
         dataset_path,
@@ -191,16 +175,12 @@ def upload_and_train_detection(
         project_type=project_type
     )
 
-    # 3) generate a new version
+    # generate new version (with fallback)
     try:
-        version_num = proj.generate_version(settings={
-            "augmentation": {},
-            "preprocessing": {}
-        })
+        version_num = proj.generate_version(settings={"augmentation": {}, "preprocessing": {}})
     except RuntimeError as e:
         msg = str(e).lower()
         if "unsupported request" in msg or "does not exist" in msg:
-            # bump slug and retry
             new_slug = real_slug + "-v2"
             proj = ws.create_project(
                 new_slug,
@@ -208,20 +188,14 @@ def upload_and_train_detection(
                 project_type=project_type,
                 project_license=project_license
             )
-            ws.upload_dataset(
-                dataset_path,
-                new_slug,
-                project_license=project_license,
-                project_type=project_type
-            )
-            version_num = proj.generate_version(settings={
-                "augmentation": {},
-                "preprocessing": {}
-            })
+            ws.upload_dataset(dataset_path, new_slug,
+                              project_license=project_license,
+                              project_type=project_type)
+            version_num = proj.generate_version(settings={"augmentation": {}, "preprocessing": {}})
         else:
             raise
 
-    # 4) wait until the dataset version is done generating, then train
+    # wait until ready, then train
     for _ in range(20):
         try:
             model = proj.version(str(version_num)).train()
@@ -241,12 +215,12 @@ 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, label="Converted Dataset Path")
+    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, label="Converted Dataset Path")
     slug_state = gr.Textbox(visible=False, label="Detection Project Slug")
 
     run_btn.click(
@@ -257,7 +231,7 @@ with gr.Blocks() as app:
 
     gr.Markdown("## 🚀 Upload & Train Detection Model")
     train_btn = gr.Button("Upload & Train")
-    url_out   = gr.Textbox(label="Hosted Model Endpoint URL")
+    url_out = gr.Textbox(label="Hosted Model Endpoint URL")
 
     train_btn.click(
         upload_and_train_detection,