Update app.py

app.py

@@ -26,14 +26,18 @@ def parse_roboflow_url(url: str):
 
 
 def convert_seg_to_bbox(api_key: str, dataset_url: str):
-    """Download a segmentation dataset from Roboflow, convert to YOLO bboxes, and return before/after galleries."""
+    """
+    1) Download segmentation dataset from Roboflow
+    2) Convert masks → YOLOv8 bboxes
+    Returns before_gallery, after_gallery, local_dataset_path, project_slug
+    """
     rf = Roboflow(api_key=api_key)
     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
 
-    # 1) Find annotation JSON
+    # find annotation JSON
     ann_file = None
     for dp, _, files in os.walk(root):
         for f in files:
@@ -51,146 +55,136 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
             if ann_file:
                 break
     if not ann_file:
-        raise FileNotFoundError("No JSON annotations found under %s" % 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)}
 
-    # 2) Prepare YOLO folders
+    # prepare YOLOv8 folders
     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)
 
-    # 3) Convert seg→bbox
+    # convert seg → bbox
     annos = {}
-    for anno in coco['annotations']:
-        img_id = anno['image_id']
-        poly = anno['segmentation'][0]
+    for a in coco['annotations']:
+        pid = a['image_id']
+        poly = a['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
-
-        iw, ih = images_info[img_id]['width'], images_info[img_id]['height']
-        line = f"{id_to_index[anno['category_id']]} {cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
-        annos.setdefault(img_id, []).append(line)
-
-    # 4) Find images and write labels
-    train_img_dir = None
+        xmin, xmax, ymin, ymax = min(xs), max(xs), min(ys), max(ys)
+        w, h = xmax - xmin, ymax - ymin
+        cx, cy = xmin + w/2, ymin + h/2
+        iw, ih = images_info[pid]['width'], images_info[pid]['height']
+        line = f"{id_to_index[a['category_id']]} {cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
+        annos.setdefault(pid, []).append(line)
+
+    # locate images and write labels
+    img_dir = None
     for dp, _, files in os.walk(root):
-        if any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in files):
-            train_img_dir = dp
+        if any(f.lower().endswith(('.jpg','.png','jpeg')) for f in files):
+            img_dir = dp
             break
-    if not train_img_dir:
-        raise FileNotFoundError("No image files found under %s" % root)
+    if not img_dir:
+        raise FileNotFoundError(f"No image files found under {root}")
 
-    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)
+    fname2id = {img['file_name']: img['id'] for img in coco['images']}
+    for fname, pid in fname2id.items():
+        src = os.path.join(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, [])))
+        with open(os.path.join(lbl_out, fname.rsplit('.',1)[0] + ".txt"), 'w') as lf:
+            lf.write("\n".join(annos.get(pid, [])))
 
-    # 5) Build galleries
+    # build preview 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)
+    sample = random.sample(list(fname2id.keys()), min(5, len(fname2id)))
+    for fn in sample:
+        img = cv2.cvtColor(cv2.imread(os.path.join(img_dir, fn)), cv2.COLOR_BGR2RGB)
 
-        # draw seg polygons
         seg_vis = img.copy()
-        for anno in coco['annotations']:
-            if anno['image_id'] != name_to_id[fname]:
+        for a in coco['annotations']:
+            if a['image_id'] != fname2id[fn]:
                 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(a['segmentation'][0], np.int32).reshape(-1,2)
+            cv2.polylines(seg_vis, [pts], True, (255,0,0), 2)
 
-        # draw boxes
         box_vis = img.copy()
-        for line in annos.get(name_to_id[fname], []):
+        for line in annos.get(fname2id[fn], []):
             _, 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)
+            iw, ih = images_info[fname2id[fn]]['width'], images_info[fname2id[fn]]['height']
+            w0, h0 = int(wnorm*iw), int(hnorm*ih)
+            x0, y0 = int(cxn*iw - w0/2), 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))
 
