File size: 7,178 Bytes
bc6cd18
 
 
 
 
 
43b1849
bc6cd18
 
43b1849
bc6cd18
eecca4c
bc6cd18
 
43b1849
bc6cd18
 
 
 
 
 
 
 
 
 
 
 
43b1849
a711e94
6c1d7d9
533c19c
eecca4c
a711e94
bc6cd18
6c1d7d9
 
bc6cd18
ac01980
 
eecca4c
6d361b8
ac01980
 
6c1d7d9
ac01980
6d361b8
43b1849
 
ac01980
 
 
6c1d7d9
ac01980
99c4ece
 
 
ac01980
a711e94
99c4ece
6c1d7d9
 
 
bc6cd18
 
eecca4c
bc6cd18
 
 
 
 
 
eecca4c
bc6cd18
6c1d7d9
 
 
6d361b8
a711e94
 
 
6c1d7d9
a711e94
 
 
eecca4c
a711e94
ac01980
6c1d7d9
a711e94
6d361b8
a711e94
6c1d7d9
228a29d
6c1d7d9
a711e94
 
6d361b8
bc6cd18
6d361b8
6c1d7d9
a711e94
bc6cd18
eecca4c
43b1849
a711e94
 
 
bc6cd18
 
6c1d7d9
 
ac01980
6c1d7d9
533c19c
bc6cd18
 
a711e94
6d361b8
6c1d7d9
533c19c
eecca4c
 
533c19c
bc6cd18
43b1849
 
bc6cd18
eecca4c
6c1d7d9
533c19c
bc6cd18
 
a711e94
ac01980
6c1d7d9
 
ac01980
 
 
 
eecca4c
6c1d7d9
 
eecca4c
6c1d7d9
 
 
 
 
 
eecca4c
ac01980
 
6c1d7d9
533c19c
a711e94
533c19c
 
 
ac01980
 
eecca4c
 
ac01980
eecca4c
6c1d7d9
533c19c
ac01980
eecca4c
533c19c
a711e94
ac01980
 
bc6cd18
6c1d7d9
 
 
 
 
 
 
 
 
 
 
 
a711e94
6c1d7d9
 
a711e94
 
6c1d7d9
 
 
 
 
a711e94
6c1d7d9
ac01980
6c1d7d9
 
7838f2e
bc6cd18
6c1d7d9
 
 
 
a711e94
bc6cd18
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
import os
import json
import random
import shutil
import tempfile
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):
    """Extract workspace, project name, and version from a Roboflow URL."""
    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):
    """
    Download a segmentation dataset from Roboflow,
    convert masks β†’ YOLOv8 bboxes,
    and return (before, after) galleries + local YOLO dataset path + 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)
    dataset = version_obj.download("coco-segmentation")
    root = dataset.location

    # 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():
                ann_file = os.path.join(dp, f)
                break
        if ann_file:
            break
    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(f"No JSON annotations found 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", []))
    id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}

    # prepare YOLOv8 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
    annos = {}
    for a in coco["annotations"]:
        pid = a["image_id"]
        poly = a["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
    for dp, _, files in os.walk(root):
        if any(f.lower().endswith((".jpg", ".png", ".jpeg")) for f in files):
            img_dir = dp
            break
    if not img_dir:
        raise FileNotFoundError(f"No images found 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)
        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, [])))

    # build preview 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)

        seg_vis = img.copy()
        for a in coco["annotations"]:
            if a["image_id"] != fname2id[fn]:
                continue
            pts = np.array(a["segmentation"][0], np.int32).reshape(-1, 2)
            cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)

        box_vis = img.copy()
        for line in annos.get(fname2id[fn], []):
            _, cxn, cyn, wnorm, hnorm = map(float, line.split())
            iw, ih = images_info[fname2id[fn]]["width"], images_info[fname2id[fn]]["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)

        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


def upload_and_train_detection(api_key: str, project_slug: str, dataset_path: str):
    """
    Given a YOLOv8 dataset folder, upload β†’ version β†’ train β†’
    return inference endpoint URL. Auto‐creates the project if needed.
    """
    rf = Roboflow(api_key=api_key)
    ws = rf.workspace()

    # get or create the detection project
    try:
        proj = ws.project(project_slug)
    except Exception:
        proj = ws.create_project(
            project_name=project_slug,
            project_type="object-detection",
            project_license="MIT"
        )

    # upload the dataset
    ws.upload_dataset(
        dataset_path,
        proj.id,
        num_workers=10,
        project_license="MIT",
        project_type="object-detection",
        batch_name=None,
        num_retries=0
    )

    # generate a new version
    new_v = proj.generate_version(settings={"preprocessing": {}, "augmentation": {}})

    # train (fast)
    version = proj.version(new_v)
    version.train(speed="fast")

    # return the hosted inference URL
    m = version.model
    return f"{m['base_url']}{m['id']}?api_key={api_key}"


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("---")

    # ─ Train UI ───────────────────────────────────────────
    btn_t = gr.Button("Upload & Train Detection Model")
    endpoint = gr.Textbox(label="Hosted Detection Endpoint URL")

    btn_t.click(
        fn=upload_and_train_detection,
        inputs=[api, state_slug, state_path],
        outputs=[endpoint]
    )

    gr.Markdown(
        "> 1) Paste your segmentation URL and Convert.  \n"
        "> 2) Then Upload & Train to instantly get your detector’s endpoint."
    )

if __name__ == "__main__":
    app.launch()