Select an image from the Dataset tab.
'
data = load_dataset()
item = next((x for x in data["images"] if x["id"] == image_id), None)
if not item:
return 'Image not found.
'
src = image_data_uri(image_id)
boxes = json.dumps(item.get("annotations", []), ensure_ascii=False)
return """
Draw boxes directly on the imageClick + drag + release = create boxChoose the class first
Drag from one corner of the object to the opposite corner, then release. The coordinates are filled automatically; click Save Box. Repeat for every object.
""" % (json.dumps(src), int(item["width"]), int(item["height"]), boxes)
def annotation_preview_image(image_id):
"""Return the real PIL image for Gradio's Image component."""
if not image_id:
return None
p = image_path(image_id)
if not p.exists():
return None
try:
return Image.open(p).convert("RGB")
except Exception:
return None
def annotation_preview_with_boxes(image_id):
image = annotation_preview_image(image_id)
if image is None:
return None
data = load_dataset()
item = next((x for x in data["images"] if x["id"] == image_id), None)
if not item:
return image
out = image.copy()
draw = ImageDraw.Draw(out)
for i, a in enumerate(item.get("annotations", []), 1):
x, y, w, h = a["box"]
draw.rectangle([x, y, x+w, y+h], outline="red", width=5)
label = f"{i}. {a['class']}"
y0 = max(0, y-24)
draw.rectangle([x, y0, x+max(130, len(label)*9), y0+24], fill="red")
draw.text((x+4, y0+4), label, fill="white")
return out
def refresh_editor(image_id):
if not image_id:
return None, "Select an image.", []
data=load_dataset()
item=next((x for x in data["images"] if x["id"]==image_id),None)
if not item:
return None,"Image not found.",[]
return annotation_preview_with_boxes(image_id), f"**{item['filename']}** — {item['width']} × {item['height']} px", item.get("annotations",[])
def draw_annotations(image_id):
"""Render the annotation canvas for the selected image."""
return annotation_canvas_html(image_id)
def add_annotation(image_id, cls, x, y, w, h):
if not image_id:return annotation_preview_with_boxes(image_id),"Select an image first.",[]
if not cls:return annotation_preview_with_boxes(image_id),"Select a class first.",[]
try:x,y,w,h=map(float,[x,y,w,h])
except:return annotation_preview_with_boxes(image_id),"Enter box coordinates first.",[]
if w<=0 or h<=0:return annotation_preview_with_boxes(image_id),"Box must have a width and height.",[]
data=load_dataset()
item=next((z for z in data["images"] if z["id"]==image_id),None)
if not item:return None,"Image not found.",[]
x=max(0,min(x,item["width"]-1)); y=max(0,min(y,item["height"]-1))
w=min(w,item["width"]-x); h=min(h,item["height"]-y)
item.setdefault("annotations",[]).append({"class":cls,"box":[x,y,w,h]})
save_dataset(data)
return annotation_preview_with_boxes(image_id),f"Saved {cls}: [{x:.0f}, {y:.0f}, {w:.0f}, {h:.0f}]",item["annotations"]
def remove_annotation(image_id,index):
if not image_id:return annotation_preview_with_boxes(image_id),"Select an image first.",[]
data=load_dataset()
item=next((z for z in data["images"] if z["id"]==image_id),None)
if not item:return None,"Image not found.",[]
try:idx=int(index)-1
except:return annotation_preview_with_boxes(image_id),"Enter an annotation number.",item.get("annotations",[])
anns=item.get("annotations",[])
if idx<0 or idx>=len(anns):return annotation_preview_with_boxes(image_id),"Annotation number not found.",anns
deleted=anns.pop(idx);save_dataset(data)
return annotation_preview_with_boxes(image_id),f"Deleted annotation {index}: {deleted['class']}",anns
def clear_annotations(image_id):
if not image_id:return annotation_preview_with_boxes(image_id),"Select an image first.",[]
data=load_dataset()
item=next((z for z in data["images"] if z["id"]==image_id),None)
if not item:return None,"Image not found.",[]
item["annotations"]=[];save_dataset(data)
return annotation_preview_with_boxes(image_id),"Annotations cleared.",[]
def save_classes(text):
classes = [x.strip() for x in (text or "").splitlines() if x.strip()]
if not classes:
return "At least one class is required.", gr.update(choices=read_classes()), dataset_status()
if len(set(classes)) != len(classes):
return "Classes must be unique.", gr.update(choices=read_classes()), dataset_status()
CLASSES_FILE.write_text("\n".join(classes) + "\n", encoding="utf-8")
data = load_dataset()
save_dataset(data)
return f"Saved {len(classes)} classes.", gr.update(choices=classes, value=classes[0]), dataset_status()
def handle_annotation_click(image_id, cls, click_state, evt: gr.SelectData):
"""Use two clicks on the real Gradio image to define a box.
First click = top-left corner, second click = opposite corner.
This avoids the unreliable HTML canvas/script path and works in Gradio itself.
"""
if not image_id:
return 0, 0, 0, 0, [], "Select an image first."
if not cls:
return 0, 0, 0, 0, [], "Select a class first."
data = load_dataset()
item = next((x for x in data["images"] if x["id"] == image_id), None)
if not item:
return 0, 0, 0, 0, [], "Image not found."
try:
point = evt.index
px, py = float(point[0]), float(point[1])
except Exception:
return 0, 0, 0, 0, click_state or [], "Could not read the image click position."
px = max(0, min(px, item["width"] - 1))
py = max(0, min(py, item["height"] - 1))
state = list(click_state or [])
if not state:
return round(px), round(py), 0, 0, [px, py], f"First corner: ({px:.0f}, {py:.0f}). Now click the opposite corner."
x0, y0 = state[:2]
x = min(x0, px); y = min(y0, py)
w = abs(px - x0); h = abs(py - y0)
if w < 2 or h < 2:
return round(x), round(y), 0, 0, [], "Box is too small. Click the first corner again."
return round(x), round(y), round(w), round(h), [], f"Box ready: [{x:.0f}, {y:.0f}, {w:.0f}, {h:.0f}] for {cls}. Click Save Box."
def build_coco():
data = load_dataset()
classes = read_classes()
if not classes:
raise RuntimeError("No classes configured.")
items = [x for x in data["images"] if x.get("annotations")]
if not items:
raise RuntimeError("Annotate at least 1 image before training.")
# With only one annotated image, use it for both training and validation so
# the first training run is possible. With 2+ images, use an 80/20 split.
if len(items) == 1:
train_items, val_items = items, items
else:
split = max(1, int(len(items) * 0.8))
if split >= len(items):
split = len(items) - 1
train_items, val_items = items[:split], items[split:]
category_id = {name: i + 1 for i, name in enumerate(classes)}
def make_coco(selected):
images, annotations = [], []
ann_id = 1
for item in selected:
images.append({
"id": item["id"],
"file_name": item["id"] + ".jpg",
"width": item["width"],
"height": item["height"],
})
for ann in item["annotations"]:
x, y, w, h = ann["box"]
annotations.append({
"id": ann_id,
"image_id": item["id"],
"category_id": category_id[ann["class"]],
"bbox": [x, y, w, h],
"area": w*h,
"iscrowd": 0,
})
ann_id += 1
return {
"images": images,
"annotations": annotations,
"categories": [{"id": i+1, "name": n} for i, n in enumerate(classes)]
}
if GENERATED_DIR.exists():
shutil.rmtree(GENERATED_DIR)
for name, selected in [("train", train_items), ("val", val_items)]:
d = GENERATED_DIR / name
(d / "images").mkdir(parents=True, exist_ok=True)
for item in selected:
shutil.copy2(image_path(item["id"]), d / "images" / f"{item['id']}.jpg")
(d / "annotations.json").write_text(
json.dumps(make_coco(selected), indent=2), encoding="utf-8"
)
def run_training(epochs, batch_size, learning_rate):
global _training, _model
try:
_training = {"running": True, "message": "building COCO dataset", "error": None}
build_coco()
_training["message"] = "training RT-DETR"
cmd = [
sys.executable, str(TRAIN_SCRIPT),
"--train-dir", str(GENERATED_DIR / "train"),
"--val-dir", str(GENERATED_DIR / "val"),
"--classes", str(CLASSES_FILE),
"--output-dir", str(MODEL_DIR),
"--epochs", str(int(epochs)),
"--batch-size", str(int(batch_size)),
"--learning-rate", str(float(learning_rate)),
]
result = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True)
if result.returncode != 0:
details = result.stderr.strip() or result.stdout.strip() or f"training process exited with code {result.returncode}"
raise RuntimeError(details[-12000:])
_model = None
_training = {"running": False, "message": "training complete", "error": None}
except Exception as e:
_training = {"running": False, "message": "training failed", "error": str(e)}
@spaces.GPU(duration=120)
def start_training(epochs, batch_size, learning_rate):
"""Start training from the Gradio event itself.
Calling the @spaces.GPU function directly is important on Hugging Face
ZeroGPU: starting it from a normal Python background thread can bypass the
GPU allocation context, making the button appear to do nothing.
"""
global _training
if _training["running"]:
return json.dumps(_training, indent=2)
# Validate parameters before requesting GPU time.
try:
epochs = max(1, int(epochs))
batch_size = max(1, int(batch_size))
learning_rate = float(learning_rate)
if learning_rate <= 0:
raise ValueError("Learning rate must be greater than 0.")
except Exception as e:
_training = {"running": False, "message": "training failed", "error": f"Invalid training settings: {e}"}
return json.dumps(_training, indent=2)
annotated = sum(bool(x.get("annotations")) for x in load_dataset()["images"])
if annotated < 1:
_training = {"running": False, "message": "training failed", "error": "Annotate at least 1 image before training."}
return json.dumps(_training, indent=2)
run_training(epochs, batch_size, learning_rate)
return json.dumps(_training, indent=2)
def training_status():
return json.dumps(_training, indent=2)
@spaces.GPU(duration=60)
def count_image(image):
if image is None:
return None, "Upload an image first.", {}
if not model_ready():
return None, "Model is not trained yet. Go to Training.", {}
try:
proc, detector = load_model()
image = image.convert("RGB") if isinstance(image, Image.Image) else Image.fromarray(image).convert("RGB")
device = next(detector.parameters()).device
inputs = proc(images=image, return_tensors="pt")
inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()}
with torch.inference_mode():
outputs = detector(**inputs)
target_sizes = torch.tensor([[image.height, image.width]], device=device)
result = proc.post_process_object_detection(
outputs, threshold=CONFIDENCE_THRESHOLD, target_sizes=target_sizes
)[0]
detections = []
counts = Counter()
for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
s = float(score.item())
cls = detector.config.id2label[int(label.item())]
coords = [round(float(v), 2) for v in box.tolist()]
detections.append({"class": cls, "confidence": round(s, 4), "box": coords})
counts[cls] += 1
out = image.copy()
draw = ImageDraw.Draw(out)
for d in detections:
x1, y1, x2, y2 = d["box"]
draw.rectangle([x1, y1, x2, y2], outline="red", width=4)
label = f"{d['class']} {d['confidence']:.2f}"
draw.rectangle([x1, max(0, y1-22), x1+max(120, len(label)*8), y1], fill="red")
draw.text((x1+3, max(0, y1-20)), label, fill="white")
response = {
"total": len(detections),
"counts": dict(sorted(counts.items())),
"detections": detections,
}
return out, json.dumps(response, indent=2), response["counts"]
except Exception as e:
return None, f"Counting failed: {e}", {}
# ----- Gradio UI -----
CSS = """
.gradio-container { max-width: 1250px !important; }
h1 { margin-bottom: 0.2rem !important; }
.anno-wrap{width:100%}.anno-toolbar{display:flex;gap:14px;flex-wrap:wrap;padding:10px 12px;margin-bottom:8px;border-radius:10px;background:#20242a}.anno-toolbar span{opacity:.85}.anno-canvas-wrap{width:100%;overflow:auto;border:1px solid #555;border-radius:10px;background:#111;padding:8px}.anno-canvas-wrap canvas{display:block;max-width:none;cursor:crosshair;touch-action:none;margin:auto}.anno-help{padding:8px 2px;opacity:.75}.anno-empty{padding:50px;text-align:center;border:1px dashed #777;border-radius:10px}
#annotation-image img { max-height: 650px !important; object-fit: contain !important; }
.status { padding: 10px 14px; border-radius: 10px; }
"""
with gr.Blocks(title="Ice Cream Dataset + Counter") as demo:
gr.Markdown("# 🍦 Ice Cream Dataset + Counter\nUpload and annotate training images, train RT-DETR, then count ice creams in new images.")
status = gr.Markdown(dataset_status(), elem_classes="status")
with gr.Tab("1 · Dataset"):
gr.Markdown("### Upload training images")
files = gr.Files(file_count="multiple", file_types=["image"], type="filepath", label="Images")
upload_btn = gr.Button("Save Images", variant="primary")
upload_msg = gr.Markdown()
# Dataset selector
image_select = gr.Dropdown(choices=image_choices(), label="Training image", interactive=True)
refresh_btn = gr.Button("Refresh Dataset")
refresh_btn.click(lambda: (dataset_status(), gr.update(choices=image_choices())), None, [status, image_select])
gr.Markdown("### Classes")
class_text = gr.Textbox(value="\n".join(read_classes()), lines=8, label="One class per line")
save_class_btn = gr.Button("Save Classes")
class_msg = gr.Markdown()
# Class dropdown is updated after the Annotate tab creates it.
with gr.Tab("2 · Annotate"):
gr.Markdown("### Annotate training images")
gr.Markdown("Select an image above, choose a class, then **click the first corner and click the opposite corner** of each object. The real uploaded image is shown below. Click **Save Box** after each box.")
with gr.Row():
with gr.Column(scale=3):
annotation_image = gr.Image(value=None, type="pil", interactive=False, label="Training image", height=650, elem_id="annotation-image")
editor_info = gr.Markdown("Select an image from the Dataset tab.")
with gr.Column(scale=1):
ann_class = gr.Dropdown(choices=read_classes(), value=(read_classes()[0] if read_classes() else None), label="Class", interactive=True)
x = gr.Number(label="X (left)", value=0, precision=0)
y = gr.Number(label="Y (top)", value=0, precision=0)
w = gr.Number(label="Width", value=0, precision=0)
h = gr.Number(label="Height", value=0, precision=0)
add_btn = gr.Button("💾 Save Box", variant="primary")
gr.Markdown("**Box method:** click corner 1 → click corner 2 → Save Box.")
delete_index = gr.Number(label="Annotation # to delete", value=1, precision=0)
delete_btn = gr.Button("Delete Box")
clear_btn = gr.Button("Clear All Boxes")
annotations = gr.JSON(label="Saved annotations")
ann_msg = gr.Markdown()
click_state = gr.State([])
def load_annotation_image(image_id):
return annotation_preview_with_boxes(image_id), refresh_editor(image_id)[1], refresh_editor(image_id)[2], []
image_select.change(load_annotation_image, image_select, [annotation_image, editor_info, annotations, click_state])
annotation_image.select(handle_annotation_click, [image_select, ann_class, click_state], [x, y, w, h, click_state, ann_msg])
add_btn.click(add_annotation, [image_select, ann_class, x, y, w, h], [annotation_image, ann_msg, annotations])
delete_btn.click(remove_annotation, [image_select, delete_index], [annotation_image, ann_msg, annotations])
clear_btn.click(clear_annotations, image_select, [annotation_image, ann_msg, annotations])
with gr.Tab("3 · Training"):
gr.Markdown("### Train RT-DETR")
gr.Markdown("Training runs in the Space process. A GPU Space is strongly recommended for practical training speed.")
with gr.Row():
epochs = gr.Number(value=int(os.getenv("EPOCHS", "30")), label="Epochs", precision=0)
batch = gr.Number(value=int(os.getenv("BATCH_SIZE", "2")), label="Batch size", precision=0)
lr = gr.Number(value=float(os.getenv("LEARNING_RATE", "1e-5")), label="Learning rate")
train_btn = gr.Button("🚀 Start Training", variant="primary")
refresh_train = gr.Button("Refresh Training Status")
train_out = gr.Code(value=training_status, language="json", label="Training status")
train_btn.click(start_training, [epochs, batch, lr], train_out)
refresh_train.click(training_status, None, train_out)
with gr.Tab("4 · Count"):
gr.Markdown("### Count ice creams")
count_in = gr.Image(type="pil", sources=["upload", "clipboard"], label="Image to count")
count_btn = gr.Button("🍦 Count", variant="primary")
count_out = gr.Image(label="Detections")
count_json = gr.Code(language="json", label="Detection details")
count_table = gr.JSON(label="Counts by class")
count_btn.click(count_image, count_in, [count_out, count_json, count_table])
# Correct the upload event now that image_select exists.
upload_btn.click(
upload_training_images,
files,
[status, image_select, upload_msg],
preprocess=False,
queue=False,
)
save_class_btn.click(save_classes, class_text, [class_msg, ann_class, status])
# The earlier placeholder event is harmlessly superseded by this real event.
demo.load(lambda: (dataset_status(), gr.update(choices=image_choices()), gr.update(choices=read_classes(), value=(read_classes()[0] if read_classes() else None))),
None, [status, image_select, ann_class])
if __name__ == "__main__":
demo.queue().launch(
server_name="0.0.0.0",
server_port=int(os.getenv("PORT", "7860")),
css=CSS,
)