-    return before, after
+    return before, after, out_root, proj  # proj is our slug
 
 
-def upload_and_train_detection(
-    api_key: str,
-    project_id: str,
-    dataset_path: str,
-    project_license: str = "MIT",
-    project_type: str = "object-detection",
-    preprocessing: dict = None,
-    augmentation: dict = None,
-    speed: str = "fast"
-):
+def upload_and_train_detection(api_key: str, project_slug: str, dataset_path: str):
     """
-    Upload a local detection dataset to Roboflow, generate+train a new version,
-    and return the hosted inference endpoint URL.
+    1) Upload local YOLOv8 dataset to Roboflow
+    2) Generate & train a new detection version
+    Returns the hosted inference endpoint URL.
     """
     rf = Roboflow(api_key=api_key)
     ws = rf.workspace()
 
-    # 1) upload
+    # upload dataset
     ws.upload_dataset(
         dataset_path,
-        project_id,
-        project_license=project_license,
-        project_type=project_type
+        project_slug,
+        project_license="MIT",
+        project_type="object-detection"
     )
 
-    # 2) generate version
-    proj = ws.project(project_id)
-    version_number = proj.generate_version(
-        preprocessing=preprocessing or {},
-        augmentation=augmentation or {}
-    )
+    # generate a new version
+    proj = ws.project(project_slug)
+    version_number = proj.generate_version(preprocessing={}, augmentation={})
 
-    # 3) train
-    proj.version(version_number).train(speed=speed)
+    # train model (fast)
+    proj.version(version_number).train(speed="fast")
 
-    # 4) fetch model endpoint info
+    # fetch hosted endpoint
     m = proj.version(str(version_number)).model
-    endpoint = f"{m['base_url']}{m['id']}?api_key={api_key}"
-    return endpoint
+    return f"{m['base_url']}{m['id']}?api_key={api_key}"
 
 
-# --- Gradio app ---
 with gr.Blocks() as app:
-    gr.Markdown("## 🔄 Segmentation → YOLOv8 Converter")
-    api_input1 = 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(label="Before (Segmentation)", columns=5)
-    after_g    = gr.Gallery(label="After (BBoxes)", columns=5)
-    run_btn.click(fn=convert_seg_to_bbox, inputs=[api_input1, url_input], outputs=[before_g, after_g])
-
-    gr.Markdown("## 🚀 Upload & Train Detection Model")
-    api_input2    = gr.Textbox(label="Roboflow API Key", type="password")
-    project_input = gr.Textbox(label="Project ID (slug)")
-    path_input    = gr.Textbox(label="Local Dataset Path")
-    train_btn     = gr.Button("Upload & Train")
-    url_output    = gr.Textbox(label="Hosted Model Endpoint URL")
+    gr.Markdown("## 🔄 Segmentation → YOLOv8 Converter + Auto Trainer")
+
+    # Converter UI
+    api_input    = gr.Textbox(label="Roboflow API Key", type="password")
+    url_input    = gr.Textbox(label="Segmentation Dataset URL")
+    convert_btn  = gr.Button("Convert to BBoxes")
+    before_gal   = gr.Gallery(label="Before (Segmentation)", columns=5)
+    after_gal    = gr.Gallery(label="After (BBoxes)", columns=5)
+    state_path   = gr.State()
+    state_slug   = gr.State()
+
+    convert_btn.click(
+        fn=convert_seg_to_bbox,
+        inputs=[api_input, url_input],
+        outputs=[before_gal, after_gal, state_path, state_slug]
+    )
+
+    # Train UI
+    train_btn     = gr.Button("Upload & Train Detection Model")
+    endpoint_text = gr.Textbox(label="Hosted Detection Endpoint URL")
+
     train_btn.click(
         fn=upload_and_train_detection,
-        inputs=[api_input2, project_input, path_input],
-        outputs=[url_output],
+        inputs=[api_input, state_slug, state_path],
+        outputs=[endpoint_text]
     )
 
+    gr.Markdown("> First convert your seg data, then click **Upload & Train** to deploy your detection model.")
+
 if __name__ == "__main__":
     app.launch()