Update app.py

app.py

@@ -26,66 +26,46 @@ def parse_roboflow_url(url: str):
 
 
 def convert_seg_to_bbox(api_key: str, dataset_url: str):
-    # 1) download the dataset
+    """Download a segmentation dataset from Roboflow, convert to YOLO bboxes, and return before/after galleries."""
     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  # e.g. "/home/user/app/ds-2"
-
-    # --- DEBUG: print out the downloaded directory tree ---
-    print(f"\n=== Downloaded dataset root: {root} ===")
-    for dirpath, dirnames, filenames in os.walk(root):
-        print(f"\nDirectory: {dirpath}")
-        for d in dirnames:
-            print(f"  [DIR ] {d}")
-        for f in filenames:
-            print(f"  [FILE] {f}")
-    print("=== end tree dump ===\n")
-
-    # 2) search for any JSON file with "train" in its name
+    root = dataset.location
+
+    # 1) Find annotation JSON
     ann_file = None
-    for dirpath, _, filenames in os.walk(root):
-        for fname in filenames:
-            if 'train' in fname.lower() and fname.lower().endswith('.json'):
-                ann_file = os.path.join(dirpath, fname)
-                print(f"Found TRAIN annotation file: {ann_file}")
+    for dp, _, files in os.walk(root):
+        for f in files:
+            if 'train' in f.lower() and f.lower().endswith('.json'):
+                ann_file = os.path.join(dp, f)
                 break
         if ann_file:
             break
-
-    # 2b) fallback: first .json anywhere
-    if ann_file is None:
-        for dirpath, _, filenames in os.walk(root):
-            for fname in filenames:
-                if fname.lower().endswith('.json'):
-                    ann_file = os.path.join(dirpath, fname)
-                    print(f"No TRAIN file—falling back to first JSON: {ann_file}")
+    if not ann_file:
+        for dp, _, files in os.walk(root):
+            for f in files:
+                if f.lower().endswith('.json'):
+                    ann_file = os.path.join(dp, f)
                     break
             if ann_file:
                 break
+    if not ann_file:
+        raise FileNotFoundError("No JSON annotations found under %s" % root)
 
-    if ann_file is None:
-        raise FileNotFoundError(f"No JSON annotations found under {root}")
-
-    # 3) load COCO annotations
-    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']}
-
-    # 4) build category→index map
     cat_ids = sorted(c['id'] for c in coco.get('categories', []))
     id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}
 
-    # 5) prepare YOLO output directories
+    # 2) 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)
-    print(f"Preparing YOLOv8 output in: {out_root}")
 
-    # 6) convert each segmentation annotation to a YOLO bbox line
+    # 3) Convert seg→bbox
     annos = {}
     for anno in coco['annotations']:
         img_id = anno['image_id']
@@ -94,59 +74,53 @@ 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
-
-        info = images_info[img_id]
-        iw, ih = info['width'], info['height']
-        line = (
-            f"{id_to_index[anno['category_id']]} "
-            f"{cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
-        )
+        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)
 
-    # 7) locate your image folder (first with any .jpg/.png)
+    # 4) Find images and write labels
     train_img_dir = None
-    for dirpath, _, files in os.walk(root):
+    for dp, _, files in os.walk(root):
         if any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in files):
-            train_img_dir = dirpath
-            print(f"Found image directory: {train_img_dir}")
+            train_img_dir = dp
             break
+    if not train_img_dir:
+        raise FileNotFoundError("No image files found under %s" % root)
 
-    if train_img_dir is None:
-        raise FileNotFoundError(f"No image files found under {root}")
-
-    # 8) 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))
-        lbl_path = os.path.join(lbl_out, fname.rsplit('.', 1)[0] + ".txt")
-        with open(lbl_path, '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, [])))
 
-    # 9) build before/after galleries
+    # 5) Build 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)
 
+        # draw seg polygons
         seg_vis = img.copy()
-        img_id = name_to_id[fname]
         for anno in coco['annotations']:
-            if anno['image_id'] != img_id: continue
-            pts = np.array(anno['segmentation'][0], dtype=np.int32).reshape(-1, 2)
+            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)
 
+        # draw boxes
         box_vis = img.copy()
-        for line in annos.get(img_id, []):
+        for line in annos.get(name_to_id[fname], []):
             _, cxn, cyn, wnorm, hnorm = map(float, line.split())
-            iw, ih = images_info[img_id]['width'], images_info[img_id]['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)
+            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))
@@ -155,19 +129,67 @@ def convert_seg_to_bbox(api_key: str, dataset_url: str):
     return before, after
 
 
+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"
+):
+    """
+    Upload a local detection dataset to Roboflow, generate+train a new version,
+    and return the hosted inference endpoint URL.
+    """
+    rf = Roboflow(api_key=api_key)
+    ws = rf.workspace()
+
+    # 1) upload
+    ws.upload_dataset(
+        dataset_path,
+        project_id,
+        project_license=project_license,
+        project_type=project_type
+    )
+
+    # 2) generate version
+    proj = ws.project(project_id)
+    version_number = proj.generate_version(
+        preprocessing=preprocessing or {},
+        augmentation=augmentation or {}
+    )
+
+    # 3) train
+    proj.version(version_number).train(speed=speed)
+
+    # 4) fetch model endpoint info
+    m = proj.version(str(version_number)).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 Converter")
-    api_input      = gr.Textbox(label="Roboflow API Key", type="password")
-    url_input      = gr.Textbox(label="Roboflow Dataset URL (Segmentation)")
-    run_btn        = gr.Button("Convert")
-    before_gallery = gr.Gallery(label="Before (Segmentation)", columns=5, height="auto")
-    after_gallery  = gr.Gallery(label="After (Bounding Boxes)", columns=5, height="auto")
-
-    run_btn.click(
-        fn=convert_seg_to_bbox,
-        inputs=[api_input, url_input],
-        outputs=[before_gallery, after_gallery]
+    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")
+    train_btn.click(
+        fn=upload_and_train_detection,
+        inputs=[api_input2, project_input, path_input],
+        outputs=[url_output],
     )
 
 if __name__ == "__main__":