Update app.py

app.py

@@ -17,7 +17,7 @@ def parse_roboflow_url(url: str):
     parsed = urlparse(url)
     parts = parsed.path.strip('/').split('/')
     workspace = parts[0]
-    project = parts[1]
+    project   = parts[1]
     try:
         version = int(parts[-1])
     except ValueError:
@@ -28,108 +28,168 @@ def parse_roboflow_url(url: str):
 def convert_seg_to_bbox(api_key: str, dataset_url: str):
     """
     1) Download segmentation dataset from Roboflow
-    2) Convert each mask to its bounding box (YOLO format)
-    3) Preserve original train/valid/test splits
-    4) Return before/after visuals plus (dataset_path, detection_slug)
+    2) Detect JSON‑vs‑mask export
+    3) Convert each mask/polygon to its bounding box (YOLO format)
+    4) Preserve original train/valid/test splits
+    5) Return before/after visuals + (dataset_path, detection_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
+    dataset     = version_obj.download("coco-segmentation")
+    root        = dataset.location
 
-    # 1) Locate all three split JSON files
-    json_files = {}
+    # scan for any .json files
+    all_json = []
     for dp, _, files in os.walk(root):
         for f in files:
-            lf = f.lower()
-            if not lf.endswith('.json'):
-                continue
-            if 'train' in lf:
-                json_files['train'] = os.path.join(dp, f)
-            elif 'valid' in lf or 'val' in lf:
-                json_files['valid'] = os.path.join(dp, f)
-            elif 'test' in lf:
-                json_files['test'] = os.path.join(dp, f)
-    if any(k not in json_files for k in ('train', 'valid', 'test')):
-        raise RuntimeError(f"Missing one of train/valid/test JSONs: {json_files}")
-
-    # 2) Build category → index mapping from the train split
-    train_coco = json.load(open(json_files['train'], 'r'))
-    cat_ids = sorted(c['id'] for c in train_coco.get('categories', []))
-    id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}
-
-    # 3) Aggregate ALL image info & annotations into global dicts
-    global_images_info = {}
-    global_annos = {}
-    for split, jf in json_files.items():
-        coco = json.load(open(jf, 'r'))
-        for img in coco['images']:
-            global_images_info[img['id']] = img
-        for anno in coco['annotations']:
-            xs = anno['segmentation'][0][0::2]
-            ys = anno['segmentation'][0][1::2]
-            xmin, xmax = min(xs), max(xs)
-            ymin, ymax = min(ys), max(ys)
-            w, h = xmax - xmin, ymax - ymin
-            cx, cy = xmin + w/2, ymin + h/2
-            iw = global_images_info[anno['image_id']]['width']
-            ih = global_images_info[anno['image_id']]['height']
-            line = f"{id_to_index[anno['category_id']]} {cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
-            global_annos.setdefault(anno['image_id'], []).append(line)
-
-    # 4) Build a quick map of filename → full path
-    name_to_id = {img['file_name']: img['id'] for img in global_images_info.values()}
-    file_paths = {
-        f: os.path.join(dp, f)
-        for dp, _, files in os.walk(root)
-        for f in files
-        if f in name_to_id
-    }
-
-    # 5) Copy images & write YOLO .txt labels, preserving original splits
-    out_root = tempfile.mkdtemp(prefix="yolov8_")
-    for split in ('train', 'valid', 'test'):
-        coco = json.load(open(json_files[split], 'r'))
-        img_dir = os.path.join(out_root, split, "images")
-        lbl_dir = os.path.join(out_root, split, "labels")
-        os.makedirs(img_dir, exist_ok=True)
-        os.makedirs(lbl_dir, exist_ok=True)
-        for img in coco['images']:
-            fname = img['file_name']
-            shutil.copy(file_paths[fname], os.path.join(img_dir, fname))
-            with open(os.path.join(lbl_dir, fname.rsplit('.', 1)[0] + ".txt"), 'w') as f:
-                f.write("\n".join(global_annos.get(img['id'], [])))
-
-    # 6) Prepare a few before/after examples (random sample across all splits)
-    before, after = [], []
-    all_ids = list(global_images_info.keys())
-    sample_ids = random.sample(all_ids, min(5, len(all_ids)))
-    for img_id in sample_ids:
-        fname = global_images_info[img_id]['file_name']
-        img = cv2.cvtColor(cv2.imread(file_paths[fname]), cv2.COLOR_BGR2RGB)
-
-        # draw segmentation outlines
-        seg_vis = img.copy()
-        for jf in json_files.values():
-            coco = json.load(open(jf, 'r'))
-            for anno in coco['annotations']:
-                if anno['image_id'] != img_id:
+            if f.lower().endswith(".json"):
+                all_json.append(os.path.join(dp, f))
+
+    if len(all_json) >= 3 and any("train" in os.path.basename(p).lower() for p in all_json):
+        # --- COCO‑JSON export branch ---
+        # locate train/valid/test JSONs
+        json_splits = {}
+        for path in all_json:
+            fn = os.path.basename(path).lower()
+            if "train" in fn:
+                json_splits["train"] = path
+            elif "val" in fn or "valid" in fn:
+                json_splits["valid"] = path
+            elif "test" in fn:
+                json_splits["test"] = path
+        if any(s not in json_splits for s in ("train", "valid", "test")):
+            raise RuntimeError(f"Missing one of train/valid/test JSONs: {json_splits}")
+
+        # build category → index from train.json
+        train_coco = json.load(open(json_splits["train"], "r"))
+        cat_ids    = sorted(c["id"] for c in train_coco.get("categories", []))
+        id2idx     = {cid: i for i, cid in enumerate(cat_ids)}
+
+        # aggregate images_info & annotations
+        images_info = {}
+        annos       = {}
+        for split, jf in json_splits.items():
+            coco = json.load(open(jf, "r"))
+            for img in coco["images"]:
+                images_info[img["id"]] = img
+            for a in coco["annotations"]:
+                xs = a["segmentation"][0][0::2]
+                ys = a["segmentation"][0][1::2]
+                xmin, xmax = min(xs), max(xs)
+                ymin, ymax = min(ys), max(ys)
+                w, h       = xmax - xmin, ymax - ymin
+                cx, cy     = xmin + w/2, ymin + h/2
+                iw = images_info[a["image_id"]]["width"]
+                ih = images_info[a["image_id"]]["height"]
+                line = (
+                    f"{id2idx[a['category_id']]} "
+                    f"{cx/iw:.6f} {cy/ih:.6f} {w/iw:.6f} {h/ih:.6f}"
+                )
+                annos.setdefault(a["image_id"], []).append(line)
+
+        # build filename → path map
+        name2id = {img["file_name"]: img["id"] for img in images_info.values()}
+        filemap = {
+            f: os.path.join(dp, f)
+            for dp, _, files in os.walk(root)
+            for f in files
+            if f in name2id
+        }
+
+        # write out per‑split folders
+        out_root = tempfile.mkdtemp(prefix="yolov8_")
+        for split in ("train", "valid", "test"):
+            coco = json.load(open(json_splits[split], "r"))
+            img_dir = os.path.join(out_root, split, "images")
+            lbl_dir = os.path.join(out_root, split, "labels")
+            os.makedirs(img_dir, exist_ok=True)
+            os.makedirs(lbl_dir, exist_ok=True)
+            for img in coco["images"]:
+                fn   = img["file_name"]
+                src  = filemap[fn]
+                dst  = os.path.join(img_dir, fn)
+                txtp = os.path.join(lbl_dir, fn.rsplit(".", 1)[0] + ".txt")
+                shutil.copy(src, dst)
+                with open(txtp, "w") as f:
+                    f.write("\n".join(annos.get(img["id"], [])))
+
+    else:
+        # --- Segmentation‐Masks export branch ---
+        splits = ["train", "valid", "test"]
+        # detect masks subfolder name
+        mask_names = ("masks", "mask", "labels")
+        out_root   = tempfile.mkdtemp(prefix="yolov8_")
+
+        for split in splits:
+            img_dir_src = os.path.join(root, split, "images")
+            # find which subdir holds the PNG masks
+            mdir = None
+            for m in mask_names:
+                candidate = os.path.join(root, split, m)
+                if os.path.isdir(candidate):
+                    mdir = candidate
+                    break
+            if mdir is None:
+                raise RuntimeError(f"No masks folder found under {split}/ (checked {mask_names})")
+
+            img_dir_dst = os.path.join(out_root, split, "images")
+            lbl_dir_dst = os.path.join(out_root, split, "labels")
+            os.makedirs(img_dir_dst, exist_ok=True)
+            os.makedirs(lbl_dir_dst, exist_ok=True)
+
+            for fn in os.listdir(img_dir_src):
+                if not fn.lower().endswith((".jpg", ".png")):
                     continue
-                pts = np.array(anno['segmentation'][0], np.int32).reshape(-1, 2)
-                cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)
+                src_img  = os.path.join(img_dir_src, fn)
+                src_mask = os.path.join(mdir, fn)
+                img      = cv2.imread(src_img)
+                h, w     = img.shape[:2]
 
-        # draw bounding boxes
-        box_vis = img.copy()
-        for line in global_annos.get(img_id, []):
-            _, cxn, cyn, wnorm, hnorm = map(float, line.split())
-            iw = global_images_info[img_id]['width']
-            ih = global_images_info[img_id]['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)
+                # read mask & binarize
+                mask = cv2.imread(src_mask, cv2.IMREAD_GRAYSCALE)
+                _, binm = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+                ys, xs  = np.nonzero(binm)
+                if len(xs) == 0:
+                    lines = []
+                else:
+                    xmin, xmax = xs.min(), xs.max()
+                    ymin, ymax = ys.min(), ys.max()
+                    bw, bh     = xmax - xmin, ymax - ymin
+                    cx, cy     = xmin + bw/2, ymin + bh/2
+                    # assume single class → index 0
+                    lines = [f"0 {cx/w:.6f} {cy/h:.6f} {bw/w:.6f} {bh/h:.6f}"]
+
+                # copy image + write YOLO text
+                dst_img = os.path.join(img_dir_dst, fn)
+                dst_txt = os.path.join(lbl_dir_dst, fn.rsplit(".",1)[0] + ".txt")
+                shutil.copy(src_img, dst_img)
+                with open(dst_txt, "w") as f:
+                    f.write("\n".join(lines))
 
+    # --- prepare before/after galleries (random sample across out_root) ---
+    before, after = [], []
+    all_imgs = []
+    for split in ("train","valid","test"):
+        for fn in os.listdir(os.path.join(out_root, split, "images")):
+            path = os.path.join(out_root, split, "images", fn)
+            all_imgs.append(path)
+    sample = random.sample(all_imgs, min(5, len(all_imgs)))
+    for img_path in sample:
+        fn       = os.path.basename(img_path)
+        img      = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
+        # draw mask outline if available (JSON branch) else read mask again
+        seg_vis  = img.copy()
+        box_vis  = img.copy()
+        # overlay all .txt bboxes
+        txtp     = img_path.replace("/images/", "/labels/").rsplit(".",1)[0] + ".txt"
+        w, h     = img.shape[1], img.shape[0]
+        for line in open(txtp):
+            _, cxn, cyn, wnorm, hnorm = map(float, line.split())
+            bw, bh = int(wnorm * w), int(hnorm * h)
+            x0      = int(cxn * w - bw/2)
+            y0      = int(cyn * h - bh/2)
+            cv2.rectangle(box_vis, (x0,y0), (x0+bw, y0+bh), (0,255,0), 2)
         before.append(Image.fromarray(seg_vis))
         after.append(Image.fromarray(box_vis))
 
@@ -142,13 +202,8 @@ def upload_and_train_detection(
     detection_slug: str,
     dataset_path: str,
     project_license: str = "MIT",
-    project_type: str = "object-detection"
+    project_type: str  = "object-detection"
 ):
-    """
-    Uploads the converted dataset (with preserved splits) to Roboflow,
-    creates or fetches a detection project, and kicks off training.
-    Returns the hosted model URL.
-    """
     rf = Roboflow(api_key=api_key)
     ws = rf.workspace()
 
@@ -166,36 +221,32 @@ def upload_and_train_detection(
         else:
             raise
 
-    # upload entire split folder
+    # upload and kick off train
     _, real_slug = proj.id.rsplit("/", 1)
-    ws.upload_dataset(
-        dataset_path,
-        real_slug,
-        project_license=project_license,
-        project_type=project_type
-    )
+    ws.upload_dataset(dataset_path, real_slug,
+                      project_license=project_license,
+                      project_type=project_type)
 
-    # generate new version (with fallback slug bump)
     try:
-        version_num = proj.generate_version(settings={"augmentation": {}, "preprocessing": {}})
+        version_num = proj.generate_version(settings={"augmentation":{}, "preprocessing":{}})
     except RuntimeError as e:
         msg = str(e).lower()
         if "unsupported request" in msg or "does not exist" in msg:
+            # slug bump fallback
             new_slug = real_slug + "-v2"
             proj = ws.create_project(
-                new_slug,
-                annotation=project_type,
+                new_slug, annotation=project_type,
                 project_type=project_type,
                 project_license=project_license
             )
             ws.upload_dataset(dataset_path, new_slug,
                               project_license=project_license,
                               project_type=project_type)
-            version_num = proj.generate_version(settings={"augmentation": {}, "preprocessing": {}})
+            version_num = proj.generate_version(settings={"augmentation":{}, "preprocessing":{}})
         else:
             raise
 
-    # wait for generation, then train
+    # wait for generation then train
     for _ in range(20):
         try:
             model = proj.version(str(version_num)).train()
@@ -204,24 +255,25 @@ def upload_and_train_detection(
             if "still generating" in str(e).lower():
                 time.sleep(5)
                 continue
-            raise
+            else:
+                raise
     else:
-        raise RuntimeError("Dataset version did not finish generating in time; try again later.")
+        raise RuntimeError("Version generation timed out, try again later.")
 
     return f"{model['base_url']}{model['id']}?api_key={api_key}"
 
 
 # --- Gradio UI ---
 with gr.Blocks() as app:
-    gr.Markdown("## 🔄 Seg→BBox + Auto-Upload/Train")
+    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, label="Converted Dataset Path")
-    slug_state = gr.Textbox(visible=False, label="Detection Project Slug")
+    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, label="Dataset Path")
+    slug_state= gr.Textbox(visible=False, label="Detection Slug")
 
     run_btn.click(
         convert_seg_to_bbox,
@@ -231,7 +283,7 @@ with gr.Blocks() as app:
 
     gr.Markdown("## 🚀 Upload & Train Detection Model")
     train_btn = gr.Button("Upload & Train")
-    url_out = gr.Textbox(label="Hosted Model Endpoint URL")
+    url_out   = gr.Textbox(label="Hosted Model URL")
 
     train_btn.click(
         upload_and_train_detection,