Update app.py

app.py

@@ -27,21 +27,22 @@ def parse_roboflow_url(url: str):
 
 def convert_seg_to_bbox(api_key: str, dataset_url: str):
     """
-    Download a segmentation dataset from Roboflow,
-    convert masks → YOLOv8 bboxes,
-    and return (before, after) galleries + local YOLO dataset path + auto slug.
+    1) Download a segmentation dataset
+    2) Convert all masks → YOLO‐style bboxes
+    3) Write out a temp YOLO dataset and return its path
+    4) Return before/after galleries + the dataset path + an auto slug
     """
     rf = Roboflow(api_key=api_key)
-    ws_name, seg_proj_slug, ver = parse_roboflow_url(dataset_url)
-    version_obj = rf.workspace(ws_name).project(seg_proj_slug).version(ver)
+    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
 
-    # find the annotation JSON
+    # Find the annotation JSON
     ann_file = None
     for dp, _, files in os.walk(root):
         for f in files:
-            if f.lower().endswith(".json") and "train" in f.lower():
+            if 'train' in f.lower() and f.lower().endswith('.json'):
                 ann_file = os.path.join(dp, f)
                 break
         if ann_file:
@@ -49,74 +50,84 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
     if not ann_file:
         for dp, _, files in os.walk(root):
             for f in files:
-                if f.lower().endswith(".json"):
+                if f.lower().endswith('.json'):
                     ann_file = os.path.join(dp, f)
                     break
             if ann_file:
                 break
     if not ann_file:
-        raise FileNotFoundError(f"No JSON annotations found under {root}")
+        raise FileNotFoundError(f"No JSON annotations 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", []))
+    with open(ann_file, 'r') as f:
+        coco = json.load(f)
+    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 YOLOv8 folders
+    # Prepare YOLO 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)
 
-    # convert seg → bbox labels
+    # Convert seg→bbox
     annos = {}
-    for a in coco["annotations"]:
-        pid = a["image_id"]
-        poly = a["segmentation"][0]
+    for anno in coco['annotations']:
+        img_id = anno['image_id']
+        poly = anno['segmentation'][0]
         xs, ys = poly[0::2], poly[1::2]
-        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
+        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']]} "
+            f"{cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
+        )
+        annos.setdefault(img_id, []).append(line)
+
+    # Find the image folder
+    train_img_dir = None
     for dp, _, files in os.walk(root):
-        if any(f.lower().endswith((".jpg", ".png", ".jpeg")) for f in files):
-            img_dir = dp
+        if any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in files):
+            train_img_dir = dp
             break
-    if not img_dir:
-        raise FileNotFoundError(f"No images found under {root}")
+    if not train_img_dir:
+        raise FileNotFoundError(f"No images under {root}")
 
-    fname2id = {img["file_name"]: img["id"] for img in coco["images"]}
-    for fname, pid in fname2id.items():
-        src = os.path.join(img_dir, fname)
+    # Copy images + write 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:
-            lf.write("\n".join(annos.get(pid, [])))
+        with open(os.path.join(lbl_out, fname.rsplit('.', 1)[0] + ".txt"), 'w') as lf:
+            lf.write("\n".join(annos.get(img_id, [])))
 
-    # build preview galleries
+    # Build before/after sample galleries
     before, after = [], []
-    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)
+    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)
 
+        # segmentation overlay
         seg_vis = img.copy()
-        for a in coco["annotations"]:
-            if a["image_id"] != fname2id[fn]:
+        for anno in coco['annotations']:
+            if anno['image_id'] != name_to_id[fname]:
                 continue
-            pts = np.array(a["segmentation"][0], np.int32).reshape(-1, 2)
+            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(fname2id[fn], []):
+        for line in annos.get(name_to_id[fname], []):
             _, cxn, cyn, wnorm, hnorm = map(float, line.split())
-            iw, ih = images_info[fname2id[fn]]["width"], images_info[fname2id[fn]]["height"]
+            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)
@@ -125,85 +136,87 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
         before.append(Image.fromarray(seg_vis))
         after.append(Image.fromarray(box_vis))
 
