asus
5c7c122105
This commit includes major additions and updates to the frontend and backend architectures, introducing new dataset management, live counting features, batch processing, and triage logic. Includes new UI pages, components, and API routes.
304 lines
12 KiB
Python
304 lines
12 KiB
Python
"""Export annotated frames as a plain YOLO dataset the user can keep.
|
|
|
|
This is deliberately independent of the master dataset. A batch can be exported
|
|
the moment it has annotations — before review, before merge — because the point
|
|
is that the user owns a copy of their own work and does not have to trust this
|
|
application to keep it.
|
|
|
|
Layout inside the zip:
|
|
|
|
images/<batch>__<frame>.jpg
|
|
labels/<batch>__<frame>.txt
|
|
data.yaml
|
|
classes.txt
|
|
MANIFEST.json what was exported, and what was skipped
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
from typing import List, Optional
|
|
|
|
from backend import batches, config, db, projects, review
|
|
|
|
|
|
class ExportError(Exception):
|
|
pass
|
|
|
|
|
|
def _label_line(class_id: int, geometry: dict, label_type: str) -> str:
|
|
from backend.dataset import _label_line as build
|
|
return build(class_id, geometry, label_type)
|
|
|
|
|
|
def annotated_frames(project_id: int, batch_ids: Optional[List[int]] = None,
|
|
approved_only: bool = False) -> list:
|
|
where = "b.project_id = ?"
|
|
args: list = [project_id]
|
|
if batch_ids:
|
|
where += f" AND b.id IN ({','.join('?' for _ in batch_ids)})"
|
|
args += list(batch_ids)
|
|
if approved_only:
|
|
where += " AND f.review_status = 'approved'"
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
f"""SELECT f.id, f.filename, f.review_status, b.id AS batch_id,
|
|
b.batch_label, b.date_label
|
|
FROM frames f JOIN batches b ON b.id = f.batch_id
|
|
WHERE {where}
|
|
ORDER BY b.id, f.idx""",
|
|
args,
|
|
)
|
|
return [dict(row) for row in cur.fetchall()]
|
|
|
|
|
|
def build_zip(project_id: int, batch_ids: Optional[List[int]] = None,
|
|
approved_only: bool = False, include_empty: bool = False) -> str:
|
|
"""Write the zip and return its path.
|
|
|
|
`include_empty` decides whether frames with no annotation are exported as
|
|
negative samples. Off by default: an empty label file is a claim that the
|
|
image contains nothing, and for a frame that was simply never annotated that
|
|
claim is false.
|
|
"""
|
|
project = projects.get(project_id)
|
|
if project is None:
|
|
raise ExportError("No such project")
|
|
|
|
frames = annotated_frames(project_id, batch_ids, approved_only)
|
|
if not frames:
|
|
raise ExportError("There is nothing to export — no frames match that selection")
|
|
|
|
staging = tempfile.mkdtemp(prefix="export-")
|
|
images_dir = os.path.join(staging, "images")
|
|
labels_dir = os.path.join(staging, "labels")
|
|
os.makedirs(images_dir)
|
|
os.makedirs(labels_dir)
|
|
|
|
written = 0
|
|
skipped_empty = 0
|
|
missing = 0
|
|
shapes = 0
|
|
per_batch: dict = {}
|
|
|
|
for frame in frames:
|
|
annotations = review.listing(frame["id"])
|
|
if not annotations and not include_empty:
|
|
skipped_empty += 1
|
|
continue
|
|
|
|
source = os.path.join(
|
|
batches.frames_dir(project["slug"], frame["batch_id"]), frame["filename"])
|
|
if not os.path.isfile(source):
|
|
missing += 1
|
|
continue
|
|
|
|
stem = f"{frame['batch_id']}__{os.path.splitext(frame['filename'])[0]}"
|
|
shutil.copyfile(source, os.path.join(images_dir, f"{stem}.jpg"))
|
|
lines = [_label_line(item["class_id"], item["geometry"], project["label_type"])
|
|
for item in annotations]
|
|
with open(os.path.join(labels_dir, f"{stem}.txt"), "w", encoding="utf-8") as handle:
|
|
handle.write("\n".join(lines) + ("\n" if lines else ""))
|
|
|
|
written += 1
|
|
shapes += len(lines)
|
|
label = f"{frame['date_label']}/{frame['batch_label']}"
|
|
per_batch[label] = per_batch.get(label, 0) + 1
|
|
|
|
if written == 0:
|
|
shutil.rmtree(staging, ignore_errors=True)
|
|
raise ExportError(
|
|
"Every matching frame was empty or missing its image, so the export would "
|
|
"contain no labels"
|
|
)
|
|
|
|
names = ", ".join(f"'{item['name']}'" for item in project["classes"])
|
|
with open(os.path.join(staging, "data.yaml"), "w", encoding="utf-8") as handle:
|
|
handle.write("path: .\ntrain: images\nval: images\n\n")
|
|
handle.write(f"nc: {len(project['classes'])}\n")
|
|
handle.write(f"names: [{names}]\n")
|
|
|
|
with open(os.path.join(staging, "classes.txt"), "w", encoding="utf-8") as handle:
|
|
handle.write("\n".join(item["name"] for item in project["classes"]) + "\n")
|
|
|
|
manifest = {
|
|
"project": project["name"],
|
|
"label_type": project["label_type"],
|
|
"exported_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
"approved_only": approved_only,
|
|
"include_empty": include_empty,
|
|
"images": written,
|
|
"shapes": shapes,
|
|
"per_batch": per_batch,
|
|
"skipped_without_annotations": skipped_empty,
|
|
"skipped_image_missing": missing,
|
|
"classes": [{"class_id": c["class_id"], "name": c["name"]} for c in project["classes"]],
|
|
}
|
|
with open(os.path.join(staging, "MANIFEST.json"), "w", encoding="utf-8") as handle:
|
|
json.dump(manifest, handle, indent=2)
|
|
|
|
project_root = config.project_dir(project["slug"])
|
|
# Each export is ~200 MB. Without this, every click leaves another copy
|
|
# behind and the project folder grows without bound.
|
|
for stale in os.listdir(project_root):
|
|
if stale.startswith("export-") and stale.endswith(".zip"):
|
|
try:
|
|
os.unlink(os.path.join(project_root, stale))
|
|
except OSError:
|
|
pass
|
|
|
|
stamp = time.strftime("%Y%m%d-%H%M%S")
|
|
path = shutil.make_archive(os.path.join(project_root, f"export-{stamp}"), "zip", staging)
|
|
shutil.rmtree(staging, ignore_errors=True)
|
|
return path
|
|
|
|
|
|
# ---- restore ------------------------------------------------------------
|
|
|
|
def _safe_members(archive) -> list:
|
|
"""Only plain files under images/ and labels/.
|
|
|
|
A zip can name `../../etc/passwd`; extracting one blindly writes wherever it
|
|
says. Every member is checked rather than trusted.
|
|
"""
|
|
picked = []
|
|
for member in archive.infolist():
|
|
if member.is_dir():
|
|
continue
|
|
name = member.filename.replace("\\", "/")
|
|
if name.startswith("/") or ".." in name.split("/"):
|
|
continue
|
|
if name.startswith(("images/", "labels/")) or name == "MANIFEST.json":
|
|
picked.append((name, member))
|
|
return picked
|
|
|
|
|
|
def _points_from_label(parts: List[str], label_type: str) -> Optional[dict]:
|
|
values = [float(v) for v in parts]
|
|
if label_type == "bbox":
|
|
if len(values) != 4:
|
|
return None
|
|
cx, cy, w, h = values
|
|
return {"type": "bbox",
|
|
"points": [cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2]}
|
|
if len(values) < 6 or len(values) % 2:
|
|
return None
|
|
return {"type": "polygon",
|
|
"points": [[values[i], values[i + 1]] for i in range(0, len(values), 2)]}
|
|
|
|
|
|
def restore_zip(project_id: int, zip_path: str, batch_label: str = "") -> dict:
|
|
"""Load an exported zip back in as a fresh batch, ready to keep reviewing.
|
|
|
|
The frames land in a new batch rather than being merged back into the ones
|
|
they came from: the originals may still exist, and silently overwriting a
|
|
batch the user is working in would destroy the very work this feature is
|
|
meant to protect.
|
|
"""
|
|
import zipfile
|
|
|
|
from PIL import Image
|
|
|
|
project = projects.get(project_id)
|
|
if project is None:
|
|
raise ExportError("No such project")
|
|
|
|
by_name = {item["name"]: item["class_id"] for item in project["classes"]}
|
|
stamp = time.strftime("%Y%m%d-%H%M%S")
|
|
label = batch_label or f"restored-{stamp}"
|
|
|
|
with zipfile.ZipFile(zip_path) as archive:
|
|
members = _safe_members(archive)
|
|
names = {name for name, _ in members}
|
|
if not any(name.startswith("images/") for name in names):
|
|
raise ExportError("This zip has no images/ folder — is it an export from this app?")
|
|
|
|
manifest = {}
|
|
if "MANIFEST.json" in names:
|
|
manifest = json.loads(archive.read("MANIFEST.json"))
|
|
source_type = manifest.get("label_type", project["label_type"])
|
|
if source_type != project["label_type"]:
|
|
raise ExportError(
|
|
f"This export holds {source_type} labels but the project is "
|
|
f"{project['label_type']} — importing it would produce wrong shapes"
|
|
)
|
|
|
|
# Classes come back by name, so an id that shifted since the export does
|
|
# not silently relabel every shape.
|
|
remap = {}
|
|
for item in manifest.get("classes", []):
|
|
if item["name"] in by_name:
|
|
remap[item["class_id"]] = by_name[item["name"]]
|
|
else:
|
|
raise ExportError(
|
|
f"The export uses class '{item['name']}', which this project does not "
|
|
"have. Add the class first, then import."
|
|
)
|
|
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
"""INSERT INTO batches (project_id, video_path, date_label, batch_label,
|
|
start_sec, end_sec, fps, status, created_at)
|
|
VALUES (?, '', 'restored', ?, 0, 0, 0, 'extracted', ?)""",
|
|
(project_id, label, time.time()),
|
|
)
|
|
batch_id = cur.lastrowid
|
|
|
|
target_dir = batches.frames_dir(project["slug"], batch_id)
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
restored, shapes, skipped = 0, 0, 0
|
|
image_members = sorted(n for n in names if n.startswith("images/"))
|
|
for index, name in enumerate(image_members):
|
|
stem = os.path.splitext(os.path.basename(name))[0]
|
|
if not stem:
|
|
continue
|
|
filename = f"{stem}.jpg"
|
|
destination = os.path.join(target_dir, filename)
|
|
with archive.open(name) as source, open(destination, "wb") as handle:
|
|
shutil.copyfileobj(source, handle)
|
|
|
|
try:
|
|
with Image.open(destination) as image:
|
|
width, height = image.size
|
|
except Exception:
|
|
os.unlink(destination)
|
|
skipped += 1
|
|
continue
|
|
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
"""INSERT INTO frames (batch_id, idx, filename, width, height, review_status)
|
|
VALUES (?, ?, ?, ?, ?, 'pending')""",
|
|
(batch_id, index, filename, width, height),
|
|
)
|
|
frame_id = cur.lastrowid
|
|
|
|
label_name = f"labels/{stem}.txt"
|
|
if label_name in names:
|
|
for line in archive.read(label_name).decode("utf-8").splitlines():
|
|
parts = line.split()
|
|
if len(parts) < 5:
|
|
continue
|
|
geometry = _points_from_label(parts[1:], project["label_type"])
|
|
if geometry is None:
|
|
continue
|
|
class_id = remap.get(int(parts[0]), int(parts[0]))
|
|
cur.execute(
|
|
"""INSERT INTO annotations (frame_id, class_id, geometry, score,
|
|
source, created_at)
|
|
VALUES (?, ?, ?, 1.0, 'manual', ?)""",
|
|
(frame_id, class_id, json.dumps(geometry), time.time()),
|
|
)
|
|
shapes += 1
|
|
restored += 1
|
|
|
|
with db.cursor() as cur:
|
|
cur.execute("UPDATE batches SET frame_count = ?, status = 'reviewing' WHERE id = ?",
|
|
(restored, batch_id))
|
|
|
|
return {"batch_id": batch_id, "batch_label": label, "frames": restored,
|
|
"shapes": shapes, "skipped": skipped}
|