"""Read-only Ultrack database helpers for napari preview tooling."""
from __future__ import annotations
import pickle
from dataclasses import dataclass
from pathlib import Path
from collections.abc import Iterable
import numpy as np
[docs]
@dataclass(frozen=True)
class HierarchyCutState:
node_ids: tuple[int, ...]
height: float | None
[docs]
@dataclass(frozen=True)
class UltrackDbPreview:
labels: np.ndarray
status: str
probabilities: dict[int, float]
label_to_node_id: dict[int, int]
node_id_to_label: dict[int, int]
node_annotations: dict[int, str]
[docs]
def as_tuple(self):
return (
self.labels,
self.status,
self.probabilities,
self.label_to_node_id,
self.node_id_to_label,
self.node_annotations,
)
def _engine(db_path: Path):
import sqlalchemy as sqla
return sqla.create_engine(
f"sqlite:///{db_path}", connect_args={"check_same_thread": False}
)
[docs]
def query_middle_frame(db_path: Path) -> int | None:
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
engine = _engine(db_path)
try:
with Session(engine) as session:
frames = sorted(int(r[0]) for r in session.query(NodeDB.t).distinct().all())
except Exception:
return None
finally:
engine.dispose()
return frames[len(frames) // 2] if frames else None
[docs]
def query_frame_range(db_path: Path) -> tuple[int, ...]:
"""Return sorted distinct frame indices present in the database."""
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
engine = _engine(db_path)
try:
with Session(engine) as session:
frames = sorted(int(r[0]) for r in session.query(NodeDB.t).distinct().all())
except Exception:
return ()
finally:
engine.dispose()
return tuple(frames)
[docs]
def query_connected_nodes(
db_path: Path, selected_node_id: int
) -> tuple[dict[int, float], dict[int, float]]:
from sqlalchemy.orm import Session
from ultrack.core.database import LinkDB
engine = _engine(db_path)
predecessors: dict[int, float] = {}
successors: dict[int, float] = {}
try:
with Session(engine) as session:
rows = (
session.query(LinkDB.source_id, LinkDB.target_id, LinkDB.weight)
.filter(
(LinkDB.source_id == int(selected_node_id))
| (LinkDB.target_id == int(selected_node_id))
)
.all()
)
for src, tgt, weight in rows:
wf = float(weight if weight is not None else 1.0)
if int(tgt) == int(selected_node_id):
src_id = int(src)
predecessors[src_id] = predecessors.get(src_id, 1.0) * wf
if int(src) == int(selected_node_id):
tgt_id = int(tgt)
successors[tgt_id] = successors.get(tgt_id, 1.0) * wf
finally:
engine.dispose()
return predecessors, successors
[docs]
@dataclass(frozen=True)
class NodeEdge:
neighbor_id: int
weight: float # raw LinkDB.weight (NOT a product)
annotation: str # annotation_name(LinkDB.annotation): REAL/FAKE/UNKNOWN
direction: str # "pred" (neighbor -> selected) or "succ" (selected -> neighbor)
neighbor_t: int # NodeDB.t of the neighbor
neighbor_prob: float # NodeDB.node_prob (1.0 if unset: NULL or ultrack's -1.0 sentinel)
neighbor_annot: str # annotation_name(NodeDB.node_annot)
weight_is_null: bool = False # raw LinkDB.weight was NULL (rendered 1.0)
[docs]
@dataclass(frozen=True)
class NodeEdges:
selected_id: int
selected_t: int
selected_prob: float
selected_annot: str
edges: tuple[NodeEdge, ...]
def _node_prob_or_unknown(prob) -> float:
"""Map an unset node probability to 1.0.
Ultrack's ``NodeDB.node_prob`` is nullable but defaults to the sentinel
``-1.0`` for nodes that were never scored, so both NULL and any negative
value mean "no probability assigned" and read back as 1.0.
"""
if prob is None:
return 1.0
value = float(prob)
return 1.0 if value < 0.0 else value
[docs]
def query_node_edges(db_path: Path, node_id: int) -> NodeEdges:
"""Return raw, un-collapsed per-link edges incident to ``node_id``.
Unlike :func:`query_connected_nodes` (which multiplies weights when a
neighbor pair shares multiple links), each ``LinkDB`` row becomes its own
:class:`NodeEdge`, so per-edge weights stay legible. Pure read.
"""
from sqlalchemy.orm import Session
from ultrack.core.database import LinkDB, NodeDB
node_id = int(node_id)
engine = _engine(db_path)
try:
with Session(engine) as session:
link_rows = (
session.query(
LinkDB.source_id, LinkDB.target_id, LinkDB.weight, LinkDB.annotation
)
.filter(
(LinkDB.source_id == node_id) | (LinkDB.target_id == node_id)
)
.all()
)
# Gather every node id we need metadata for (selected + neighbors).
wanted: set[int] = {node_id}
for src, tgt, _weight, _annot in link_rows:
wanted.add(int(src))
wanted.add(int(tgt))
node_rows = (
session.query(
NodeDB.id, NodeDB.t, NodeDB.node_prob, NodeDB.node_annot
)
.filter(NodeDB.id.in_(sorted(wanted)))
.all()
)
node_meta: dict[int, tuple[int, float, str]] = {}
for nid, t, prob, annot in node_rows:
node_meta[int(nid)] = (
int(t),
_node_prob_or_unknown(prob),
annotation_name(annot),
)
edges: list[NodeEdge] = []
for src, tgt, weight, annot in link_rows:
src_id = int(src)
tgt_id = int(tgt)
if src_id == tgt_id:
continue # skip self / degenerate links
weight_is_null = weight is None
wf = float(weight if weight is not None else 1.0)
if tgt_id == node_id:
neighbor_id, direction = src_id, "pred"
else:
neighbor_id, direction = tgt_id, "succ"
n_t, n_prob, n_annot = node_meta.get(
neighbor_id, (0, 1.0, "UNKNOWN")
)
edges.append(
NodeEdge(
neighbor_id=neighbor_id,
weight=wf,
annotation=annotation_name(annot),
direction=direction,
neighbor_t=n_t,
neighbor_prob=n_prob,
neighbor_annot=n_annot,
weight_is_null=weight_is_null,
)
)
sel_t, sel_prob, sel_annot = node_meta.get(node_id, (0, 1.0, "UNKNOWN"))
finally:
engine.dispose()
return NodeEdges(
selected_id=node_id,
selected_t=sel_t,
selected_prob=sel_prob,
selected_annot=sel_annot,
edges=tuple(edges),
)
def _empty_annotation_counts() -> dict[str, int]:
return {"REAL": 0, "FAKE": 0, "UNKNOWN": 0}
def _annotation_counts_from_rows(rows) -> dict[str, int]:
counts = _empty_annotation_counts()
for annotation, count in rows:
name = annotation_name(annotation)
counts[name] = counts.get(name, 0) + int(count or 0)
return counts
[docs]
def link_annotation_counts(
db_path: Path, selected_node_id: int | None = None
) -> dict[str, int]:
"""Return link annotation counts for the whole DB or one node's incident links."""
engine = _engine(db_path)
try:
from sqlalchemy import func
from sqlalchemy.orm import Session
from ultrack.core.database import LinkDB
with Session(engine) as session:
query = session.query(LinkDB.annotation, func.count(LinkDB.source_id))
if selected_node_id is not None:
node_id = int(selected_node_id)
query = query.filter(
(LinkDB.source_id == node_id) | (LinkDB.target_id == node_id)
)
rows = query.group_by(LinkDB.annotation).all()
return _annotation_counts_from_rows(rows)
except Exception:
return _empty_annotation_counts()
finally:
engine.dispose()
def _finite_values(rows) -> list[float]:
values: list[float] = []
for row in rows:
try:
raw = row[0]
except (TypeError, KeyError, IndexError):
raw = row
if raw is None:
continue
try:
value = float(raw)
except (TypeError, ValueError):
continue
if np.isfinite(value):
values.append(value)
return values
def _stats_text(label: str, values: list[float]) -> str | None:
if not values:
return None
arr = np.asarray(values, dtype=np.float64)
return (
f"{label}: "
f"min {float(np.min(arr)):.3f}, "
f"median {float(np.median(arr)):.3f}, "
f"mean {float(np.mean(arr)):.3f}, "
f"max {float(np.max(arr)):.3f}"
)
[docs]
def summary_text(db_path: Path, frame: int) -> str:
from sqlalchemy import func
from sqlalchemy.orm import Session
from ultrack.core.database import LinkDB, NodeDB, VarAnnotation
try:
from ultrack.core.database import OverlapDB
except (ImportError, AttributeError):
# Capability probe only: an older ultrack without OverlapDB. A runtime
# error must not be silently turned into "no overlaps".
OverlapDB = None
engine = _engine(db_path)
try:
with Session(engine) as session:
n_nodes = int(session.query(func.count(NodeDB.id)).scalar() or 0)
n_links = int(session.query(func.count(LinkDB.source_id)).scalar() or 0)
node_prob_values = _finite_values(
session.query(NodeDB.node_prob)
.filter(NodeDB.node_prob.isnot(None))
.all()
)
link_weight_values = _finite_values(
session.query(LinkDB.weight)
.filter(LinkDB.weight.isnot(None))
.all()
)
n_real = int(
session.query(func.count(NodeDB.id))
.filter(NodeDB.node_annot == VarAnnotation.REAL)
.scalar()
or 0
)
n_fake = int(
session.query(func.count(NodeDB.id))
.filter(NodeDB.node_annot == VarAnnotation.FAKE)
.scalar()
or 0
)
link_annotation_rows = (
session.query(LinkDB.annotation, func.count(LinkDB.source_id))
.group_by(LinkDB.annotation)
.all()
)
link_annotations = _annotation_counts_from_rows(link_annotation_rows)
frame_nodes = session.query(NodeDB).filter(NodeDB.t == frame).all()
selected = sum(1 for n in frame_nodes if getattr(n, "selected", False))
node_ids = [int(n.id) for n in frame_nodes]
overlaps = 0
if node_ids:
if OverlapDB is not None:
try:
overlaps = int(
session.query(func.count(OverlapDB.node_id))
.filter(
OverlapDB.node_id.in_(node_ids)
| OverlapDB.ancestor_id.in_(node_ids)
)
.scalar()
or 0
)
except Exception:
overlaps = 0
stats_parts = [
text
for text in (
_stats_text("node prob", node_prob_values),
_stats_text("edge weight", link_weight_values),
)
if text is not None
]
lines = [
f"Database: {n_nodes} nodes, {n_links} links",
f"Node annotations: REAL {n_real}, FAKE {n_fake}",
(
f"Link annotations: REAL links {link_annotations['REAL']}, "
f"FAKE links {link_annotations['FAKE']}"
),
f"Frame {frame}: {selected} selected, {overlaps} overlaps",
]
if stats_parts:
lines.extend(stats_parts)
return "\n".join(lines)
finally:
engine.dispose()
[docs]
def query_distinct_heights(db_path: Path) -> tuple[float, ...]:
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
engine = _engine(db_path)
try:
with Session(engine) as session:
return tuple(
float(r[0])
for r in session.query(NodeDB.height)
.distinct()
.order_by(NodeDB.height)
.all()
if r[0] is not None
)
finally:
engine.dispose()
[docs]
def query_union_sizes(db_path: Path, frame: int) -> tuple[int, ...]:
"""Distinct atom-union sizes (``height`` = number of merged atoms) present in
``frame``, sorted ascending. Drives the vertical "union size" slider: index 0
is the finest (1 atom = individual atoms), higher indices merge more atoms.
"""
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
engine = _engine(db_path)
try:
with Session(engine) as session:
rows = (
session.query(NodeDB.height)
.filter(NodeDB.t == frame)
.distinct()
.order_by(NodeDB.height)
.all()
)
finally:
engine.dispose()
return tuple(int(round(float(r[0]))) for r in rows if r[0] is not None)
[docs]
def query_union_color_classes(
db_path: Path, frame: int, union_size: int
) -> tuple[tuple[int, ...], ...]:
"""Group the size-``union_size`` candidates of ``frame`` into color classes of
mutually non-overlapping candidates (a greedy graph coloring of the shared-atom
overlap graph from ``OverlapDB``).
Each returned class is a tuple of node ids that can be painted together in one
full-frame partition without conflict, so every candidate of this size is shown
in exactly one class while keeping the number of classes (= horizontal slider
positions) small. Returns ``()`` if no candidates of this size exist.
"""
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
try:
from ultrack.core.database import OverlapDB
except (ImportError, AttributeError):
# Capability probe only: an older ultrack without OverlapDB. A runtime
# error must not be silently turned into "no overlaps".
OverlapDB = None
engine = _engine(db_path)
try:
with Session(engine) as session:
node_ids = [
int(r[0])
for r in session.query(NodeDB.id)
.filter(NodeDB.t == frame)
.filter(NodeDB.height == float(union_size))
.order_by(NodeDB.id)
.all()
]
id_set = set(node_ids)
edges: list[tuple[int, int]] = []
if OverlapDB is not None and node_ids:
# OverlapDB stores node_id < ancestor_id, so every intra-set edge has
# its node_id endpoint in node_ids; filtering on that side is exact.
rows = (
session.query(OverlapDB.node_id, OverlapDB.ancestor_id)
.filter(OverlapDB.node_id.in_(node_ids))
.all()
)
edges = [(int(a), int(b)) for a, b in rows if int(b) in id_set]
finally:
engine.dispose()
return greedy_color_classes(node_ids, edges)
[docs]
def greedy_color_classes(
node_ids: list[int], edges: Iterable[tuple[int, int]]
) -> tuple[tuple[int, ...], ...]:
"""Partition ``node_ids`` into classes with no intra-class edge (greedy graph
coloring). Welsh–Powell order (highest degree first) keeps the class count low;
ids preserve their original order within each class. Pure helper, no DB."""
if not node_ids:
return ()
adj: dict[int, set[int]] = {nid: set() for nid in node_ids}
id_set = set(node_ids)
for a, b in edges:
if a != b and a in id_set and b in id_set:
adj[a].add(b)
adj[b].add(a)
order = sorted(node_ids, key=lambda n: (-len(adj[n]), n))
colors: dict[int, int] = {}
for nid in order:
used = {colors[m] for m in adj[nid] if m in colors}
c = 0
while c in used:
c += 1
colors[nid] = c
n_colors = max(colors.values()) + 1
classes: list[list[int]] = [[] for _ in range(n_colors)]
for nid in node_ids: # stable id order within each class
classes[colors[nid]].append(nid)
return tuple(tuple(c) for c in classes)
def _paint_union_partition(
union_nodes: list, fill_nodes: list, plane_shape: tuple[int, int]
) -> tuple[np.ndarray, list]:
"""Paint a full-frame partition: the merged ``union_nodes`` (the selected merge
group) take priority, then ``fill_nodes`` cover the leftover territory.
``fill_nodes`` should be ordered most-merged first; each is painted all-or-nothing
only if its whole mask is still free, so a leftover atom is shown inside the
largest union that still fits rather than dropping back to an individual atom.
Returns ``(labels, ordered_nodes)`` where ``ordered_nodes[i]`` is painted with
label ``i + 1`` — matching ``node_preview_metadata`` so click selection lines up.
"""
parsed: list[tuple] = []
max_y = max_x = 0
for node in union_nodes:
result = node_mask_and_bbox(node)
if result is None:
continue
(y0, x0, y1, x1), mask = result
max_y, max_x = max(max_y, y1), max(max_x, x1)
parsed.append((node, (y0, x0, y1, x1), mask, True))
for node in fill_nodes:
result = node_mask_and_bbox(node)
if result is None:
continue
(y0, x0, y1, x1), mask = result
max_y, max_x = max(max_y, y1), max(max_x, x1)
parsed.append((node, (y0, x0, y1, x1), mask, False))
base_y, base_x = plane_shape
labels = np.zeros((max(base_y, max_y, 1), max(base_x, max_x, 1)), dtype=np.uint32)
ordered: list = []
for node, (y0, x0, y1, x1), mask, priority in parsed:
region = labels[y0:y1, x0:x1]
if region.shape != mask.shape:
continue
if priority:
# Unions never overlap within a class, so they paint freely.
sub = mask
else:
# All-or-nothing: only paint a leftover union/atom if its whole footprint
# is still free, so larger unions win and nothing is split mid-region.
if (mask & (region != 0)).any():
continue
sub = mask
if not sub.any():
continue
region[sub] = len(ordered) + 1
ordered.append(node)
return labels, ordered
[docs]
def render_union_partition(
db_path: Path,
frame: int,
color_node_ids: Iterable[int],
*,
plane_shape: tuple[int, int],
union_size: int | None = None,
) -> UltrackDbPreview:
"""Full-frame partition for one horizontal slider position: the candidates in
``color_node_ids`` painted as merged regions, with the leftover territory filled
by the largest available unions (size ≤ ``union_size``) so regions outside this
merge group still show at their most-merged state instead of as raw atoms.
``union_size`` caps how merged the leftover fill may be; when ``None`` it is taken
from the selected candidates (their ``height``), defaulting to atoms only.
"""
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
color_node_ids = tuple(int(n) for n in color_node_ids)
engine = _engine(db_path)
try:
with Session(engine) as session:
union_nodes: list = []
if color_node_ids:
rows = (
session.query(NodeDB)
.filter(NodeDB.t == frame)
.filter(NodeDB.id.in_(color_node_ids))
.all()
)
by_id = {int(n.id): n for n in rows}
union_nodes = [by_id[i] for i in color_node_ids if i in by_id]
if union_size is None:
heights = [
float(n.height) for n in union_nodes if n.height is not None
]
union_size = int(round(max(heights))) if heights else 1
# Leftover fill: every non-selected candidate up to the current level,
# most-merged first, so larger unions claim free territory before atoms.
fill_nodes = (
session.query(NodeDB)
.filter(NodeDB.t == frame)
.filter(NodeDB.height <= float(union_size))
.filter(NodeDB.id.notin_(color_node_ids) if color_node_ids else True)
.order_by(NodeDB.height.desc(), NodeDB.id)
.all()
)
finally:
engine.dispose()
labels, ordered = _paint_union_partition(union_nodes, fill_nodes, plane_shape)
if not ordered:
return _empty_preview(plane_shape, f"No candidates for frame {frame}.")
probabilities, label_to_node_id, node_id_to_label = node_preview_metadata(ordered)
annotations = node_annotation_metadata(ordered)
n_merged = len(union_nodes)
return UltrackDbPreview(
labels=labels,
status=(
f"Frame {frame}: {len(ordered)} region(s), "
f"{n_merged} merged candidate(s)."
),
probabilities=probabilities,
label_to_node_id=label_to_node_id,
node_id_to_label=node_id_to_label,
node_annotations=annotations,
)
[docs]
def query_hierarchy_cut_states(
db_path: Path, frame: int, *, source_index: int | None = None
) -> tuple[HierarchyCutState, ...]:
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
from ultrack.utils.constants import NO_PARENT
source_node_ids: tuple[int, ...] = ()
if source_index is not None:
try:
from itasc.tracking_ultrack.multi_threshold import query_source_node_ids
source_node_ids = query_source_node_ids(db_path, int(source_index))
except Exception:
source_node_ids = ()
engine = _engine(db_path)
try:
with Session(engine) as session:
query = (
session.query(NodeDB.id, NodeDB.hier_parent_id, NodeDB.height)
.filter(NodeDB.t == frame)
.order_by(NodeDB.height, NodeDB.id)
)
if source_index is not None:
query = query.filter(NodeDB.id.in_(source_node_ids))
rows = []
for nid, pid, height in query.all():
if height is None:
continue
parent_id = NO_PARENT if pid is None else int(pid)
rows.append((int(nid), parent_id, float(height)))
except Exception:
return tuple(HierarchyCutState((), float(h)) for h in query_distinct_heights(db_path))
finally:
engine.dispose()
if not rows:
return ()
node_ids = {nid for nid, _, _ in rows}
heights_by_id = {nid: height for nid, _, height in rows}
parent_by_id = {
nid: pid for nid, pid, _ in rows if pid != NO_PARENT and pid in node_ids
}
children: dict[int, set[int]] = {}
for child_id, parent_id in parent_by_id.items():
children.setdefault(parent_id, set()).add(child_id)
active = {nid for nid, _, _ in rows if nid not in children}
if not active:
active = set(node_ids)
states: list[HierarchyCutState] = []
seen: set[tuple[int, ...]] = set()
def append_state() -> None:
ordered = tuple(sorted(active, key=lambda n: (heights_by_id[n], n)))
if ordered in seen:
return
seen.add(ordered)
height = max((heights_by_id[n] for n in ordered), default=None)
states.append(HierarchyCutState(ordered, height))
append_state()
while True:
promotable = [
parent_id
for parent_id, child_ids in children.items()
if parent_id not in active and child_ids and child_ids.issubset(active)
]
if not promotable:
break
min_height = min(heights_by_id[parent_id] for parent_id in promotable)
for parent_id in sorted(
p for p in promotable if heights_by_id[p] == min_height
):
active.difference_update(children[parent_id])
active.add(parent_id)
append_state()
return tuple(states)
[docs]
def render_hierarchy_cut(
db_path: Path,
frame: int,
height: float,
*,
plane_shape: tuple[int, int],
) -> UltrackDbPreview:
from sqlalchemy.orm import Session, aliased
from ultrack.core.database import NodeDB
from ultrack.utils.constants import NO_PARENT
engine = _engine(db_path)
try:
with Session(engine) as session:
parent = aliased(NodeDB)
child = aliased(NodeDB)
same_child = (
session.query(child.id)
.where(child.hier_parent_id == NodeDB.id)
.where(child.height == NodeDB.height)
.where(NodeDB.height == height)
.exists()
)
nodes = (
session.query(NodeDB)
.outerjoin(parent, NodeDB.hier_parent_id == parent.id)
.where(NodeDB.t == frame)
.where(NodeDB.height <= height)
.where(
(NodeDB.hier_parent_id == NO_PARENT)
| ((NodeDB.height < height) & (parent.height > height))
| ((NodeDB.height == height) & (parent.height >= height))
)
.where(~same_child)
.all()
)
finally:
engine.dispose()
return finalize_hierarchy_nodes(
nodes,
frame,
plane_shape=plane_shape,
empty_msg=f"No segments at this threshold for frame {frame}.",
status_suffix=f"at h={height:.2f}",
)
[docs]
def render_hierarchy_cut_state(
db_path: Path,
frame: int,
state: HierarchyCutState,
*,
plane_shape: tuple[int, int],
) -> UltrackDbPreview:
if not state.node_ids:
return render_hierarchy_cut(
db_path,
frame,
float(state.height or 0.0),
plane_shape=plane_shape,
)
from sqlalchemy.orm import Session
from ultrack.core.database import NodeDB
engine = _engine(db_path)
try:
with Session(engine) as session:
rows = (
session.query(NodeDB)
.where(NodeDB.t == frame)
.where(NodeDB.id.in_(state.node_ids))
.all()
)
finally:
engine.dispose()
by_id = {int(n.id): n for n in rows}
nodes = [by_id[nid] for nid in state.node_ids if nid in by_id]
height = "—" if state.height is None else f"{state.height:.2f}"
return finalize_hierarchy_nodes(
nodes,
frame,
plane_shape=plane_shape,
empty_msg=f"No hierarchy state segments for frame {frame}.",
status_suffix=f"at cut state h={height}",
)
[docs]
def finalize_hierarchy_nodes(
nodes: Iterable,
frame: int,
*,
plane_shape: tuple[int, int],
empty_msg: str,
status_suffix: str,
) -> UltrackDbPreview:
nodes = list(nodes)
if not nodes:
return _empty_preview(plane_shape, empty_msg)
labels = paint_nodes(nodes, plane_shape)
probabilities, label_to_node_id, node_id_to_label = node_preview_metadata(nodes)
annotations = node_annotation_metadata(nodes)
return UltrackDbPreview(
labels=labels,
status=f"Frame {frame}: {len(nodes)} segment(s) {status_suffix}.",
probabilities=probabilities,
label_to_node_id=label_to_node_id,
node_id_to_label=node_id_to_label,
node_annotations=annotations,
)
def _empty_preview(plane_shape: tuple[int, int], status: str) -> UltrackDbPreview:
return UltrackDbPreview(
labels=np.zeros(plane_shape, dtype=np.uint32),
status=status,
probabilities={},
label_to_node_id={},
node_id_to_label={},
node_annotations={},
)
[docs]
def annotation_name(value) -> str:
if value is None:
return "UNKNOWN"
for raw in (getattr(value, "name", None), value, getattr(value, "value", None)):
if raw is None:
continue
if raw in {0, "0"}:
return "UNKNOWN"
if raw in {1, "1"}:
return "REAL"
if raw in {2, "2"}:
return "FAKE"
name = str(raw).split(".")[-1].upper()
if name in {"REAL", "FAKE", "UNKNOWN"}:
return name
return "UNKNOWN"
[docs]
def paint_nodes(nodes, plane_shape: tuple[int, int]) -> np.ndarray:
masks: list[tuple[int, tuple[int, int, int, int], np.ndarray]] = []
max_y = max_x = 0
for label, node in enumerate(nodes, start=1):
parsed = node_mask_and_bbox(node)
if parsed is None:
continue
bbox, mask = parsed
y0, x0, y1, x1 = bbox
max_y = max(max_y, y1)
max_x = max(max_x, x1)
masks.append((label, bbox, mask))
base_y, base_x = plane_shape
labels = np.zeros((max(base_y, max_y, 1), max(base_x, max_x, 1)), dtype=np.uint32)
for label, (y0, x0, y1, x1), mask in masks:
target = labels[y0:y1, x0:x1]
if target.shape != mask.shape:
continue
target[mask.astype(bool)] = label
return labels
[docs]
def node_mask_and_bbox(node):
try:
node_obj = node.pickle
if isinstance(node_obj, (bytes, memoryview)):
node_obj = pickle.loads(bytes(node_obj))
if node_obj is None:
return None
except Exception:
return None
if isinstance(node_obj, dict):
bbox, mask = node_obj.get("bbox"), node_obj.get("mask")
elif isinstance(node_obj, tuple) and len(node_obj) >= 2:
bbox, mask = node_obj[0], node_obj[1]
else:
bbox = getattr(node_obj, "bbox", None)
mask = getattr(node_obj, "mask", None)
if bbox is None or mask is None:
return None
bbox_array = np.asarray(bbox, dtype=int).ravel()
if bbox_array.size >= 6:
y0, x0, y1, x1 = (
int(bbox_array[1]),
int(bbox_array[2]),
int(bbox_array[4]),
int(bbox_array[5]),
)
elif bbox_array.size >= 4:
y0, x0, y1, x1 = (int(v) for v in bbox_array[:4])
else:
return None
mask_array = np.asarray(mask)
if mask_array.ndim == 3 and mask_array.shape[0] == 1:
mask_array = mask_array[0]
elif mask_array.ndim > 2:
mask_array = np.squeeze(mask_array)
if mask_array.ndim != 2:
return None
if mask_array.shape != (y1 - y0, x1 - x0):
return None
return (y0, x0, y1, x1), mask_array.astype(bool, copy=False)