Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import random | |
| import shutil | |
| import tempfile | |
| import time | |
| from urllib.parse import urlparse | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| import gradio as gr | |
| from roboflow import Roboflow | |
| def parse_roboflow_url(url: str): | |
| parsed = urlparse(url) | |
| parts = parsed.path.strip('/').split('/') | |
| workspace = parts[0] | |
| project = parts[1] | |
| try: | |
| version = int(parts[-1]) | |
| except ValueError: | |
| version = int(parts[-2]) | |
| return workspace, project, version | |
| def convert_seg_to_bbox(api_key: str, dataset_url: str): | |
| """ | |
| 1) Download segmentation dataset from Roboflow | |
| 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 | |
| # scan for any .json files | |
| all_json = [] | |
| for dp, _, files in os.walk(root): | |
| for f in files: | |
| 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 | |
| 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] | |
| # 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)) | |
| detection_slug = proj_name + "-detection" | |
| return before, after, out_root, detection_slug | |
| def upload_and_train_detection( | |
| api_key: str, | |
| detection_slug: str, | |
| dataset_path: str, | |
| project_license: str = "MIT", | |
| project_type: str = "object-detection" | |
| ): | |
| rf = Roboflow(api_key=api_key) | |
| ws = rf.workspace() | |
| # get or create project | |
| try: | |
| proj = ws.project(detection_slug) | |
| except Exception as e: | |
| if "does not exist" in str(e).lower(): | |
| proj = ws.create_project( | |
| detection_slug, | |
| annotation=project_type, | |
| project_type=project_type, | |
| project_license=project_license | |
| ) | |
| else: | |
| raise | |
| # 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) | |
| try: | |
| 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, | |
| 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":{}}) | |
| else: | |
| raise | |
| # wait for generation then train | |
| for _ in range(20): | |
| try: | |
| model = proj.version(str(version_num)).train() | |
| break | |
| except RuntimeError as e: | |
| if "still generating" in str(e).lower(): | |
| time.sleep(5) | |
| continue | |
| else: | |
| raise | |
| else: | |
| 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") | |
| 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="Dataset Path") | |
| slug_state= gr.Textbox(visible=False, label="Detection Slug") | |
| run_btn.click( | |
| convert_seg_to_bbox, | |
| inputs=[api_input, url_input], | |
| outputs=[before_g, after_g, ds_state, slug_state] | |
| ) | |
| gr.Markdown("## 🚀 Upload & Train Detection Model") | |
| train_btn = gr.Button("Upload & Train") | |
| url_out = gr.Textbox(label="Hosted Model URL") | |
| train_btn.click( | |
| upload_and_train_detection, | |
| inputs=[api_input, slug_state, ds_state], | |
| outputs=[url_out] | |
| ) | |
| if __name__ == "__main__": | |
| app.launch() | |