463 lines
16 KiB
Python
463 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Anotátor dlaždic pro trénink YOLOv8
|
||
Formát: YOLO class cx cy w h (normalizovaný 0–1)
|
||
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">
|
||
1–4 vybrat třídu<br>
|
||
S uložit<br>
|
||
Del smazat vybraný<br>
|
||
← → předchozí / další<br>
|
||
Z zpět (undo)<br>
|
||
Esc 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é 0–1
|
||
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.")
|