Files
su2-img/scripts/annotator.py
Lukáš Trkan f4f2352ec3 update
2026-04-29 08:53:16 +02:00

463 lines
16 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
Anotátor dlaždic pro trénink YOLOv8
Formát: YOLO class cx cy w h (normalizovaný 01)
Anotace se ukládají do data/hk_custom/train/labels/
Obrázky se kopírují do data/hk_custom/train/images/
"""
import json, shutil, urllib.parse
from http.server import HTTPServer, BaseHTTPRequestHandler
from pathlib import Path
TILES_DIR = Path("tiles_vrchlabi")
OUT_IMG_DIR = Path("data/vrchlabi_custom/train/images")
OUT_LBL_DIR = Path("data/vrchlabi_custom/train/labels")
OUT_IMG_DIR.mkdir(parents=True, exist_ok=True)
OUT_LBL_DIR.mkdir(parents=True, exist_ok=True)
TILE_SIZE = 256
CLASS_NAMES = ["car", "van", "truck", "bus"]
CLASS_COLORS = ["#00dc00", "#ffdc00", "#dc0000", "#0078ff"]
HTML = r"""<!DOCTYPE html>
<html lang="cs">
<head>
<meta charset="UTF-8">
<title>Anotátor vozidel</title>
<style>
* { box-sizing: border-box; margin: 0; padding: 0; }
body { background: #1a1a1a; color: #eee; font-family: monospace; display: flex; height: 100vh; overflow: hidden; }
#sidebar {
width: 260px; min-width: 260px; background: #252525; border-right: 1px solid #333;
display: flex; flex-direction: column; padding: 12px; gap: 10px; overflow-y: auto;
}
#main { flex: 1; display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 10px; padding: 10px; }
h2 { font-size: 13px; color: #aaa; text-transform: uppercase; letter-spacing: 1px; margin-bottom: 4px; }
.cls-btn {
display: flex; align-items: center; gap: 8px; padding: 8px 10px;
border: 2px solid transparent; border-radius: 6px; cursor: pointer;
background: #333; color: #eee; font-family: monospace; font-size: 13px; width: 100%;
}
.cls-btn:hover { background: #444; }
.cls-btn.active { border-color: #fff; background: #3a3a3a; }
.cls-dot { width: 14px; height: 14px; border-radius: 3px; flex-shrink: 0; }
.cls-key { margin-left: auto; color: #777; font-size: 11px; }
#tile-name { font-size: 11px; color: #777; word-break: break-all; }
#ann-count { font-size: 12px; color: #aaa; }
#progress { font-size: 11px; color: #666; }
.btn {
padding: 7px 10px; border-radius: 5px; border: none; cursor: pointer;
font-family: monospace; font-size: 12px; width: 100%;
}
.btn-save { background: #1a6b2a; color: #eee; }
.btn-save:hover { background: #218838; }
.btn-skip { background: #333; color: #aaa; }
.btn-skip:hover { background: #444; }
.btn-del { background: #5c1a1a; color: #eee; }
.btn-del:hover { background: #8b2020; }
hr { border: none; border-top: 1px solid #333; }
#canvas-wrap {
position: relative; cursor: crosshair;
box-shadow: 0 0 0 2px #444;
}
canvas { display: block; }
#nav { display: flex; gap: 8px; align-items: center; }
.nav-btn {
padding: 6px 16px; background: #333; border: none; color: #eee;
border-radius: 5px; cursor: pointer; font-family: monospace; font-size: 13px;
}
.nav-btn:hover { background: #444; }
#tile-idx { font-size: 12px; color: #777; min-width: 80px; text-align: center; }
#filter-wrap { display: flex; gap: 6px; align-items: center; font-size: 12px; }
#filter-wrap label { color: #aaa; }
#filter-wrap select { background: #333; color: #eee; border: 1px solid #555; border-radius: 4px; padding: 3px 6px; font-family: monospace; font-size: 12px; }
#shortcut-help { font-size: 10px; color: #555; line-height: 1.7; }
.saved-badge { color: #00dc00; font-size: 11px; display: none; }
</style>
</head>
<body>
<div id="sidebar">
<h2>Třída vozidla</h2>
<div id="cls-buttons"></div>
<hr>
<div id="filter-wrap">
<label>Zobrazit:</label>
<select id="filter">
<option value="all">vše</option>
<option value="unannotated" selected>neanotované</option>
<option value="annotated">anotované</option>
</select>
</div>
<hr>
<h2>Aktuální dlaždice</h2>
<div id="tile-name">—</div>
<div id="ann-count">0 anotací</div>
<div id="progress">—</div>
<span class="saved-badge" id="saved-badge">✓ uloženo</span>
<hr>
<button class="btn btn-save" onclick="save()">💾 Uložit [S]</button>
<button class="btn btn-del" onclick="deleteSelected()">🗑 Smazat vybraný [Del]</button>
<button class="btn btn-skip" onclick="next()">Přeskočit [→]</button>
<hr>
<div id="shortcut-help">
14 &nbsp; vybrat třídu<br>
S &nbsp;&nbsp;&nbsp; uložit<br>
Del &nbsp; smazat vybraný<br>
← → &nbsp; předchozí / další<br>
Z &nbsp;&nbsp;&nbsp; zpět (undo)<br>
Esc &nbsp; zrušit kresbu
</div>
</div>
<div id="main">
<div id="nav">
<button class="nav-btn" onclick="prev()">◀</button>
<span id="tile-idx">0 / 0</span>
<button class="nav-btn" onclick="next()">▶</button>
</div>
<div id="canvas-wrap">
<canvas id="canvas"></canvas>
</div>
</div>
<script>
const SCALE = 3; // zvětšení 256→768 px
const TILE_PX = 256;
const DISP_PX = TILE_PX * SCALE;
const CLS_NAMES = ["car","van","truck","bus"];
const CLS_COLORS = ["#00dc00","#ffdc00","#dc0000","#0078ff"];
let tiles = [];
let tileIdx = 0;
let annotations = []; // [{cls, cx, cy, w, h}] normalizované 01
let selectedIdx = -1;
let drawing = false;
let drawStart = null;
let drawCurrent = null;
let activeCls = 0;
let dirty = false;
const canvas = document.getElementById("canvas");
const ctx = canvas.getContext("2d");
canvas.width = DISP_PX;
canvas.height = DISP_PX;
// --- Init class buttons ---
const clsDiv = document.getElementById("cls-buttons");
CLS_NAMES.forEach((name, i) => {
const btn = document.createElement("button");
btn.className = "cls-btn" + (i === 0 ? " active" : "");
btn.id = "cls-btn-" + i;
btn.innerHTML = `<span class="cls-dot" style="background:${CLS_COLORS[i]}"></span>${name}<span class="cls-key">[${i+1}]</span>`;
btn.onclick = () => setClass(i);
clsDiv.appendChild(btn);
});
function setClass(i) {
activeCls = i;
document.querySelectorAll(".cls-btn").forEach((b,j) => b.classList.toggle("active", j===i));
}
// --- Load tile list ---
async function loadTiles() {
const filter = document.getElementById("filter").value;
const res = await fetch("/api/tiles?filter=" + filter);
tiles = await res.json();
tileIdx = 0;
loadTile();
}
document.getElementById("filter").addEventListener("change", loadTiles);
async function loadTile() {
if (!tiles.length) { document.getElementById("tile-name").textContent = "Žádné dlaždice"; return; }
const name = tiles[tileIdx];
document.getElementById("tile-name").textContent = name;
document.getElementById("tile-idx").textContent = `${tileIdx+1} / ${tiles.length}`;
// Načti existující anotace
const res = await fetch("/api/ann/" + encodeURIComponent(name));
annotations = await res.json();
selectedIdx = -1;
drawing = false;
dirty = false;
document.getElementById("saved-badge").style.display = "none";
updateAnnCount();
// Nakresli dlaždici
const img = new Image();
img.onload = () => { ctx.drawImage(img, 0, 0, DISP_PX, DISP_PX); drawAnnotations(); };
img.src = "/tiles/" + encodeURIComponent(name);
}
function updateAnnCount() {
document.getElementById("ann-count").textContent = annotations.length + " anotací";
const total = tiles.length;
document.getElementById("progress").textContent = `Dlaždice ${tileIdx+1}/${total}`;
}
// --- Draw ---
function drawAnnotations(preview) {
const img = new Image();
img.onload = () => {
ctx.drawImage(img, 0, 0, DISP_PX, DISP_PX);
annotations.forEach((ann, i) => {
const x = (ann.cx - ann.w/2) * DISP_PX;
const y = (ann.cy - ann.h/2) * DISP_PX;
const w = ann.w * DISP_PX;
const h = ann.h * DISP_PX;
const color = CLS_COLORS[ann.cls];
ctx.strokeStyle = color;
ctx.lineWidth = i === selectedIdx ? 3 : 2;
ctx.strokeRect(x, y, w, h);
if (i === selectedIdx) {
ctx.fillStyle = color + "33";
ctx.fillRect(x, y, w, h);
}
ctx.fillStyle = color;
ctx.font = "bold 11px monospace";
ctx.fillText(CLS_NAMES[ann.cls], x + 2, y + 12);
});
if (preview) {
const {x1,y1,x2,y2} = preview;
ctx.strokeStyle = CLS_COLORS[activeCls];
ctx.lineWidth = 2;
ctx.setLineDash([4,3]);
ctx.strokeRect(x1, y1, x2-x1, y2-y1);
ctx.setLineDash([]);
}
};
img.src = "/tiles/" + encodeURIComponent(tiles[tileIdx]);
}
// --- Mouse ---
function getPos(e) {
const r = canvas.getBoundingClientRect();
return { x: e.clientX - r.left, y: e.clientY - r.top };
}
canvas.addEventListener("mousedown", e => {
if (e.button !== 0) return;
const p = getPos(e);
// Check click on existing annotation
let hit = -1;
for (let i = annotations.length-1; i >= 0; i--) {
const ann = annotations[i];
const x = (ann.cx - ann.w/2) * DISP_PX;
const y = (ann.cy - ann.h/2) * DISP_PX;
const w = ann.w * DISP_PX;
const h = ann.h * DISP_PX;
if (p.x >= x && p.x <= x+w && p.y >= y && p.y <= y+h) { hit = i; break; }
}
if (hit >= 0) { selectedIdx = hit; drawAnnotations(); return; }
selectedIdx = -1;
drawing = true;
drawStart = p;
drawCurrent = p;
});
canvas.addEventListener("mousemove", e => {
if (!drawing) return;
drawCurrent = getPos(e);
const x1 = Math.min(drawStart.x, drawCurrent.x);
const y1 = Math.min(drawStart.y, drawCurrent.y);
const x2 = Math.max(drawStart.x, drawCurrent.x);
const y2 = Math.max(drawStart.y, drawCurrent.y);
drawAnnotations({x1,y1,x2,y2});
});
canvas.addEventListener("mouseup", e => {
if (!drawing) return;
drawing = false;
const p = getPos(e);
const x1 = Math.min(drawStart.x, p.x) / DISP_PX;
const y1 = Math.min(drawStart.y, p.y) / DISP_PX;
const x2 = Math.max(drawStart.x, p.x) / DISP_PX;
const y2 = Math.max(drawStart.y, p.y) / DISP_PX;
const w = x2 - x1, h = y2 - y1;
if (w < 0.005 || h < 0.005) { drawAnnotations(); return; }
annotations.push({ cls: activeCls, cx: x1+w/2, cy: y1+h/2, w, h });
selectedIdx = annotations.length - 1;
dirty = true;
updateAnnCount();
drawAnnotations();
});
// --- Actions ---
async function save() {
if (!tiles.length) return;
const name = tiles[tileIdx];
await fetch("/api/ann/" + encodeURIComponent(name), {
method: "POST",
headers: {"Content-Type":"application/json"},
body: JSON.stringify(annotations)
});
dirty = false;
const badge = document.getElementById("saved-badge");
badge.style.display = "inline";
setTimeout(() => badge.style.display = "none", 1500);
// Reload tile list (annotated status may change)
const filter = document.getElementById("filter").value;
if (filter === "unannotated") {
const res = await fetch("/api/tiles?filter=unannotated");
const newTiles = await res.json();
// Stay at same index if possible
tiles = newTiles;
if (tileIdx >= tiles.length) tileIdx = Math.max(0, tiles.length - 1);
loadTile();
}
}
function deleteSelected() {
if (selectedIdx < 0) return;
annotations.splice(selectedIdx, 1);
selectedIdx = -1;
dirty = true;
updateAnnCount();
drawAnnotations();
}
function next() { if (dirty) save(); tileIdx = (tileIdx+1) % tiles.length; loadTile(); }
function prev() { if (dirty) save(); tileIdx = (tileIdx - 1 + tiles.length) % tiles.length; loadTile(); }
// --- Keyboard ---
document.addEventListener("keydown", e => {
if (e.key >= "1" && e.key <= "4") { setClass(parseInt(e.key)-1); return; }
if (e.key === "s" || e.key === "S") { save(); return; }
if (e.key === "Delete" || e.key === "Backspace") { deleteSelected(); return; }
if (e.key === "ArrowRight") { next(); return; }
if (e.key === "ArrowLeft") { prev(); return; }
if (e.key === "Escape") { drawing = false; drawAnnotations(); return; }
if ((e.key === "z" || e.key === "Z") && (e.ctrlKey || e.metaKey)) {
annotations.pop(); dirty = true; selectedIdx = -1; updateAnnCount(); drawAnnotations(); return;
}
});
loadTiles();
</script>
</body>
</html>
"""
class Handler(BaseHTTPRequestHandler):
def log_message(self, format, *args): pass # potlač logy
def send_json(self, data, code=200):
body = json.dumps(data).encode()
self.send_response(code)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.send_header("Access-Control-Allow-Origin", "*")
self.end_headers()
self.wfile.write(body)
def do_GET(self):
path = urllib.parse.urlparse(self.path)
qs = urllib.parse.parse_qs(path.query)
if path.path == "/":
body = HTML.encode()
self.send_response(200)
self.send_header("Content-Type", "text/html; charset=utf-8")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
elif path.path == "/api/tiles":
filt = qs.get("filter", ["all"])[0]
all_tiles = sorted(p.name for p in TILES_DIR.glob("*.jpg"))
if filt == "annotated":
all_tiles = [t for t in all_tiles if (OUT_LBL_DIR / (Path(t).stem + ".txt")).exists()]
elif filt == "unannotated":
all_tiles = [t for t in all_tiles if not (OUT_LBL_DIR / (Path(t).stem + ".txt")).exists()]
self.send_json(all_tiles)
elif path.path.startswith("/api/ann/"):
name = urllib.parse.unquote(path.path[len("/api/ann/"):])
lbl = OUT_LBL_DIR / (Path(name).stem + ".txt")
if not lbl.exists():
self.send_json([])
else:
anns = []
for line in lbl.read_text().splitlines():
parts = line.strip().split()
if len(parts) == 5:
cls, cx, cy, w, h = int(parts[0]), *map(float, parts[1:])
anns.append({"cls": cls, "cx": cx, "cy": cy, "w": w, "h": h})
self.send_json(anns)
elif path.path.startswith("/tiles/"):
name = urllib.parse.unquote(path.path[len("/tiles/"):])
img_path = TILES_DIR / name
if not img_path.exists():
self.send_response(404); self.end_headers(); return
data = img_path.read_bytes()
self.send_response(200)
self.send_header("Content-Type", "image/jpeg")
self.send_header("Content-Length", str(len(data)))
self.end_headers()
self.wfile.write(data)
else:
self.send_response(404); self.end_headers()
def do_POST(self):
if self.path.startswith("/api/ann/"):
name = urllib.parse.unquote(self.path[len("/api/ann/"):])
length = int(self.headers.get("Content-Length", 0))
body = json.loads(self.rfile.read(length))
# Ulož label
lbl = OUT_LBL_DIR / (Path(name).stem + ".txt")
lines = [f"{a['cls']} {a['cx']:.6f} {a['cy']:.6f} {a['w']:.6f} {a['h']:.6f}" for a in body]
lbl.write_text("\n".join(lines) + ("\n" if lines else ""))
# Zkopíruj obrázek pokud ještě není
dst = OUT_IMG_DIR / name
if not dst.exists():
shutil.copy2(TILES_DIR / name, dst)
self.send_json({"ok": True})
else:
self.send_response(404); self.end_headers()
def do_OPTIONS(self):
self.send_response(200)
self.send_header("Access-Control-Allow-Origin", "*")
self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
self.send_header("Access-Control-Allow-Headers", "Content-Type")
self.end_headers()
if __name__ == "__main__":
import webbrowser, threading
port = 8765
server = HTTPServer(("127.0.0.1", port), Handler)
print(f"Anotátor běží na http://127.0.0.1:{port}")
print(f"Anotace → {OUT_LBL_DIR}")
print(f"Obrázky → {OUT_IMG_DIR}")
print("Ctrl+C pro zastavení")
threading.Timer(0.5, lambda: webbrowser.open(f"http://127.0.0.1:{port}")).start()
try:
server.serve_forever()
except KeyboardInterrupt:
print("\nZastaveno.")