feat: add counting bench, triage, and dataset modules
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.
This commit is contained in:
asus committed
1 parent
8285400254
commit
5c7c122105
80 files changed
+20074
-1412
No files matched your search
+393
-129
@@ -11,48 +11,115 @@ Label files are plain YOLO:
|
||||
segment class_id x1 y1 x2 y2 … (normalized polygon)
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from backend import batches, config, db, jobs, projects, review
|
||||
from backend import batches, config, datasets, db, jobs, projects, review, triage
|
||||
|
||||
DatasetError = datasets.DatasetError
|
||||
|
||||
|
||||
class DatasetError(Exception):
|
||||
pass
|
||||
def dataset_dir(project_slug: str, dataset_id: int) -> str:
|
||||
return datasets.dataset_root(project_slug, dataset_id)
|
||||
|
||||
|
||||
def dataset_dir(project_slug: str) -> str:
|
||||
return os.path.join(config.project_dir(project_slug), "dataset")
|
||||
def runs_dir(project_slug: str) -> str:
|
||||
"""Where a training run's assembled view lives.
|
||||
|
||||
It sits outside any one dataset because a run may combine several, and the
|
||||
combined view belongs to the run, not to any of its sources.
|
||||
"""
|
||||
return os.path.join(config.project_dir(project_slug), "runs")
|
||||
|
||||
|
||||
def approve(batch_id: int) -> dict:
|
||||
"""Sign a batch off and queue its merge (REQ-045, REQ-050)."""
|
||||
batch = batches.get(batch_id)
|
||||
if batch is None:
|
||||
raise DatasetError("No such batch")
|
||||
if batch["status"] == "merged":
|
||||
raise DatasetError("This batch is already in the master dataset")
|
||||
if batch["review"]["pending"] > 0:
|
||||
raise DatasetError(
|
||||
f"{batch['review']['pending']} frame(s) still need a decision before this "
|
||||
"batch can be approved"
|
||||
def approve(batch_ids, dataset_id: Optional[int] = None,
|
||||
dataset_name: str = "") -> dict:
|
||||
"""Sign a selection off and queue one merge into one named dataset (REQ-131).
|
||||
|
||||
The rules in force right now are frozen onto the target dataset (REQ-132):
|
||||
the merge runs under them, and a later rule edit cannot rewrite what this
|
||||
dataset claims to be.
|
||||
|
||||
Without `dataset_id` a new dataset is created, so merging the same batches
|
||||
again never collides with the earlier result — it produces a second dataset
|
||||
holding them as they look now.
|
||||
"""
|
||||
ids = triage.as_ids(batch_ids)
|
||||
if not ids:
|
||||
raise DatasetError("Pick at least one batch to merge")
|
||||
selected = []
|
||||
for batch_id in ids:
|
||||
batch = batches.get(batch_id)
|
||||
if batch is None:
|
||||
raise DatasetError("No such batch")
|
||||
if batch["review"]["approved"] == 0:
|
||||
raise DatasetError(
|
||||
f"No frame in {batch['date_label']}/{batch['batch_label']} is approved "
|
||||
"— there is nothing to merge")
|
||||
selected.append(batch)
|
||||
if len({batch["project_id"] for batch in selected}) > 1:
|
||||
raise DatasetError("Those batches are not all in the same project")
|
||||
project_id = selected[0]["project_id"]
|
||||
|
||||
# Frames that are not approved — rejected or never looked at — are simply
|
||||
# left behind. Only what the user signed off on enters the dataset, so a
|
||||
# partly-reviewed batch can be merged for the part that is done.
|
||||
with db.cursor() as cur:
|
||||
placeholders = ",".join("?" for _ in ids)
|
||||
cur.execute(
|
||||
f"""SELECT 1 FROM jobs WHERE batch_id IN ({placeholders}) AND type = 'merge'
|
||||
AND status IN ('queued', 'running')""",
|
||||
ids,
|
||||
)
|
||||
if batch["review"]["approved"] == 0:
|
||||
raise DatasetError("Every frame was rejected — there is nothing to merge")
|
||||
if cur.fetchone() is not None:
|
||||
raise DatasetError("A merge for one of these batches is already queued")
|
||||
|
||||
batches.set_status(batch_id, "approved")
|
||||
resolver = triage.Resolver(project_id)
|
||||
if dataset_id is None:
|
||||
target = datasets.create(project_id, name=dataset_name,
|
||||
rule_version=resolver.version(), rules=resolver.rules)
|
||||
dataset_id = target["id"]
|
||||
else:
|
||||
target = datasets.get(dataset_id)
|
||||
if target is None:
|
||||
raise DatasetError("No such dataset")
|
||||
if all(_unmerged_approved(batch["id"], dataset_id) == 0 for batch in selected):
|
||||
raise DatasetError(
|
||||
f"Every approved frame of this selection is already in \u201c{target['name']}\u201d")
|
||||
# An existing dataset keeps the rules it was cut under; a second merge
|
||||
# into it must not re-cut the frames already there under new ones.
|
||||
if not target["rules"]:
|
||||
datasets.snapshot_rules(dataset_id, resolver.rules, resolver.version())
|
||||
|
||||
for batch in selected:
|
||||
batches.set_status(batch["id"], "approved")
|
||||
labels = ", ".join(f"{b['date_label']}/{b['batch_label']}" for b in selected)
|
||||
job = jobs.create(
|
||||
"merge",
|
||||
params={"batch_id": batch_id},
|
||||
project_id=batch["project_id"],
|
||||
batch_id=batch_id,
|
||||
message=f"{batch['date_label']}/{batch['batch_label']}",
|
||||
params={"batch_ids": ids, "dataset_id": dataset_id},
|
||||
project_id=project_id,
|
||||
batch_id=ids[0],
|
||||
message=labels,
|
||||
)
|
||||
return job.to_dict()
|
||||
|
||||
|
||||
def _unmerged_approved(batch_id: int, dataset_id: int) -> int:
|
||||
"""Approved frames of this batch not yet in *this* dataset."""
|
||||
with db.cursor() as cur:
|
||||
cur.execute(
|
||||
"""SELECT COUNT(*) FROM frames f
|
||||
LEFT JOIN dataset_items d
|
||||
ON d.frame_id = f.id AND d.dataset_id = ?
|
||||
WHERE f.batch_id = ? AND f.review_status = 'approved' AND d.id IS NULL""",
|
||||
(dataset_id, batch_id),
|
||||
)
|
||||
return cur.fetchone()[0]
|
||||
|
||||
|
||||
def _label_line(class_id: int, geometry: dict, label_type: str) -> str:
|
||||
if label_type == "bbox":
|
||||
x0, y0, x1, y1 = review.to_box(geometry)
|
||||
@@ -66,106 +133,245 @@ def _label_line(class_id: int, geometry: dict, label_type: str) -> str:
|
||||
return f"{class_id} {coords}"
|
||||
|
||||
|
||||
def _next_split(cur, project_id: int, val_every: int) -> str:
|
||||
"""Continue the every-Nth pattern from wherever the last merge left off."""
|
||||
def split_for(project_id: int, batch_id: int, stem: str, val_every: int) -> str:
|
||||
"""Which split a frame belongs to, derived from its identity rather than from
|
||||
how many rows happen to precede it.
|
||||
|
||||
A positional every-Nth rule makes membership depend on insertion history, so
|
||||
deleting or re-merging a batch silently reshuffles every later frame — and a
|
||||
frame that was in `val` for the last comparison could land in `train` for the
|
||||
next one. Hashing the identity makes the stable-val-split invariant true by
|
||||
construction: the same frame always lands in the same split, whatever else
|
||||
happened to the dataset. Rows already in `dataset_items` keep the split they
|
||||
were recorded with; nothing recomputes them.
|
||||
"""
|
||||
if val_every <= 0:
|
||||
return "train"
|
||||
cur.execute("SELECT COUNT(*) FROM dataset_items WHERE project_id = ?", (project_id,))
|
||||
position = cur.fetchone()[0]
|
||||
return "val" if position % val_every == val_every - 1 else "train"
|
||||
digest = hashlib.sha1(f"{project_id}/{batch_id}/{stem}".encode("utf-8")).hexdigest()
|
||||
return "val" if int(digest[:8], 16) % val_every == 0 else "train"
|
||||
|
||||
|
||||
def sync_labels(project_id: int, selected_class_ids: Optional[List[int]] = None) -> dict:
|
||||
"""Re-sync label files on disk for all merged frames in the project dataset."""
|
||||
project = projects.get(project_id)
|
||||
root = dataset_dir(project["slug"])
|
||||
def resync(dataset_id: int) -> dict:
|
||||
"""Rewrite one dataset's labels from the current annotations and rules.
|
||||
|
||||
This used to run automatically before every training run, which quietly
|
||||
undid the triage applied at merge: a dataset merged under "reclass small
|
||||
boxes" had its labels rebuilt from the raw annotations on the next run, so
|
||||
the files stopped matching the `rule_version` stamped on them. It is now a
|
||||
deliberate act, and it re-stamps that version so the dataset never claims a
|
||||
rule set it is not in.
|
||||
|
||||
Only frames a human signed off on are written. A merged frame whose batch was
|
||||
auto-annotated again drops back to `pending`, and rewriting its label from
|
||||
fresh model output would push predictions nobody checked into the dataset.
|
||||
"""
|
||||
target = datasets.get(dataset_id)
|
||||
if target is None:
|
||||
raise DatasetError("No such dataset")
|
||||
project = projects.get(target["project_id"])
|
||||
resolver = triage.Resolver(project["id"])
|
||||
root = dataset_dir(project["slug"], dataset_id)
|
||||
|
||||
with db.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT d.frame_id, d.label_rel FROM dataset_items d WHERE d.project_id = ?",
|
||||
(project_id,),
|
||||
"""SELECT d.frame_id, d.label_rel FROM dataset_items d
|
||||
JOIN frames f ON f.id = d.frame_id
|
||||
WHERE d.dataset_id = ? AND f.review_status = 'approved'""",
|
||||
(dataset_id,),
|
||||
)
|
||||
items = cur.fetchall()
|
||||
|
||||
class_map = None
|
||||
if selected_class_ids is not None and len(selected_class_ids) > 0:
|
||||
class_map = {cid: idx for idx, cid in enumerate(sorted(selected_class_ids))}
|
||||
|
||||
synced_files = 0
|
||||
total_lines = 0
|
||||
written = 0
|
||||
emptied = 0
|
||||
for frame_id, label_rel in items:
|
||||
annotations = review.listing(frame_id)
|
||||
if class_map is not None:
|
||||
annotations = [a for a in annotations if a["class_id"] in class_map]
|
||||
|
||||
lines = []
|
||||
for item in annotations:
|
||||
mapped_cid = class_map[item["class_id"]] if class_map is not None else item["class_id"]
|
||||
lines.append(_label_line(mapped_cid, item["geometry"], project["label_type"]))
|
||||
|
||||
resolved = resolver.resolve_shapes(annotations)
|
||||
if resolved is None:
|
||||
# Every shape was dropped. The image stays in the dataset but an
|
||||
# empty label would claim it is empty, so the file is left as it was
|
||||
# and the count is reported (REQ-104).
|
||||
emptied += 1
|
||||
continue
|
||||
lines = [_label_line(item["class_id"], item["geometry"], project["label_type"])
|
||||
for item in resolved]
|
||||
path = os.path.join(root, label_rel)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write("\n".join(lines) + ("\n" if lines else ""))
|
||||
synced_files += 1
|
||||
total_lines += len(lines)
|
||||
_write_atomic(path, "\n".join(lines) + ("\n" if lines else ""))
|
||||
written += 1
|
||||
|
||||
return {"synced_files": synced_files, "total_lines": total_lines}
|
||||
# Resync is the one deliberate way an existing dataset adopts today's rules,
|
||||
# so the snapshot moves with the labels (REQ-132).
|
||||
datasets.snapshot_rules(dataset_id, resolver.rules, resolver.version())
|
||||
|
||||
return {"labels_written": written, "frames_left_alone": emptied,
|
||||
"rule_version": resolver.version()}
|
||||
|
||||
|
||||
def write_data_yaml(project: dict, batch_ids: list = None, selected_class_ids: Optional[List[int]] = None) -> str:
|
||||
"""Rebuild data.yaml from the project's classes (REQ-051)."""
|
||||
sync_labels(project["id"], selected_class_ids=selected_class_ids)
|
||||
root = dataset_dir(project["slug"])
|
||||
os.makedirs(root, exist_ok=True)
|
||||
counts = summary(project["id"])["splits"]
|
||||
def _write_atomic(path: str, text: str) -> None:
|
||||
"""Write via temp file + rename, so a training run never reads a half-written
|
||||
label file or a truncated data.yaml."""
|
||||
tmp = f"{path}.tmp"
|
||||
with open(tmp, "w", encoding="utf-8") as handle:
|
||||
handle.write(text)
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def _build_selected_tree(run_root: str, rows: list, class_map: Optional[dict]) -> tuple:
|
||||
"""Materialise the run's view of the chosen datasets under `runs/selected/`.
|
||||
|
||||
Labels are copied from what each dataset holds on disk — not re-derived from
|
||||
the live annotations. A dataset is the snapshot of a batch as it was merged,
|
||||
under the triage rules recorded in its `rule_version`; re-resolving here
|
||||
would train on today's rules while the dataset claims yesterday's, and two
|
||||
runs over the same dataset could then disagree. Change the rules and merge
|
||||
again into a new dataset, or resync this one on purpose.
|
||||
|
||||
The only thing this does apply is a per-run class filter, which renumbers ids
|
||||
into a contiguous 0..k-1 space. That contradicts `project_classes`, so it
|
||||
cannot be written back into the dataset's own label files.
|
||||
"""
|
||||
selected_root = os.path.join(run_root, "selected")
|
||||
if os.path.isdir(selected_root):
|
||||
shutil.rmtree(selected_root)
|
||||
listed = {"train": [], "val": []}
|
||||
excluded = 0
|
||||
for row in rows:
|
||||
split, source_image, source_label = row["split"], row["source_image"], row["source_label"]
|
||||
lines = _read_label(source_label)
|
||||
|
||||
if class_map is not None:
|
||||
kept = []
|
||||
for line in lines:
|
||||
head, _, rest = line.partition(" ")
|
||||
try:
|
||||
current = int(head)
|
||||
except ValueError:
|
||||
continue
|
||||
if current in class_map:
|
||||
kept.append(f"{class_map[current]} {rest}")
|
||||
# A frame that had shapes but none of the chosen classes is not a
|
||||
# negative sample of those classes — it is a frame full of things the
|
||||
# run was told to ignore, and an empty label would teach exactly that.
|
||||
if lines and not kept:
|
||||
excluded += 1
|
||||
continue
|
||||
lines = kept
|
||||
|
||||
stem = os.path.basename(source_image)
|
||||
image_dst = os.path.join(selected_root, "images", split, stem)
|
||||
label_dst = os.path.join(selected_root, "labels", split,
|
||||
os.path.splitext(stem)[0] + ".txt")
|
||||
os.makedirs(os.path.dirname(image_dst), exist_ok=True)
|
||||
os.makedirs(os.path.dirname(label_dst), exist_ok=True)
|
||||
if not os.path.exists(image_dst):
|
||||
os.symlink(source_image, image_dst)
|
||||
_write_atomic(label_dst, "\n".join(lines) + ("\n" if lines else ""))
|
||||
listed[split].append(image_dst)
|
||||
listed["excluded"] = excluded
|
||||
return selected_root, listed
|
||||
|
||||
|
||||
def _read_label(path: str) -> List[str]:
|
||||
if not os.path.isfile(path):
|
||||
return []
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
return [line for line in handle.read().splitlines() if line.strip()]
|
||||
|
||||
|
||||
def write_data_yaml(project: dict, dataset_ids: List[int], batch_ids: list = None,
|
||||
selected_class_ids: Optional[List[int]] = None,
|
||||
require_val: bool = False,
|
||||
base_dataset_ids: Optional[List[int]] = None) -> str:
|
||||
"""Assemble the chosen datasets into one data.yaml for a run (REQ-051, REQ-110).
|
||||
|
||||
Always via the `selected/` tree of symlinks, even for a single dataset with
|
||||
no filters. The alternative — pointing YOLO at a dataset folder directly —
|
||||
only works while a run uses exactly one dataset, and it puts a per-run class
|
||||
renumbering into the shared label files. One assembly path is easier to
|
||||
trust than two that diverge the moment a second dataset is picked.
|
||||
"""
|
||||
if not dataset_ids and not base_dataset_ids:
|
||||
raise DatasetError("Pick at least one dataset to train on")
|
||||
run_root = runs_dir(project["slug"])
|
||||
os.makedirs(run_root, exist_ok=True)
|
||||
|
||||
target_classes = project["classes"]
|
||||
class_map = None
|
||||
if selected_class_ids is not None and len(selected_class_ids) > 0:
|
||||
target_classes = [c for c in project["classes"] if c["class_id"] in selected_class_ids]
|
||||
class_map = {cid: idx for idx, cid in enumerate(sorted(selected_class_ids))}
|
||||
|
||||
names = ", ".join(f"'{item['name']}'" for item in target_classes)
|
||||
|
||||
items = datasets.combined_items(project["id"], dataset_ids)
|
||||
if batch_ids:
|
||||
with db.cursor() as cur:
|
||||
placeholders = ",".join("?" for _ in batch_ids)
|
||||
cur.execute(
|
||||
f"""SELECT d.image_rel, d.split FROM dataset_items d
|
||||
JOIN frames f ON f.id = d.frame_id
|
||||
WHERE d.project_id = ? AND f.batch_id IN ({placeholders})""",
|
||||
[project["id"]] + list(batch_ids),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
keep = _frames_of_batches(set(batch_ids))
|
||||
items = [item for item in items if item["frame_id"] in keep]
|
||||
|
||||
train_files = [row[0] for row in rows if row[1] == "train"]
|
||||
val_files = [row[0] for row in rows if row[1] == "val"] or train_files
|
||||
rows = []
|
||||
for item in items:
|
||||
root = dataset_dir(project["slug"], item["dataset_id"])
|
||||
rows.append({
|
||||
"frame_id": item["frame_id"],
|
||||
"split": item["split"],
|
||||
"source_image": os.path.join(root, item["image_rel"]),
|
||||
"source_label": os.path.join(root, item["label_rel"]),
|
||||
})
|
||||
|
||||
train_txt = os.path.join(root, "selected_train.txt")
|
||||
val_txt = os.path.join(root, "selected_val.txt")
|
||||
with open(train_txt, "w", encoding="utf-8") as handle:
|
||||
handle.write("\n".join(os.path.join(root, rel) for rel in train_files) + "\n")
|
||||
with open(val_txt, "w", encoding="utf-8") as handle:
|
||||
handle.write("\n".join(os.path.join(root, rel) for rel in val_files) + "\n")
|
||||
# Base datasets are appended, never merged into the dedupe above: they carry
|
||||
# no frame_id, and they are always train-only (REQ-122).
|
||||
if base_dataset_ids:
|
||||
from backend import base_dataset
|
||||
rows.extend(base_dataset.rows(project["id"], base_dataset_ids, project["slug"]))
|
||||
|
||||
path = os.path.join(root, "selected_data.yaml")
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(f"path: {root}\n")
|
||||
handle.write(f"train: {train_txt}\n")
|
||||
handle.write(f"val: {val_txt}\n\n")
|
||||
handle.write(f"nc: {len(project['classes'])}\n")
|
||||
handle.write(f"names: [{names}]\n")
|
||||
return path
|
||||
selected_root, listed = _build_selected_tree(run_root, rows, class_map)
|
||||
if require_val:
|
||||
_require_val(len(listed["val"]), "the selected dataset(s)")
|
||||
|
||||
path = os.path.join(root, "data.yaml")
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(f"path: {root}\n")
|
||||
handle.write("train: images/train\n")
|
||||
handle.write(f"val: images/{'val' if counts['val'] > 0 else 'train'}\n\n")
|
||||
handle.write(f"nc: {len(project['classes'])}\n")
|
||||
handle.write(f"names: [{names}]\n")
|
||||
train_txt = os.path.join(run_root, "selected_train.txt")
|
||||
val_txt = os.path.join(run_root, "selected_val.txt")
|
||||
_write_atomic(train_txt, "\n".join(listed["train"]) + "\n")
|
||||
_write_atomic(val_txt, "\n".join(listed["val"]) + "\n")
|
||||
|
||||
path = os.path.join(run_root, "selected_data.yaml")
|
||||
_write_atomic(path,
|
||||
f"path: {selected_root}\n"
|
||||
f"train: {train_txt}\n"
|
||||
f"val: {val_txt}\n\n"
|
||||
f"nc: {len(target_classes)}\n"
|
||||
f"names: [{names}]\n")
|
||||
return path
|
||||
|
||||
|
||||
def _frames_of_batches(batch_ids: set) -> set:
|
||||
with db.cursor() as cur:
|
||||
cur.execute(
|
||||
f"SELECT id FROM frames WHERE batch_id IN ({','.join('?' for _ in batch_ids)})",
|
||||
list(batch_ids),
|
||||
)
|
||||
return {row[0] for row in cur.fetchall()}
|
||||
|
||||
|
||||
def _require_val(count: int, subject: str) -> None:
|
||||
"""Refuse to build a dataset with an empty val split.
|
||||
|
||||
Falling back to the training images produces a base-vs-new mAP measured on
|
||||
data the model was fitted to — a number that looks fine and means nothing.
|
||||
For a system whose whole purpose is answering "did retraining help?", this
|
||||
has to fail loudly.
|
||||
"""
|
||||
if count == 0:
|
||||
raise DatasetError(
|
||||
f"There are no validation images in {subject}, so a base-vs-new comparison "
|
||||
"would be measured on the training images. Merge more frames, or lower the "
|
||||
"project's val_every."
|
||||
)
|
||||
|
||||
|
||||
def summary(project_id: int) -> dict:
|
||||
"""Counts only. The per-shape size analytics that used to live here walked
|
||||
every annotation in the project on every page load — Data Prep already
|
||||
serves that, per batch, from `triage`."""
|
||||
with db.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT split, COUNT(*) FROM dataset_items WHERE project_id = ? GROUP BY split",
|
||||
@@ -185,7 +391,12 @@ def summary(project_id: int) -> dict:
|
||||
(project_id,),
|
||||
)
|
||||
merged = [dict(row) for row in cur.fetchall()]
|
||||
return {"splits": splits, "total": splits["train"] + splits["val"], "batches": merged}
|
||||
|
||||
return {
|
||||
"splits": splits,
|
||||
"total": splits["train"] + splits["val"],
|
||||
"batches": merged,
|
||||
}
|
||||
|
||||
|
||||
def drop_class_from_labels(project: dict, class_id: int) -> dict:
|
||||
@@ -196,16 +407,15 @@ def drop_class_from_labels(project: dict, class_id: int) -> dict:
|
||||
class that is now `1` — labels that quietly name the wrong thing are worse
|
||||
than labels that are missing.
|
||||
"""
|
||||
root = dataset_dir(project["slug"])
|
||||
with db.cursor() as cur:
|
||||
cur.execute("SELECT label_rel FROM dataset_items WHERE project_id = ?",
|
||||
cur.execute("SELECT label_rel, dataset_id FROM dataset_items WHERE project_id = ?",
|
||||
(project["id"],))
|
||||
label_files = [row[0] for row in cur.fetchall()]
|
||||
label_files = [(row[0], row[1]) for row in cur.fetchall()]
|
||||
|
||||
rewritten = 0
|
||||
dropped = 0
|
||||
for rel in label_files:
|
||||
path = os.path.join(root, rel)
|
||||
for rel, dataset_id in label_files:
|
||||
path = os.path.join(dataset_dir(project["slug"], dataset_id), rel)
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
@@ -240,74 +450,128 @@ def drop_class_from_labels(project: dict, class_id: int) -> dict:
|
||||
return {"label_files_rewritten": rewritten, "dataset_lines_removed": dropped}
|
||||
|
||||
|
||||
def zip_path(project: dict) -> str:
|
||||
"""Zip the master dataset for download (REQ-054)."""
|
||||
root = dataset_dir(project["slug"])
|
||||
def zip_path(project: dict, dataset_id: int) -> str:
|
||||
"""Zip one dataset for download (REQ-054)."""
|
||||
root = dataset_dir(project["slug"], dataset_id)
|
||||
if not os.path.isdir(os.path.join(root, "images")):
|
||||
raise DatasetError("This project's dataset is still empty")
|
||||
archive = os.path.join(config.project_dir(project["slug"]), "dataset")
|
||||
raise DatasetError("This dataset is still empty")
|
||||
archive = os.path.join(config.project_dir(project["slug"]), f"dataset-{dataset_id}")
|
||||
return shutil.make_archive(archive, "zip", root)
|
||||
|
||||
|
||||
@jobs.handler("merge")
|
||||
def _run_merge(job) -> None:
|
||||
batch = batches.get(job.params["batch_id"])
|
||||
if batch is None:
|
||||
raise DatasetError("The batch disappeared before the merge started")
|
||||
project = projects.get(batch["project_id"])
|
||||
root = dataset_dir(project["slug"])
|
||||
ids = job.params.get("batch_ids") or [job.params["batch_id"]]
|
||||
selected = [batches.get(bid) for bid in ids]
|
||||
if any(batch is None for batch in selected):
|
||||
raise DatasetError("A batch disappeared before the merge started")
|
||||
project = projects.get(selected[0]["project_id"])
|
||||
dataset_id = job.params["dataset_id"]
|
||||
target = datasets.get(dataset_id)
|
||||
if target is None:
|
||||
raise DatasetError("The target dataset disappeared before the merge started")
|
||||
root = dataset_dir(project["slug"], dataset_id)
|
||||
for split in ("train", "val"):
|
||||
os.makedirs(os.path.join(root, "images", split), exist_ok=True)
|
||||
os.makedirs(os.path.join(root, "labels", split), exist_ok=True)
|
||||
|
||||
frames = [f for f in batches.frames(batch["id"]) if f["review_status"] == "approved"]
|
||||
source_dir = batches.frames_dir(project["slug"], batch["id"])
|
||||
job.progress(0, len(frames))
|
||||
job.log(f"Merging {len(frames)} approved frame(s) into the master dataset")
|
||||
work = []
|
||||
for batch in selected:
|
||||
frames = [f for f in batches.frames(batch["id"]) if f["review_status"] == "approved"]
|
||||
work.extend((batch, frame) for frame in frames)
|
||||
job.progress(0, len(work))
|
||||
job.log(f"Merging {len(work)} approved frame(s) from {len(selected)} batch(es) "
|
||||
f"into \u201c{target['name']}\u201d")
|
||||
|
||||
# Triage gates the merge (REQ-104), under the rules frozen onto this dataset
|
||||
# when it was created (REQ-132) — not under whatever the project says now.
|
||||
resolver = triage.Resolver(project["id"], frozen=target["rules"])
|
||||
gating = bool(resolver.rules or resolver.overrides)
|
||||
if gating:
|
||||
job.log(f"Applying {len(resolver.rules)} triage rule(s), version {resolver.version()}")
|
||||
|
||||
added = {"train": 0, "val": 0}
|
||||
skipped = 0
|
||||
for index, frame in enumerate(frames):
|
||||
triaged_out = 0
|
||||
cancelled = False
|
||||
for index, (batch, frame) in enumerate(work):
|
||||
if job.cancelled:
|
||||
job.log(f"Cancelled after {index} frame(s)")
|
||||
cancelled = True
|
||||
break
|
||||
|
||||
annotations = review.listing(frame["id"])
|
||||
if gating:
|
||||
resolved = resolver.resolve_shapes(annotations)
|
||||
if resolved is None:
|
||||
triaged_out += 1
|
||||
job.progress(index + 1, len(work))
|
||||
continue
|
||||
annotations = resolved
|
||||
|
||||
with db.cursor() as cur:
|
||||
cur.execute("SELECT 1 FROM dataset_items WHERE frame_id = ?", (frame["id"],))
|
||||
cur.execute("SELECT 1 FROM dataset_items WHERE dataset_id = ? AND frame_id = ?",
|
||||
(dataset_id, frame["id"]))
|
||||
if cur.fetchone() is not None:
|
||||
skipped += 1
|
||||
job.progress(index + 1, len(frames))
|
||||
job.progress(index + 1, len(work))
|
||||
continue
|
||||
split = _next_split(cur, project["id"], project["val_every"])
|
||||
|
||||
stem = f"{batch['id']}__{os.path.splitext(frame['filename'])[0]}"
|
||||
# A frame's split is decided once for the whole project and every
|
||||
# later dataset inherits it. Letting each dataset re-decide would put
|
||||
# the same image in `val` for one run and `train` for the next, so a
|
||||
# base-vs-new mAP would be measured on images the new model had been
|
||||
# fitted to. The hash agrees with itself, but rows merged before the
|
||||
# hash existed carry a positional split — those have to be honoured,
|
||||
# not recomputed.
|
||||
cur.execute(
|
||||
"SELECT split FROM dataset_items WHERE frame_id = ? LIMIT 1",
|
||||
(frame["id"],),
|
||||
)
|
||||
previous = cur.fetchone()
|
||||
split = previous[0] if previous else split_for(
|
||||
project["id"], batch["id"], stem, project["val_every"])
|
||||
image_rel = f"images/{split}/{stem}.jpg"
|
||||
label_rel = f"labels/{split}/{stem}.txt"
|
||||
shutil.copyfile(os.path.join(source_dir, frame["filename"]),
|
||||
os.path.join(root, image_rel))
|
||||
shutil.copyfile(
|
||||
os.path.join(batches.frames_dir(project["slug"], batch["id"]), frame["filename"]),
|
||||
os.path.join(root, image_rel))
|
||||
|
||||
lines = [_label_line(item["class_id"], item["geometry"], project["label_type"])
|
||||
for item in review.listing(frame["id"])]
|
||||
for item in annotations]
|
||||
# An approved frame with nothing on it is a negative sample, and an
|
||||
# empty .txt is how YOLO spells that (REQ-033).
|
||||
with open(os.path.join(root, label_rel), "w", encoding="utf-8") as handle:
|
||||
handle.write("\n".join(lines) + ("\n" if lines else ""))
|
||||
|
||||
cur.execute(
|
||||
"""INSERT INTO dataset_items (project_id, frame_id, split, image_rel,
|
||||
label_rel, added_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)""",
|
||||
(project["id"], frame["id"], split, image_rel, label_rel, time.time()),
|
||||
"""INSERT INTO dataset_items (project_id, dataset_id, frame_id, split,
|
||||
image_rel, label_rel, added_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)""",
|
||||
(project["id"], dataset_id, frame["id"], split, image_rel, label_rel,
|
||||
time.time()),
|
||||
)
|
||||
added[split] += 1
|
||||
job.progress(index + 1, len(frames))
|
||||
job.progress(index + 1, len(work))
|
||||
|
||||
if cancelled:
|
||||
# Leaving them 'merged' would be a lie: the frames after the break point
|
||||
# have no dataset_items rows and no files, and approve() refuses to
|
||||
# re-merge a merged batch, so they could never be added. The per-frame
|
||||
# dataset_items guard already makes re-running the merge idempotent.
|
||||
job.log("Batches left approved — re-approve them to finish the merge")
|
||||
return
|
||||
|
||||
with db.cursor() as cur:
|
||||
cur.execute("UPDATE batches SET status = 'merged', merged_at = ? WHERE id = ?",
|
||||
(time.time(), batch["id"]))
|
||||
cur.executemany(
|
||||
"UPDATE batches SET status = 'merged', merged_at = ? WHERE id = ?",
|
||||
[(time.time(), bid) for bid in ids],
|
||||
)
|
||||
|
||||
path = write_data_yaml(projects.get(project["id"]))
|
||||
totals = summary(project["id"])["splits"]
|
||||
totals = datasets.get(dataset_id)["splits"]
|
||||
job.log(f"Added {added['train']} train / {added['val']} val"
|
||||
+ (f", skipped {skipped} already merged" if skipped else ""))
|
||||
job.log(f"Master dataset now {totals['train']} train / {totals['val']} val — {path}")
|
||||
+ (f", skipped {skipped} already in this dataset" if skipped else "")
|
||||
+ (f", held back {triaged_out} by triage" if triaged_out else ""))
|
||||
job.log(f"\u201c{target['name']}\u201d now holds "
|
||||
f"{totals['train']} train / {totals['val']} val")
|
||||
Reference in new issue
Block a user