Object detection (centroid)

pip install label-studio
label-studio start             # opens http://localhost:8080
# Export → "YOLO" → unzip, then copy to the Jetson:
scp -r yolo_export/images yolo_export/labels <your-user>@<jetson-ip>:~/sutraai/data/raw/detect/

import numpy as np, timm, torch, torch.nn as nn
from pathlib import Path

class CentroidDetector(nn.Module):
    """MobileNetV2 cut at stride 8 + a 1x1 head: one class prediction per grid cell. Index 0 = background."""
    def __init__(self, n_classes):
        super().__init__()
        self.backbone = timm.create_model("mobilenetv2_050", pretrained=True, features_only=True, out_indices=(2,))
        ch = self.backbone.feature_info.channels()[-1]
        self.head = nn.Sequential(nn.Conv2d(ch, 32, 1), nn.ReLU(), nn.Conv2d(32, n_classes + 1, 1))
    def forward(self, x):
        return self.head(self.backbone(x)[0])            # [B, n_classes+1, H/8, W/8]

def centroid_target(label_path, grid):
    t = np.zeros((grid, grid), dtype=np.int64)          # 0 = background
    if Path(label_path).exists():
        for line in open(label_path):
            if not line.strip(): continue
            c, cx, cy, w, h = map(float, line.split())
            gx, gy = min(int(cx * grid), grid - 1), min(int(cy * grid), grid - 1)
            t[gy, gx] = int(c) + 1                        # objects are 1..N
    return t

def cells_to_points(logits, threshold=0.5):
    """logits [C+1, G, G] → list of (gx, gy, class) where an object is predicted."""
    prob = torch.softmax(logits, 0)
    conf, cls = prob[1:].max(0)
    ys, xs = torch.where(conf > threshold)
    return [(int(x), int(y), int(cls[y, x])) for y, x in zip(ys, xs)]

def match(pred, true, max_dist=1):
    tp, used = 0, set()
    for p in pred:
        for j, t in enumerate(true):
            if j not in used and p[2] == t[2] and max(abs(p[0] - t[0]), abs(p[1] - t[1])) <= max_dist:
                tp += 1; used.add(j); break
    return tp, len(pred) - tp, len(true) - tp            # tp, fp, fn

import sys, random, numpy as np, torch
from pathlib import Path
from PIL import Image
from torchvision import transforms as T
ROOT = Path.home() / "sutraai"; sys.path.insert(0, str(ROOT))
from sutra_core.detect import centroid_target

SIZE = 96; GRID = SIZE // 8                              # try SIZE = 160 if eyes are too small at 96
IMG, LAB = ROOT / "data/raw/detect/images", ROOT / "data/raw/detect/labels"
CLASSES = ["eye_open", "eye_closed", "yawn", "no_yawn"]   # same order as in Label Studio's classes.txt
imgs = sorted([*IMG.glob("*.jpg"), *IMG.glob("*.png")])
random.seed(42); random.shuffle(imgs)
n = len(imgs); split = {"train": imgs[:int(.7 * n)], "val": imgs[int(.7 * n):int(.85 * n)], "test": imgs[int(.85 * n):]}
tf = T.Compose([T.Grayscale(3), T.Resize((SIZE, SIZE)), T.ColorJitter(0.3, 0.3), T.ToTensor(),
                T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

class CentroidDataset(torch.utils.data.Dataset):
    def __init__(self, paths): self.paths = paths
    def __len__(self): return len(self.paths)
    def __getitem__(self, i):
        p = self.paths[i]
        x = tf(Image.open(p).convert("RGB"))
        return x, torch.tensor(centroid_target(LAB / (p.stem + ".txt"), GRID))

ds = {k: CentroidDataset(v) for k, v in split.items()}
counts = np.zeros(len(CLASSES), int)
for p in imgs:
    t = centroid_target(LAB / (p.stem + ".txt"), GRID); counts += np.bincount(t.ravel(), minlength=len(CLASSES) + 1)[1:]
print({k: len(v) for k, v in ds.items()}, dict(zip(CLASSES, counts)))

%run 40_detect_data.ipynb
import json, torch.nn.functional as F
from torch.utils.data import DataLoader
from sutra_core.detect import CentroidDetector, cells_to_points, match

torch.manual_seed(42)
model = CentroidDetector(len(CLASSES)).cuda()
weights = torch.tensor([1.0] + [20.0] * len(CLASSES)).cuda()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
dl = DataLoader(ds["train"], batch_size=32, shuffle=True, num_workers=4)

def f1_on(dataset, thr=0.5):
    model.eval(); tp = fp = fn = 0
    with torch.no_grad():
        for x, t in DataLoader(dataset, batch_size=32):
            out = model(x.cuda()).cpu()
            for o, ti in zip(out, t):
                true = [(int(gx), int(gy), int(ti[gy, gx]) - 1) for gy, gx in zip(*torch.where(ti > 0))]
                a, b, c = match(cells_to_points(o, thr), true); tp += a; fp += b; fn += c
    return 2 * tp / max(1, 2 * tp + fp + fn)

best = -1
for ep in range(60):
    model.train()
    for x, t in dl:
        loss = F.cross_entropy(model(x.cuda()), t.cuda(), weight=weights)
        opt.zero_grad(); loss.backward(); opt.step()
    f1 = f1_on(ds["val"])
    print(json.dumps({"epoch": ep + 1, "loss": round(loss.item(), 4), "val_f1": round(f1, 4)}))
    if f1 > best:
        best = f1; (ROOT / "models/detect_v1").mkdir(parents=True, exist_ok=True)
        torch.save(model.state_dict(), ROOT / "models/detect_v1/best.pt")
model.load_state_dict(torch.load(ROOT / "models/detect_v1/best.pt"))
test_f1 = f1_on(ds["test"]); print("test F1:", round(test_f1, 3))

import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 6, figsize=(16, 3))
model.eval()
for ax, p in zip(axes, split["test"][:6]):
    x, _ = ds["test"][split["test"].index(p)]
    with torch.no_grad(): pts = cells_to_points(model(x[None].cuda())[0].cpu())
    ax.imshow(Image.open(p).convert("L").resize((SIZE, SIZE)), cmap="gray")
    for gx, gy, c in pts:
        ax.plot((gx + .5) * 8, (gy + .5) * 8, "o", ms=8); ax.text((gx + .5) * 8, (gy + .5) * 8 - 4, CLASSES[c], fontsize=7, color="yellow")
    ax.axis("off")
plt.show()

%run 40_detect_data.ipynb
from sutra_core.detect import CentroidDetector
EXP = ROOT / "exports"
model = CentroidDetector(len(CLASSES)); model.load_state_dict(torch.load(ROOT / "models/detect_v1/best.pt", map_location="cpu")); model.eval()
torch.onnx.export(model, torch.randn(1, 3, SIZE, SIZE), str(EXP / "detect_v1.onnx"),
                  input_names=["input"], output_names=["cells"], opset_version=17, dynamo=False)
!/usr/src/tensorrt/bin/trtexec --onnx={EXP}/detect_v1.onnx --saveEngine={EXP}/detect_v1_fp16.engine --fp16 2>&1 | grep -E "Latency|PASSED|FAILED"