Update app.py

app.py

@@ -24,7 +24,7 @@ def parse_roboflow_url(url: str):
     return workspace, project, version
 
 
-def convert_seg_to_bbox(api_key: str, dataset_url: str):
+def convert_seg_to_bbox(api_key: str, dataset_url: str, split_ratios=(0.8, 0.1, 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)
@@ -56,14 +56,14 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
     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 dirs
+    # prepare flat YOLO dirs
     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)
 
-    # convert seg→bbox
+    # convert segmentation → bounding‐box labels
     annos = {}
     for anno in coco['annotations']:
         img_id = anno['image_id']
@@ -72,7 +72,7 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
         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
+        cx, cy = x_min + w / 2, y_min + h / 2
 
         iw, ih = images_info[img_id]['width'], images_info[img_id]['height']
         line = (
@@ -81,26 +81,52 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
         )
         annos.setdefault(img_id, []).append(line)
 
-    # locate images
+    # locate raw images folder
     train_img_dir = None
     for dp, _, files in os.walk(root):
-        if any(f.lower().endswith(('.jpg','.png','.jpeg')) for f in files):
+        if any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in files):
             train_img_dir = dp
             break
     if not train_img_dir:
         raise FileNotFoundError(f"No images under {root}")
 
-    # copy images + write labels
+    # copy images + write flat 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:
+        with open(os.path.join(lbl_out, fname.rsplit('.', 1)[0] + ".txt"), 'w') as lf:
             lf.write("\n".join(annos.get(img_id, [])))
 
-    # build before/after samples
+    # split into train/valid/test
+    all_images = sorted([f for f in os.listdir(img_out) if f.lower().endswith(('.jpg', '.png', '.jpeg'))])
+    random.shuffle(all_images)
+    n = len(all_images)
+    n_train = int(n * split_ratios[0])
+    n_valid = int(n * split_ratios[1])
+    splits = {
+        "train": all_images[:n_train],
+        "valid": all_images[n_train:n_train + n_valid],
+        "test":  all_images[n_train + n_valid:]
+    }
+
+    for split_name, files in splits.items():
+        img_dir = os.path.join(out_root, split_name, "images")
+        lbl_dir = os.path.join(out_root, split_name, "labels")
+        os.makedirs(img_dir, exist_ok=True)
+        os.makedirs(lbl_dir, exist_ok=True)
+        for fname in files:
+            shutil.move(os.path.join(img_out, fname), os.path.join(img_dir, fname))
+            lbl_fname = fname.rsplit(".", 1)[0] + ".txt"
+            shutil.move(os.path.join(lbl_out, lbl_fname), os.path.join(lbl_dir, lbl_fname))
+
+    # remove flat dirs
+    shutil.rmtree(img_out)
+    shutil.rmtree(lbl_out)
+
+    # build before/after galleries for a few samples
     before, after = [], []
     sample = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))
     for fname in sample:
@@ -111,17 +137,17 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
         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))
@@ -140,7 +166,7 @@ def upload_and_train_detection(
     rf = Roboflow(api_key=api_key)
     ws = rf.workspace()
 
-    # get-or-create detection project
+    # get or create project
     try:
         proj = ws.project(project_slug)
     except Exception:
@@ -151,7 +177,7 @@ def upload_and_train_detection(
             project_license=project_license
         )
 
-    # upload the YOLO dataset
+    # upload folder with train/valid/test
     ws.upload_dataset(
         dataset_path,
         project_slug,
@@ -159,16 +185,16 @@ def upload_and_train_detection(
         project_type=project_type
     )
 
-    # generate a new version (must pass both keys—even if empty)
+    # create new version
     version_num = proj.generate_version(settings={
         "augmentation": {},
         "preprocessing": {},
     })
 
-    # train it
+    # enqueue training (now sees splits)
     proj.version(str(version_num)).train()
 
-    # return hosted endpoint
+    # return endpoint
     m = proj.version(str(version_num)).model
     return f"{m['base_url']}{m['id']}?api_key={api_key}"
 
@@ -177,13 +203,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,