diff --git a/README.md b/README.md index 53ecb9b..144b0cc 100644 --- a/README.md +++ b/README.md @@ -222,8 +222,32 @@ Notes: - Raw ids are remapped to a contiguous `0..n-1` per head (so `scale`'s `1`/`2`/`3` become `0`/`1`/`2`); `categories` may be omitted, in which case the id set is inferred. -- A missing per-annotation value defaults to `0`; a JSON with **no** `attributes` ⇒ - detection-only, exactly as before. +- A **missing or out-of-vocab** per-annotation value is *ignored* (`-100`), not trained as + class `0` — so a **partially tagged** dataset is valid: each head learns only from the + instances that actually carry its value. A JSON with **no** `attributes` ⇒ detection-only, + exactly as before. + +**Class-conditional heads.** Give an attribute definition an optional `applies_to` list of +primary-class names or ids, and that head is only trained on — and only emitted for — +instances of those classes (hard routing on the primary class). Omit it and the head applies +to every class. So different attributes can attach to different classes, each with its own +label set, in one model: + +```jsonc +"attributes": [ + {"name": "posture", "categories": [...], "applies_to": ["cat", "dog"]} +] +``` + +At inference a scoped head reports `ids = -1` ("not applicable") for detections whose class it +does not cover. The scope is stored in the checkpoint, so it survives reload. + +**Sidecar format (optional).** You can keep the COCO JSON as plain, standard COCO and put the +attributes beside it instead of inside it: an `attributes.yaml` schema in the dataset root plus +`attributes/.json` values keyed by annotation id (`{ann_id: {head: value}}`). If present +(and the JSON has no embedded `attributes`), it is merged in memory at load — so a plain COCO +dataset always works and the sidecar is picked up automatically when you add it. Embedded +`attributes` in the JSON take precedence. A tiny, self-contained example (two heads, including a non-contiguous id set) lives in [sample_data/](sample_data/). diff --git a/eomt/aux_cls.py b/eomt/aux_cls.py index 64c32fa..e72c3b8 100644 --- a/eomt/aux_cls.py +++ b/eomt/aux_cls.py @@ -1,8 +1,8 @@ """Secondary per-instance classification heads ("attributes") for EoMT. The primary task (instance segmentation over ``nc`` classes) is unchanged. Each -aux head predicts an extra attribute per *detected instance* — typology, -laterality, severity, … — read from the matched query's embedding. +aux head predicts an extra attribute per *detected instance* — color, material, +size, … — read from the matched query's embedding. The supervision reuses EoMT's own Hungarian matcher (``model.eomt.criterion.matcher``) so each attribute is trained on the **same** diff --git a/eomt/config.py b/eomt/config.py index 3cfd637..8c75765 100644 --- a/eomt/config.py +++ b/eomt/config.py @@ -35,19 +35,29 @@ class AuxHeadSpec: ``name`` keys the head everywhere (model ``ModuleDict``, COCO ``attributes`` field, checkpoint metadata). ``names`` maps the contiguous ``0..num_classes-1`` ids to human labels. Several specs ⇒ several independent heads. + + ``applies_to`` optionally scopes the head to a subset of **primary** classes + (contiguous ids): the head is only supervised on — and only emitted for — + instances whose primary class is in the set. ``None`` (the default) ⇒ the head + applies to every class (the original behaviour). This is the hard class-routing + that lets different attributes attach to different primary classes. """ name: str num_classes: int names: dict[int, str] = field(default_factory=dict) + applies_to: frozenset[int] | None = None def aux_specs_to_meta(specs: list[AuxHeadSpec] | None) -> list[dict]: """Serialize aux-head specs for a checkpoint.""" - return [ - {"name": s.name, "num_classes": int(s.num_classes), "names": dict(s.names)} - for s in (specs or []) - ] + out: list[dict] = [] + for s in specs or []: + d = {"name": s.name, "num_classes": int(s.num_classes), "names": dict(s.names)} + if s.applies_to is not None: + d["applies_to"] = sorted(int(c) for c in s.applies_to) + out.append(d) + return out def aux_specs_from_meta(meta: list[dict] | None) -> list[AuxHeadSpec]: @@ -55,7 +65,9 @@ def aux_specs_from_meta(meta: list[dict] | None) -> list[AuxHeadSpec]: out: list[AuxHeadSpec] = [] for d in meta or []: names = {int(k): str(v) for k, v in (d.get("names") or {}).items()} - out.append(AuxHeadSpec(str(d["name"]), int(d["num_classes"]), names)) + raw = d.get("applies_to") + applies_to = frozenset(int(c) for c in raw) if raw is not None else None + out.append(AuxHeadSpec(str(d["name"]), int(d["num_classes"]), names, applies_to)) return out diff --git a/eomt/data/coco.py b/eomt/data/coco.py index 49cc088..0e6bb6c 100644 --- a/eomt/data/coco.py +++ b/eomt/data/coco.py @@ -35,20 +35,159 @@ def _build_category_maps(coco): return cat2contig, contig2cat, names, len(cat_ids) -def _build_attribute_maps(coco, only: list[str] | None = None): +def _resolve_applies_to(raw, cat2contig, contig_names, attr_name): + """Resolve a definition's ``applies_to`` to a frozenset of contiguous primary ids. + + Entries may be category **names** (matched against ``contig_names``) or raw COCO + category **ids** (mapped through ``cat2contig``). ``None`` ⇒ head applies to every + primary class (returns ``None``). Unresolvable entries are warned about and dropped. + """ + if raw is None: + return None + name2contig = {v: k for k, v in contig_names.items()} + out: set[int] = set() + for entry in raw: + if isinstance(entry, str) and entry in name2contig: + out.add(name2contig[entry]) + elif isinstance(entry, int) and entry in cat2contig: + out.add(cat2contig[entry]) + else: + warnings.warn( + f"attribute {attr_name!r} 'applies_to' entry {entry!r} matches no " + "category (by name or raw id); ignoring it.", + stacklevel=2, + ) + return frozenset(out) + + +def _find_dataset_root(json_file: Path) -> Path: + """Dataset root that holds optional sidecar attribute files. + + The JSON's directory, or its parent when the JSON sits in an ``annotations/`` + subfolder (the COCO convention: ``/annotations/instances_train.json``). + """ + p = json_file.parent + return p.parent if p.name == "annotations" else p + + +def _split_name(json_file: Path) -> str: + """Infer the split name from a COCO json filename (``train`` / ``val`` / stem).""" + stem = json_file.stem.lower() + if "train" in stem: + return "train" + if "val" in stem: + return "val" + return stem + + +def _load_attr_schema(root: Path): + """Read ``/attributes.yaml`` into a normalized top-level ``attributes`` list. + + Each head is ``{name, categories:[{id,name}], applies_to?}``. ``categories`` may be + written as bare labels (auto-id 0..n-1) or explicit ``{id,name}`` dicts. Returns + ``None`` when the schema file is absent. + """ + import yaml + + schema_path = root / "attributes.yaml" + if not schema_path.exists(): + return None + doc = yaml.safe_load(schema_path.read_text()) or {} + norm = [] + for d in doc.get("attributes") or []: + cats = [] + for i, c in enumerate(d.get("categories") or []): + if isinstance(c, dict): + cats.append({"id": int(c["id"]), "name": str(c.get("name", c["id"]))}) + else: # bare label -> auto id by position + cats.append({"id": i, "name": str(c)}) + entry = {"name": d["name"], "categories": cats} + if d.get("applies_to") is not None: + entry["applies_to"] = list(d["applies_to"]) + norm.append(entry) + return norm + + +def _load_attr_values(root: Path, json_file: Path): + """Per-annotation attribute values keyed by **annotation id** (order-independent). + + Reads ``/attributes/.json`` (``{ann_id: {head: value}}``); values may + be category labels or raw ids. Returns ``{int ann_id: {head: value}}`` or ``None``. + """ + import json + + cand = root / "attributes" / f"{_split_name(json_file)}.json" + if not cand.exists(): + return None + raw = json.loads(cand.read_text()) + return {int(k): v for k, v in raw.items()} + + +def _merge_attribute_sidecar(coco, json_file) -> None: + """Merge an optional on-disk attribute sidecar into the in-memory COCO handle. + + No-op unless ``/attributes.yaml`` exists **and** the JSON has no embedded + ``attributes`` (embedded definitions win — one source of truth). When it applies, + the schema becomes the top-level ``attributes`` list and per-annotation values + (keyed by annotation id) are written into each ``ann["attributes"]`` as raw + category ids, after which discovery/loading is identical to the embedded case. + So a plain COCO dataset always works; drop in the sidecar and it is picked up. + """ + dataset = getattr(coco, "dataset", {}) or {} + if dataset.get("attributes"): # embedded schema present -> ignore sidecar + return + schema = _load_attr_schema(_find_dataset_root(Path(json_file))) + if not schema: + return + dataset["attributes"] = schema + name2id = {d["name"]: {c["name"]: c["id"] for c in d["categories"]} for d in schema} + valid_ids = {d["name"]: {c["id"] for c in d["categories"]} for d in schema} + values = _load_attr_values(_find_dataset_root(Path(json_file)), Path(json_file)) or {} + for ann in dataset.get("annotations", []) or []: + vals = values.get(ann["id"]) + if not vals: + continue + merged = {} + for head, v in vals.items(): + if head not in name2id: + continue + rid = name2id[head].get(v, v) # value may be a label or a raw id + if isinstance(rid, int) and rid in valid_ids[head]: + merged[head] = rid + if merged: + ann["attributes"] = {**ann.get("attributes", {}), **merged} + + +def _attr_label(spec: AuxHeadSpec, id_maps: dict, ann_attrs: dict, primary_cls: int) -> int: + """Contiguous attribute id for one instance, or ``-100`` (ignored by the aux CE). + + ``-100`` when the attribute is missing / out-of-vocab (never silently trained as + class 0), **or** when the head is class-scoped (``applies_to``) and this instance's + primary class is out of scope — the hard-routing choke point, so the loss / + accuracy paths need no change (they already skip ``-100``). + """ + if spec.applies_to is not None and primary_cls not in spec.applies_to: + return -100 + return id_maps[spec.name].get(ann_attrs.get(spec.name), -100) + + +def _build_attribute_maps(coco, only: list[str] | None = None, *, cat2contig=None, contig_names=None): """Discover secondary attributes from a COCO handle. Reads the (non-standard but valid) top-level ``attributes`` list:: "attributes": [ - {"name": "typology", "categories": [{"id": 0, "name": "scratch"}, ...]}, - {"name": "severity", "categories": [...]}, + {"name": "color", "categories": [{"id": 0, "name": "red"}, ...]}, + {"name": "posture", "categories": [...], "applies_to": ["cat", 7]}, ] - and the per-annotation ``"attributes": {"typology": 0, "severity": 2}`` field. + and the per-annotation ``"attributes": {"color": 0, "posture": 2}`` field. Falls back to inferring each attribute's id set from the annotations when a - definition omits ``categories``. Returns ``(specs, id_maps)`` where ``id_maps`` - is ``{attr_name: {raw_id: contiguous_id}}``. + definition omits ``categories``. The optional ``applies_to`` scopes a head to a + subset of primary classes (category names or raw ids); it is resolved against + ``cat2contig`` / ``contig_names`` (pass them from ``_build_category_maps``). + Returns ``(specs, id_maps)`` where ``id_maps`` is ``{attr_name: {raw_id: + contiguous_id}}``. """ dataset = getattr(coco, "dataset", {}) or {} defs = dataset.get("attributes") or [] @@ -57,6 +196,8 @@ def _build_attribute_maps(coco, only: list[str] | None = None): if not defs: return [], {} + cat2contig = cat2contig or {} + contig_names = contig_names or {} anns = dataset.get("annotations") or [] specs: list[AuxHeadSpec] = [] id_maps: dict[str, dict] = {} @@ -81,7 +222,10 @@ def _build_attribute_maps(coco, only: list[str] | None = None): raw2contig = {r: i for i, r in enumerate(raw_ids)} disp = {r: str(r) for r in raw_ids} names = {raw2contig[r]: disp.get(r, str(r)) for r in raw_ids} - specs.append(AuxHeadSpec(name=name, num_classes=len(raw2contig), names=names)) + applies_to = _resolve_applies_to(d.get("applies_to"), cat2contig, contig_names, name) + specs.append( + AuxHeadSpec(name=name, num_classes=len(raw2contig), names=names, applies_to=applies_to) + ) id_maps[name] = raw2contig return specs, id_maps @@ -112,6 +256,7 @@ def __init__( self.img_dir = Path(img_dir) self.imgsz = imgsz self.coco = COCO(str(json_file)) + _merge_attribute_sidecar(self.coco, json_file) self.cat2contig, self.contig2cat, self.names, self.num_classes = ( _build_category_maps(self.coco) ) @@ -124,7 +269,9 @@ def __init__( else: only = None if attributes is True else ([] if attributes is False else attributes) self.aux_specs, self._attr_id_maps = ( - ([], {}) if only == [] else _build_attribute_maps(self.coco, only=only) + ([], {}) if only == [] else _build_attribute_maps( + self.coco, only=only, cat2contig=self.cat2contig, contig_names=self.names + ) ) self.ids = [ i @@ -156,13 +303,11 @@ def __getitem__(self, idx: int): if m.sum() == 0: continue masks.append(torch.from_numpy(m)) - classes.append(self.cat2contig[ann["category_id"]]) + cls = self.cat2contig[ann["category_id"]] + classes.append(cls) ann_attrs = ann.get("attributes", {}) for spec in self.aux_specs: - raw = ann_attrs.get(spec.name) - # Missing / out-of-vocab → -100 (ignored by the aux CE), never - # silently trained as class 0. - attrs[spec.name].append(self._attr_id_maps[spec.name].get(raw, -100)) + attrs[spec.name].append(_attr_label(spec, self._attr_id_maps, ann_attrs, cls)) image_tv = tv_tensors.Image( torch.from_numpy(np.array(img)).permute(2, 0, 1) # (3, H, W) uint8 @@ -207,11 +352,6 @@ class CocoDetection(Dataset): **normalized ``cxcywh`` boxes** ``(num_inst, 4)`` in ``[0, 1]`` (relative to the square input) and their contiguous class ids. Boxes ride the **same** transforms as masks via ``tv_tensors.BoundingBoxes``, so LSJ/flip/crop apply identically. - - .. warning:: - Horizontal flip swaps left/right, which **corrupts laterality-style aux - labels** exactly as it does for the seg dataset. Train laterality detection - runs with ``flip_prob=0`` (see ``scripts/train_parts_large.sh``). """ def __init__( @@ -232,6 +372,7 @@ def __init__( self.img_dir = Path(img_dir) self.imgsz = imgsz self.coco = COCO(str(json_file)) + _merge_attribute_sidecar(self.coco, json_file) self.cat2contig, self.contig2cat, self.names, self.num_classes = ( _build_category_maps(self.coco) ) @@ -240,7 +381,9 @@ def __init__( else: only = None if attributes is True else ([] if attributes is False else attributes) self.aux_specs, self._attr_id_maps = ( - ([], {}) if only == [] else _build_attribute_maps(self.coco, only=only) + ([], {}) if only == [] else _build_attribute_maps( + self.coco, only=only, cat2contig=self.cat2contig, contig_names=self.names + ) ) self.ids = [ i @@ -279,11 +422,11 @@ def __getitem__(self, idx: int): if w <= 0 or h <= 0: continue boxes.append([x, y, x + w, y + h]) # xyxy - classes.append(self.cat2contig[ann["category_id"]]) + cls = self.cat2contig[ann["category_id"]] + classes.append(cls) ann_attrs = ann.get("attributes", {}) for spec in self.aux_specs: - raw = ann_attrs.get(spec.name) - attrs[spec.name].append(self._attr_id_maps[spec.name].get(raw, -100)) + attrs[spec.name].append(_attr_label(spec, self._attr_id_maps, ann_attrs, cls)) H, W = img.height, img.width image_tv = tv_tensors.Image( @@ -357,6 +500,7 @@ def __init__( self.mean = mean self.std = std self.coco = COCO(str(json_file)) + _merge_attribute_sidecar(self.coco, json_file) self.cat2contig, self.contig2cat, self.names, self.num_classes = ( _build_category_maps(self.coco) ) @@ -365,7 +509,9 @@ def __init__( else: only = None if attributes is True else ([] if attributes is False else attributes) self.aux_specs, self._attr_id_maps = ( - ([], {}) if only == [] else _build_attribute_maps(self.coco, only=only) + ([], {}) if only == [] else _build_attribute_maps( + self.coco, only=only, cat2contig=self.cat2contig, contig_names=self.names + ) ) self.ids = sorted(self.coco.getImgIds()) diff --git a/eomt/engine/predict.py b/eomt/engine/predict.py index e46cc5e..c47935f 100644 --- a/eomt/engine/predict.py +++ b/eomt/engine/predict.py @@ -50,13 +50,22 @@ def predict_image( ) tensor = torch.from_numpy(chw).unsqueeze(0).to(device) out = model(tensor) + # Class-scoped aux heads: pass their primary-class scope so postprocess emits + # ``ids = -1`` for detections the head does not apply to (the inference side of + # the hard class-routing). Unscoped heads (``applies_to=None``) are omitted here. + aux_scopes = { + s.name: s.applies_to + for s in getattr(model, "aux_specs", []) + if s.applies_to is not None + } if getattr(model, "family", "instance") == "detect": return postprocess_detection( out, conf_thres, (orig_w, orig_h), max_det=max_det, preprocess_meta=meta, + aux_scopes=aux_scopes, ) return postprocess_instance( out, conf_thres, (orig_w, orig_h), max_det=max_det, - mask_thresh=mask_thresh, preprocess_meta=meta, + mask_thresh=mask_thresh, preprocess_meta=meta, aux_scopes=aux_scopes, ) diff --git a/eomt/engine/track.py b/eomt/engine/track.py index dc71b24..8bc3d03 100644 --- a/eomt/engine/track.py +++ b/eomt/engine/track.py @@ -3,7 +3,7 @@ Runs the still-image inference (:func:`~eomt.engine.predict.predict_image`) on every frame of a video, associates detections across frames into persistent ``track_id``s with ByteTrack, and (optionally) writes an annotated video. For models with auxiliary -attribute heads (e.g. ``laterality`` / ``typology``) the tracker associates on the +attribute heads (e.g. ``color`` / ``material``) the tracker associates on the **main class + box only** — attributes flicker frame-to-frame and must not drive identity — while each track keeps a running mean of the attribute probabilities and reports a temporally-smoothed label (``aux_track``). diff --git a/eomt/engine/train.py b/eomt/engine/train.py index f9484cf..4a9dfad 100644 --- a/eomt/engine/train.py +++ b/eomt/engine/train.py @@ -159,11 +159,19 @@ def _compute_aux_class_weights(train_ds) -> dict: annotations (no image decode); missing / out-of-vocab values are ignored. """ anns = getattr(train_ds.coco, "dataset", {}).get("annotations", []) or [] + cat2contig = train_ds.cat2contig out: dict[str, torch.Tensor] = {} for spec in train_ds.aux_specs: id_map = train_ds._attr_id_maps[spec.name] counts = torch.zeros(spec.num_classes) for a in anns: + # Skip out-of-scope instances for class-scoped heads: they are routed to + # -100 and never contribute to the head's loss, so they must not skew its + # class-frequency weights either. + if spec.applies_to is not None: + cls = cat2contig.get(a.get("category_id")) + if cls is None or cls not in spec.applies_to: + continue cid = id_map.get(a.get("attributes", {}).get(spec.name), -100) if 0 <= cid < spec.num_classes: counts[cid] += 1 @@ -172,6 +180,43 @@ def _compute_aux_class_weights(train_ds) -> dict: return out +def _aux_coverage(train_ds) -> dict[str, tuple[int, int]]: + """Per head ``{name: (supervised, in_scope)}`` instance counts from train annotations. + + ``in_scope`` counts instances whose primary class the head applies to (every + instance for an unscoped head); ``supervised`` counts those of them carrying a + tagged, in-vocab value (not ignored as ``-100``). Diagnostic only — mirrors the + routing / ignore logic in :func:`eomt.data.coco._attr_label` without decoding images. + """ + anns = getattr(train_ds.coco, "dataset", {}).get("annotations", []) or [] + cat2contig = train_ds.cat2contig + out: dict[str, tuple[int, int]] = {} + for spec in train_ds.aux_specs: + id_map = train_ds._attr_id_maps[spec.name] + supervised = in_scope = 0 + for a in anns: + cls = cat2contig.get(a.get("category_id")) + if cls is None: + continue + if spec.applies_to is not None and cls not in spec.applies_to: + continue + in_scope += 1 + if id_map.get(a.get("attributes", {}).get(spec.name), -100) >= 0: + supervised += 1 + out[spec.name] = (supervised, in_scope) + return out + + +def _aux_scope_str(spec, names: dict) -> str: + """Human-readable ``applies_to`` scope for the banner ("" when unscoped).""" + if spec.applies_to is None: + return "" + if len(spec.applies_to) > 6: + return f" on {len(spec.applies_to)} classes" + labels = sorted(names.get(c, str(c)) for c in spec.applies_to) + return " on {" + ",".join(labels) + "}" + + def _write_run_config(path: Path, cfg: dict) -> None: """Persist the resolved run hyper-parameters (incl. model ``size``) to YAML. @@ -478,9 +523,14 @@ def train( aux_specs = train_ds.aux_specs print(f"[data] train: {len(train_ds)} images, {nc} classes") if aux_specs: + cov = _aux_coverage(train_ds) print( "[data] aux heads: " - + ", ".join(f"{s.name}({s.num_classes})" for s in aux_specs) + + "; ".join( + f"{s.name}({s.num_classes}){_aux_scope_str(s, train_ds.names)}: " + f"{cov[s.name][0]}/{cov[s.name][1]} tagged" + for s in aux_specs + ) ) train_loader = DataLoader( train_ds, diff --git a/eomt/postprocess.py b/eomt/postprocess.py index 86b8e0b..648c3b5 100644 --- a/eomt/postprocess.py +++ b/eomt/postprocess.py @@ -15,6 +15,33 @@ import torch.nn.functional as F # noqa: N812 +def _build_aux_result(aux_logits: dict, sel, classes, aux_scopes: dict | None) -> dict: + """Per-head ``{name: {"ids", "probs"}}`` for the kept queries, class-scope-gated. + + For a class-scoped head (``aux_scopes[name]`` a set of primary class ids), any kept + detection whose class is out of scope gets ``ids = -1`` (the "not applicable" + sentinel) and zeroed ``probs`` — the inference side of the hard routing, so + downstream (plotting / track) can skip heads that don't apply to a class. + """ + aux_scopes = aux_scopes or {} + res: dict = {} + for name, lg in aux_logits.items(): + probs = lg[0].float().softmax(dim=-1)[sel] # (N, ns) + ids = probs.argmax(dim=-1) + scope = aux_scopes.get(name) + if scope is not None and ids.numel(): + out_of_scope = torch.tensor( + [int(c) not in scope for c in classes.tolist()], + dtype=torch.bool, device=ids.device, + ) + ids = ids.clone() + ids[out_of_scope] = -1 + probs = probs.clone() + probs[out_of_scope] = 0.0 + res[name] = {"ids": ids, "probs": probs} + return res + + def boxes_from_masks(masks: torch.Tensor) -> torch.Tensor: """Derive ``xyxy`` boxes (pixel coords) from boolean masks ``(N, H, W)``.""" n = masks.shape[0] @@ -89,6 +116,7 @@ def postprocess_detection( *, max_det: int = 100, preprocess_meta: dict | None = None, + aux_scopes: dict | None = None, **_: object, ) -> dict: """Convert raw EoMT **detection** output to the canonical instance dict (no masks). @@ -146,10 +174,7 @@ def postprocess_detection( "classes": classes.long(), } if aux_logits is not None: - result["aux"] = {} - for name, lg in aux_logits.items(): - probs = lg[0].float().softmax(dim=-1)[sel] # (N, ns) - result["aux"][name] = {"ids": probs.argmax(dim=-1), "probs": probs} + result["aux"] = _build_aux_result(aux_logits, sel, classes, aux_scopes) return result @@ -162,6 +187,7 @@ def postprocess_instance( mask_thresh: float = 0.5, min_mask_area: float = 0.0, preprocess_meta: dict | None = None, + aux_scopes: dict | None = None, **_: object, ) -> dict: """Convert raw EoMT output to the canonical instance-seg dict. @@ -246,8 +272,5 @@ def postprocess_instance( "masks": masks, } if aux_logits is not None: - result["aux"] = {} - for name, lg in aux_logits.items(): - probs = lg[0].float().softmax(dim=-1)[sel] # (N, ns) - result["aux"][name] = {"ids": probs.argmax(dim=-1), "probs": probs} + result["aux"] = _build_aux_result(aux_logits, sel, classes, aux_scopes) return result diff --git a/eomt/visualize.py b/eomt/visualize.py index c7e434e..6651a23 100644 --- a/eomt/visualize.py +++ b/eomt/visualize.py @@ -48,6 +48,8 @@ def _aux_label( rows = [] for head, pred in aux.items(): idx = int(pred["ids"][i]) + if idx < 0: # -1 sentinel: class-scoped head does not apply to this instance + continue prob = float(pred["probs"][i][idx]) label = str(idx) if aux_names and head in aux_names: diff --git a/pyproject.toml b/pyproject.toml index f69d115..15d89fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "attr-eomt" -version = "1.0.0" +version = "1.1.0" description = "Standalone EoMT (Encoder-only Mask Transformer) for instance segmentation, with DINOv2 init, COCO training/validation and inference." readme = "README.md" requires-python = ">=3.10" diff --git a/sample_data/data.yaml b/sample_data/data.yaml index 69473bd..8622c6a 100644 --- a/sample_data/data.yaml +++ b/sample_data/data.yaml @@ -25,6 +25,29 @@ # # The attribute names used here (`scale`, `viewpoint`) are only examples — use # whatever attributes your task needs, and as many heads as you like. +# +# Class-conditional heads (optional): give an attribute definition an +# `applies_to` list of primary-class names or ids, and that head is only trained +# on — and only emitted for — instances of those classes. This is hard routing on +# the primary class, so different attributes can attach to different classes and +# each keeps its own label vocabulary. Omit `applies_to` and the head applies to +# every class (the original behaviour): +# +# "attributes": [ +# {"name": "posture", "categories": [...], "applies_to": ["cat", "dog"]} +# ] +# +# Missing / untagged values are always fine: an instance with no value for a head +# (or an out-of-scope instance for an `applies_to` head) is ignored by that head's +# loss (-100), never trained as class 0. So partially-tagged datasets just work. +# +# Sidecar format (optional): instead of embedding attributes in the COCO JSON you +# may keep the JSON as plain COCO and drop, in the dataset root, an +# `attributes.yaml` schema + `attributes/.json` values keyed by annotation +# id (`{ann_id: {head: value}}`, value = label or raw id). If present (and the JSON +# has no embedded `attributes`), it is merged in memory before loading — so a plain +# COCO dataset always works, and the sidecar is picked up automatically if you add +# it. Embedded `attributes` in the JSON take precedence. path: . train_images: images/train diff --git a/site/index.html b/site/index.html index 0aef0d4..a09770a 100644 --- a/site/index.html +++ b/site/index.html @@ -415,7 +415,7 @@

The training recipe, and the knobs

⚑ -

Laterality gotcha. If an attribute encodes left/right, train with flip_prob=0 — a horizontal flip mirrors the pixels but not the label, silently corrupting it.

+

Orientation gotcha. If an attribute encodes orientation, train with flip_prob=0 — a horizontal flip mirrors the pixels but not the label, silently corrupting it.

diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 3d94f7d..31183e6 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -62,8 +62,8 @@ def test_aux_heads_train_and_infer(): from eomt.config import AuxHeadSpec torch.manual_seed(0) - specs = [AuxHeadSpec("typology", 4, {0: "a", 1: "b", 2: "c", 3: "d"}), - AuxHeadSpec("severity", 3, {0: "lo", 1: "mid", 2: "hi"})] + specs = [AuxHeadSpec("color", 4, {0: "a", 1: "b", 2: "c", 3: "d"}), + AuxHeadSpec("material", 3, {0: "lo", 1: "mid", 2: "hi"})] model = build_model("s", nc=NC, imgsz=IMGSZ, aux_heads=specs).train() x = torch.randn(2, 3, IMGSZ, IMGSZ) mask_labels = [ @@ -72,32 +72,32 @@ def test_aux_heads_train_and_infer(): ] class_labels = [torch.tensor([0, 1]), torch.tensor([2])] aux_labels = { - "typology": [torch.tensor([1, 3]), torch.tensor([0])], - "severity": [torch.tensor([2, 0]), torch.tensor([1])], + "color": [torch.tensor([1, 3]), torch.tensor([0])], + "material": [torch.tensor([2, 0]), torch.tensor([1])], } out = model(x, mask_labels=mask_labels, class_labels=class_labels) total = out["loss"] a_loss, per_head = aux_loss(model, out, mask_labels, class_labels, aux_labels) - assert set(per_head) == {"typology", "severity"} + assert set(per_head) == {"color", "material"} (total + a_loss).backward() # default head is a small MLP -> assert grads flow through *some* head param - assert any(p.grad is not None for p in model.aux_heads["typology"].parameters()) + assert any(p.grad is not None for p in model.aux_heads["color"].parameters()) acc = aux_accuracy(model, out, mask_labels, class_labels, aux_labels) - assert set(acc) == {"typology", "severity"} + assert set(acc) == {"color", "material"} # inference forward exposes per-head logits; postprocess attaches them model.eval() with torch.no_grad(): out = model(x) - assert out["aux_queries_logits"]["typology"].shape == (2, EOMT_CONFIGS["s"].num_queries, 4) + assert out["aux_queries_logits"]["color"].shape == (2, EOMT_CONFIGS["s"].num_queries, 4) res = postprocess_instance( {k: v[:1] if torch.is_tensor(v) else {n: t[:1] for n, t in v.items()} for k, v in out.items() if k != "query_embed"}, conf_thres=0.0, original_size=(20, 15), max_det=5, ) - assert set(res["aux"]) == {"typology", "severity"} - assert res["aux"]["typology"]["probs"].shape[1] == 4 + assert set(res["aux"]) == {"color", "material"} + assert res["aux"]["color"]["probs"].shape[1] == 4 def test_aux_ignore_index(): @@ -127,6 +127,131 @@ def test_aux_ignore_index(): assert aux_accuracy(model, out2, None, None, all_ignored, indices=indices)["typ"] == (0, 0) +def test_aux_applies_to_roundtrip(): + """A class-scoped head's applies_to survives checkpoint (de)serialization.""" + from eomt.config import AuxHeadSpec, aux_specs_from_meta, aux_specs_to_meta + + specs = [ + AuxHeadSpec("posture", 2, {0: "sit", 1: "stand"}, frozenset({1, 3})), + AuxHeadSpec("coat", 3, {0: "short", 1: "long", 2: "curly"}, None), + ] + meta = aux_specs_to_meta(specs) + assert meta[0]["applies_to"] == [1, 3] + assert "applies_to" not in meta[1] # unscoped -> key omitted + back = aux_specs_from_meta(meta) + assert back[0].applies_to == frozenset({1, 3}) + assert back[1].applies_to is None + + +def _write_mini_coco(tmp_path, with_applies_to: bool): + """Write a 2-instance COCO (cat + dog) + one image; return (img_dir, json_path).""" + import json + + from PIL import Image + + img_dir = tmp_path / "images" + img_dir.mkdir() + Image.fromarray(np.zeros((20, 20, 3), dtype=np.uint8)).save(img_dir / "im1.png") + + def sq(x0, y0, x1, y1): # polygon segmentation for a rectangle + return [[x0, y0, x1, y0, x1, y1, x0, y1]] + + attr_def = { + "name": "posture", + "categories": [{"id": 0, "name": "sit"}, {"id": 1, "name": "stand"}, {"id": 2, "name": "lie"}], + } + if with_applies_to: + attr_def["applies_to"] = ["cat"] # head applies to 'cat' only + coco = { + "images": [{"id": 1, "file_name": "im1.png", "width": 20, "height": 20}], + "categories": [{"id": 1, "name": "cat"}, {"id": 2, "name": "dog"}], + "attributes": [attr_def], + "annotations": [ + {"id": 10, "image_id": 1, "category_id": 1, "iscrowd": 0, + "bbox": [1, 1, 8, 8], "area": 64, "segmentation": sq(1, 1, 9, 9), + "attributes": {"posture": 1}}, # cat, tagged 'stand' + {"id": 11, "image_id": 1, "category_id": 2, "iscrowd": 0, + "bbox": [10, 10, 8, 8], "area": 64, "segmentation": sq(10, 10, 18, 18), + "attributes": {"posture": 2}}, # dog, tagged 'lie' (out of scope) + ], + } + jf = tmp_path / "instances_train.json" + jf.write_text(json.dumps(coco)) + return img_dir, jf + + +def test_aux_class_routing_in_dataset(tmp_path): + """applies_to routes out-of-scope instances to -100 at label construction.""" + from eomt.data.coco import CocoInstanceSeg + from eomt.data.transforms import build_val_transform + + # cat ids sorted -> cat=contig 0, dog=contig 1; head 'posture' applies to cat. + img_dir, jf = _write_mini_coco(tmp_path, with_applies_to=True) + ds = CocoInstanceSeg(img_dir, jf, imgsz=IMGSZ, transform=build_val_transform(IMGSZ)) + assert ds.aux_specs[0].applies_to == frozenset({0}) # 'cat' resolved to contig 0 + _, _, classes, attrs = ds[0] + posture = attrs["posture"] + # cat instance keeps its tagged id (1); dog instance is routed to -100. + cat_pos = (classes == 0).nonzero(as_tuple=True)[0] + dog_pos = (classes == 1).nonzero(as_tuple=True)[0] + assert posture[cat_pos].item() == 1 + assert posture[dog_pos].item() == -100 + + +def test_aux_no_applies_to_trains_all_classes(tmp_path): + """Without applies_to (regression), the head is supervised on every class.""" + from eomt.data.coco import CocoInstanceSeg + from eomt.data.transforms import build_val_transform + + img_dir, jf = _write_mini_coco(tmp_path, with_applies_to=False) + ds = CocoInstanceSeg(img_dir, jf, imgsz=IMGSZ, transform=build_val_transform(IMGSZ)) + assert ds.aux_specs[0].applies_to is None + _, _, classes, attrs = ds[0] + # both instances keep their tagged posture (no routing to -100) + assert set(attrs["posture"].tolist()) == {1, 2} + + +def test_postprocess_gates_aux_by_class(): + """Class-scoped aux heads emit ids=-1 for detections whose class is out of scope.""" + from eomt.postprocess import _build_aux_result + + aux_logits = {"posture": torch.tensor([[[2.0, 0.0], [0.0, 3.0]]])} # (1, Q=2, ns=2) + sel = torch.tensor([0, 1]) + classes = torch.tensor([0, 1]) # det0 -> class 0 (in scope), det1 -> class 1 (out) + res = _build_aux_result(aux_logits, sel, classes, {"posture": frozenset({0})}) + assert res["posture"]["ids"].tolist() == [0, -1] + assert float(res["posture"]["probs"][1].sum()) == 0.0 + # unscoped -> unchanged + res2 = _build_aux_result(aux_logits, sel, classes, {}) + assert res2["posture"]["ids"].tolist() == [0, 1] + + +def test_attribute_sidecar_merge(tmp_path): + """A plain COCO + attributes.yaml + attributes/.json is merged in memory.""" + import json + + from eomt.data.coco import _merge_attribute_sidecar + + (tmp_path / "annotations").mkdir() + (tmp_path / "attributes").mkdir() + (tmp_path / "attributes.yaml").write_text( + "attributes:\n - name: posture\n categories: [sit, stand, lie]\n applies_to: [cat]\n" + ) + (tmp_path / "attributes" / "train.json").write_text( + json.dumps({"10": {"posture": "stand"}, "11": {"posture": 2}}) + ) + + class FakeCoco: + dataset = {"annotations": [{"id": 10, "category_id": 1}, {"id": 11, "category_id": 2}]} + + coco = FakeCoco() + _merge_attribute_sidecar(coco, tmp_path / "annotations" / "instances_train.json") + assert coco.dataset["attributes"][0]["name"] == "posture" + by_id = {a["id"]: a for a in coco.dataset["annotations"]} + assert by_id[10]["attributes"] == {"posture": 1} # label 'stand' -> raw id 1 + assert by_id[11]["attributes"] == {"posture": 2} # raw id passthrough + + def test_resolve_checkpoint_folder(tmp_path): """A run/weights folder resolves to best.pt (infer) or last.pt (resume).""" from eomt.serialization import resolve_checkpoint diff --git a/tools/attr_tagger/README.md b/tools/attr_tagger/README.md index 8b8ddcb..a8528d8 100644 --- a/tools/attr_tagger/README.md +++ b/tools/attr_tagger/README.md @@ -57,7 +57,9 @@ Standard COCO plus a top-level `attributes` schema and a per-annotation ```jsonc "attributes": [ - { "name": "color", "categories": [{ "id": 0, "name": "red" }, { "id": 1, "name": "green" }] } + { "name": "color", "categories": [{ "id": 0, "name": "red" }, { "id": 1, "name": "green" }] }, + // optional: scope a head to a subset of primary classes (by name or id) + { "name": "posture", "categories": [...], "applies_to": ["cat", "dog"] } ], "annotations": [ { "id": 1, "image_id": 1, "category_id": 7, "bbox": [...], "segmentation": [...], @@ -67,3 +69,9 @@ Standard COCO plus a top-level `attributes` schema and a per-annotation Untagged instances simply omit the value; the training loader treats those as ignored (`-100`), so a partially-tagged file is still valid to train on. + +**Class-conditional heads.** Add an optional `applies_to` list to an attribute +definition and that head is only trained on — and only predicted for — instances +of those primary classes (hard routing on the class). Omit it and the head +applies to every class. Different attributes can therefore attach to different +classes, each with its own label set.