-    # auto‐slug for the detection project
-    detection_slug = f"{seg_proj_slug}-detection"
-    return before, after, out_root, detection_slug
+    # auto-generated detection project slug
+    project_slug = f"{proj_name}-detection"
+    return before, after, out_root, project_slug
 
 
-def upload_and_train_detection(api_key: str, project_slug: str, dataset_path: str):
+def upload_and_train_detection(
+    api_key: str,
+    project_slug: str,
+    dataset_path: str,
+    project_license: str = "MIT",
+    project_type: str = "object-detection"
+):
     """
-    Given a YOLOv8 dataset folder, upload → version → train →
-    return inference endpoint URL. Auto‐creates the project if needed.
+    1) (re)create a Detection project
+    2) upload the YOLO dataset
+    3) generate & train a new version
+    4) return the hosted endpoint URL
     """
     rf = Roboflow(api_key=api_key)
     ws = rf.workspace()
 
-    # get or create the detection project
+    # 1) get-or-create project (need annotation arg)
     try:
         proj = ws.project(project_slug)
     except Exception:
+        # annotation must be provided as the 2nd positional arg
         proj = ws.create_project(
-            project_name=project_slug,
-            project_type="object-detection",
-            project_license="MIT"
+            project_slug,
+            annotation=project_type,
+            project_type=project_type,
+            project_license=project_license
         )
 
-    # upload the dataset
+    # 2) upload dataset
     ws.upload_dataset(
         dataset_path,
-        proj.id,
-        num_workers=10,
-        project_license="MIT",
-        project_type="object-detection",
-        batch_name=None,
-        num_retries=0
+        project_slug,
+        project_license=project_license,
+        project_type=project_type
     )
 
-    # generate a new version
-    new_v = proj.generate_version(settings={"preprocessing": {}, "augmentation": {}})
-
-    # train (fast)
-    version = proj.version(new_v)
-    version.train(speed="fast")
+    # 3) generate a new version (no args = default preprocessing/augmentation)
+    version_num = proj.generate_version()
+    # 4) train it
+    proj.version(str(version_num)).train()
 
-    # return the hosted inference URL
-    m = version.model
-    return f"{m['base_url']}{m['id']}?api_key={api_key}"
+    # 5) grab its hosted endpoint
+    m = proj.version(str(version_num)).model
+    endpoint = f"{m['base_url']}{m['id']}?api_key={api_key}"
+    return endpoint
 
 
+# --- Gradio app ---
 with gr.Blocks() as app:
-    gr.Markdown("## 🔄 Segmentation → YOLOv8 + 📡 Auto‑Deploy Detector")
-
-    # ─ Convert UI ─────────────────────────────────────────
-    api    = gr.Textbox(label="Roboflow API Key", type="password")
-    segurl = gr.Textbox(label="Segmentation Dataset URL")
-    btn_c  = gr.Button("Convert to YOLOv8 BBoxes")
-    out_b  = gr.Gallery(label="Before (masks)")
-    out_a  = gr.Gallery(label="After (bboxes)")
-    state_path = gr.State()
-    state_slug = gr.State()
-
-    btn_c.click(
-        fn=convert_seg_to_bbox,
-        inputs=[api, segurl],
-        outputs=[out_b, out_a, state_path, state_slug]
-    )
+    gr.Markdown("## 🔄 Segmentation → YOLOv8 Converter + Auto‐Upload")
 
-    gr.Markdown("---")
+    # single API key input
+    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(label="Before (Segmentation)", columns=5, height="auto")
+    after_g  = gr.Gallery(label="After (BBoxes)", columns=5, height="auto")
 
-    # ─ Train UI ───────────────────────────────────────────
-    btn_t = gr.Button("Upload & Train Detection Model")
-    endpoint = gr.Textbox(label="Hosted Detection Endpoint URL")
+    # hidden states for the YOLO dataset path & auto‐slug
+    dataset_path_state  = gr.Textbox(visible=False)
+    project_slug_state  = gr.Textbox(visible=False)
 
-    btn_t.click(
-        fn=upload_and_train_detection,
-        inputs=[api, state_slug, state_path],
-        outputs=[endpoint]
+    run_btn.click(
+        fn=convert_seg_to_bbox,
+        inputs=[api_input, url_input],
+        outputs=[before_g, after_g, dataset_path_state, project_slug_state]
     )
 
-    gr.Markdown(
-        "> 1) Paste your segmentation URL and Convert.  \n"
-        "> 2) Then Upload & Train to instantly get your detector’s endpoint."
+    gr.Markdown("## 🚀 Upload & Train Detection Model")
+    train_btn  = gr.Button("Upload & Train Detection Model")
+    url_output = gr.Textbox(label="Model Endpoint URL")
+
+    train_btn.click(
+        fn=upload_and_train_detection,
+        inputs=[api_input, project_slug_state, dataset_path_state],
+        outputs=[url_output]
     )
 
 if __name__ == "__main__":