Update app.py

app.py

@@ -13,7 +13,6 @@ from roboflow import Roboflow
 
 
 def parse_roboflow_url(url: str):
-    """Extract workspace, project name, and version from a Roboflow URL."""
     parsed = urlparse(url)
     parts = parsed.path.strip('/').split('/')
     workspace = parts[0]
@@ -26,19 +25,13 @@ def parse_roboflow_url(url: str):
 
 
 def convert_seg_to_bbox(api_key: str, dataset_url: str):
-    """
-    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, 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 annotation JSON
     ann_file = None
     for dp, _, files in os.walk(root):
         for f in files:
@@ -58,20 +51,19 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
     if not ann_file:
         raise FileNotFoundError(f"No JSON annotations under {root}")
 
-    with open(ann_file, 'r') as f:
-        coco = json.load(f)
+    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)}
 
-    # Prepare YOLO folders
+    # prepare 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 seg→bbox
     annos = {}
     for anno in coco['annotations']:
         img_id = anno['image_id']
@@ -89,54 +81,51 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
         )
         annos.setdefault(img_id, []).append(line)
 
-    # Find the image folder
+    # locate images
     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 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 sample galleries
+    # build before/after samples
     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)
 
-        # segmentation overlay
         seg_vis = img.copy()
         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)
 
-        # bbox overlay
         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))
 
-    # auto-generated detection project slug
     project_slug = f"{proj_name}-detection"
     return before, after, out_root, project_slug
 
@@ -148,20 +137,13 @@ def upload_and_train_detection(
     project_license: str = "MIT",
     project_type: str = "object-detection"
 ):
-    """
-    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()
 
-    # 1) get-or-create project (need annotation arg)
+    # get-or-create detection project
     try:
         proj = ws.project(project_slug)
     except Exception:
-        # annotation must be provided as the 2nd positional arg
         proj = ws.create_project(
             project_slug,
             annotation=project_type,
@@ -169,54 +151,48 @@ def upload_and_train_detection(
             project_license=project_license
         )
 
-    # 2) upload dataset
-    ws.upload_dataset(
-        dataset_path,
-        project_slug,
-        project_license=project_license,
-        project_type=project_type
-    )
+    # upload the YOLO dataset
+    ws.upload_dataset(dataset_path, project_slug,
+                      project_license=project_license,
+                      project_type=project_type)
+
+    # generate a new version (must pass settings arg—even if empty)
+    version_num = proj.generate_version(settings={})
 
-    # 3) generate a new version (no args = default preprocessing/augmentation)
-    version_num = proj.generate_version()
-    # 4) train it
+    # train it
     proj.version(str(version_num)).train()
 
-    # 5) grab its hosted endpoint
+    # return hosted endpoint
     m = proj.version(str(version_num)).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 ---
+# --- Gradio UI ---
 with gr.Blocks() as app:
-    gr.Markdown("## 🔄 Segmentation → YOLOv8 Converter + Auto‐Upload")
+    gr.Markdown("## 🔄 Seg→BBox + Auto‐Upload/Train")
 
-    # 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")
-
-    # hidden states for the YOLO dataset path & auto‐slug
-    dataset_path_state  = gr.Textbox(visible=False)
-    project_slug_state  = gr.Textbox(visible=False)
+    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(
-        fn=convert_seg_to_bbox,
+        convert_seg_to_bbox,
         inputs=[api_input, url_input],
-        outputs=[before_g, after_g, dataset_path_state, project_slug_state]
+        outputs=[before_g, after_g, ds_state, slug_state]
     )
 
     gr.Markdown("## 🚀 Upload & Train Detection Model")
-    train_btn  = gr.Button("Upload & Train Detection Model")
-    url_output = gr.Textbox(label="Model Endpoint URL")
+    train_btn = gr.Button("Upload & Train")
+    url_out   = gr.Textbox(label="Hosted Model Endpoint URL")
 
     train_btn.click(
-        fn=upload_and_train_detection,
-        inputs=[api_input, project_slug_state, dataset_path_state],
-        outputs=[url_output]
+        upload_and_train_detection,
+        inputs=[api_input, slug_state, ds_state],
+        outputs=[url_out]
     )
 
 if __name__ == "__main__":