Your first image model

C:\sutraai-photos\
├── clean\     60 photos (.jpg)
├── stain\     60 photos
└── hole\      60 photos

ssh <your-user>@<jetson-ip> "mkdir -p ~/sutraai/data/raw/images"
scp -r C:\sutraai-photos\clean C:\sutraai-photos\stain C:\sutraai-photos\hole <your-user>@<jetson-ip>:~/sutraai/data/raw/images/

cd ~/sutraai/data/raw/images
for d in */; do printf "%-10s %s photos\n" "$d" "$(find "$d" -maxdepth 1 -type f \( -iname '*.jpg' -o -iname '*.jpeg' -o -iname '*.png' \) | wc -l)"; done
du -sh .        # 180 phone photos ≈ 0.5–1 GB
# expect: clean/ 60 photos, stain/ 60 photos, hole/ 60 photos (at least 50 each, roughly equal)

mkdir -p ~/sutraai/data/raw/_parked
mv -n ~/sutraai/data/raw/images/good ~/sutraai/data/raw/images/defect ~/sutraai/data/raw/_parked/
ls ~/sutraai/data/raw/images          # expect: clean  hole  stain

import random, shutil, hashlib
from pathlib import Path
import matplotlib.pyplot as plt
from PIL import Image, ImageOps

ROOT = Path.home() / "sutraai"
RAW, OUT = ROOT / "data/raw/images", ROOT / "data/processed/images"
CLASSES = ["clean", "stain", "hole"]             # Option B: ["screw", "nut", "washer"]. Only these folders are used.
EXTS = {".jpg", ".jpeg", ".png"}                 # compared in lower case, so .JPG from iPhones counts too
MAX_SIDE = 512
random.seed(42)
if OUT.exists(): shutil.rmtree(OUT)

def save_small(src, dst):
    img = ImageOps.exif_transpose(Image.open(src)).convert("RGB")   # apply the phone's rotation flag
    img.thumbnail((MAX_SIDE, MAX_SIDE))                              # longest side 512 px, keeps the shape
    img.save(dst.with_suffix(".jpg"), quality=92)

found = sorted(p.name for p in RAW.iterdir() if p.is_dir())
ignored, missing = [c for c in found if c not in CLASSES], [c for c in CLASSES if c not in found]
if ignored: print(f"⚠ ignoring folders not in CLASSES: {ignored} (park them outside data/raw/images, substep 4)")
assert not missing, f"no folder for {missing} in {RAW}"

counts, first_seen, conflicts = {}, {}, []
for cls_dir in [RAW / c for c in CLASSES]:
    files = sorted(f for f in cls_dir.iterdir() if f.is_file() and f.suffix.lower() in EXTS)
    skipped = [f.name for f in cls_dir.iterdir() if f.is_file() and f.suffix.lower() not in EXTS]
    if skipped: print(f"{cls_dir.name}: skipped {len(skipped)} files, e.g. {skipped[0]} (convert HEIC to JPG)")
    unique = []                                       # the same photo copied twice would leak across splits
    for f in files:
        h = hashlib.sha256(f.read_bytes()).hexdigest()
        if h in first_seen:
            other = first_seen[h]
            if other.parent == f.parent: print(f"  {cls_dir.name}: duplicate skipped, {f.name} = {other.name}")
            else: conflicts.append(f"{other.parent.name}/{other.name} = {f.parent.name}/{f.name}")
            continue
        first_seen[h] = f; unique.append(f)
    files = unique
    random.shuffle(files)
    n = len(files); n_tr, n_va = int(n * 0.70), int(n * 0.15)
    assert n >= 10, f"'{cls_dir.name}': only {n} JPG/PNG photos found in {cls_dir}. Need 50+ (aim for 60)."
    if n < 50: print(f"⚠ '{cls_dir.name}': {n} photos is enough for a dry run, not for a real model (take 50+)")
    for split, chunk in [("train", files[:n_tr]), ("val", files[n_tr:n_tr + n_va]), ("test", files[n_tr + n_va:])]:
        dst = OUT / split / cls_dir.name; dst.mkdir(parents=True, exist_ok=True)
        for f in chunk: save_small(f, dst / f.name)
    counts[cls_dir.name] = n
    print(f"{cls_dir.name:10} {n:4d} photos → train {n_tr}, val {n_va}, test {n - n_tr - n_va}")
assert not conflicts, "the same photo is in two class folders, delete the wrong copy: " + "; ".join(conflicts[:5])
print(counts)                                                      # e.g. {'clean': 60, 'hole': 60, 'stain': 60}
plt.bar(counts.keys(), counts.values()); plt.title("images per class"); plt.show()

# the same photo must never be in two splits
seen = {}
for p in OUT.rglob("*.jpg"):
    seen.setdefault(hashlib.sha256(p.read_bytes()).hexdigest(), []).append(p.relative_to(OUT))
leaks = [paths for paths in seen.values() if len({q.parts[0] for q in paths}) > 1]
print("duplicate photos across splits:", len(leaks))                # expect: 0
for paths in leaks[:5]: print("   ", " = ".join(map(str, paths)))   # names, if any, so you can delete one copy

# look at a few samples per class
fig, axes = plt.subplots(len(counts), 5, figsize=(12, 2.6 * len(counts)), squeeze=False)
for r, cls in enumerate(counts):
    for c, f in enumerate(sorted((OUT / "train" / cls).iterdir())[:5]):
        axes[r][c].imshow(Image.open(f)); axes[r][c].set_title(cls); axes[r][c].axis("off")
plt.show()

import json, timm, torch, numpy as np
from pathlib import Path
from torchvision import datasets, transforms as T
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix, classification_report

ROOT = Path.home() / "sutraai"
CFG = dict(size=96, arch="mobilenetv2_050", epochs=30, lr=1e-3, batch=32, seed=42, run="image_v1")
torch.manual_seed(CFG["seed"]); np.random.seed(CFG["seed"])
(ROOT / "models" / CFG["run"]).mkdir(parents=True, exist_ok=True)
(ROOT / "reports" / CFG["run"]).mkdir(parents=True, exist_ok=True)

mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
train_tf = T.Compose([T.RandomResizedCrop(CFG["size"], scale=(0.7, 1.0)), T.RandomHorizontalFlip(),
                      T.ColorJitter(0.3, 0.3, 0.2), T.ToTensor(), T.Normalize(mean, std)])
eval_tf = T.Compose([T.Resize(int(CFG["size"] * 1.15)), T.CenterCrop(CFG["size"]), T.ToTensor(), T.Normalize(mean, std)])
D = ROOT / "data/processed/images"
for split in ["train", "val", "test"]:                         # friendly check before ImageFolder complains
    empty = [c.name for c in sorted((D / split).iterdir()) if not any(c.iterdir())]
    assert not empty, f"{split}/ has no photos for {empty}: too few photos per class. Add photos, re-run NB10."
train_ds = datasets.ImageFolder(D / "train", train_tf)
val_ds, test_ds = datasets.ImageFolder(D / "val", eval_tf), datasets.ImageFolder(D / "test", eval_tf)
classes = train_ds.classes
loader = lambda ds, shuffle: DataLoader(ds, batch_size=CFG["batch"], shuffle=shuffle, num_workers=4)

model = timm.create_model(CFG["arch"], pretrained=True, num_classes=len(classes)).cuda()
opt = torch.optim.AdamW(model.parameters(), lr=CFG["lr"])
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, CFG["epochs"])
scaler = torch.amp.GradScaler()

def predict(ds):
    model.eval(); preds, labels = [], []
    with torch.no_grad():
        for x, y in loader(ds, False):
            preds += model(x.cuda()).argmax(1).cpu().tolist(); labels += y.tolist()
    return np.array(labels), np.array(preds)

best = 0.0
for epoch in range(CFG["epochs"]):
    model.train()
    for x, y in loader(train_ds, True):
        x, y = x.cuda(), y.cuda()
        with torch.autocast("cuda", dtype=torch.float16):          # mixed precision: faster, less memory
            loss = torch.nn.functional.cross_entropy(model(x), y)
        opt.zero_grad(); scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
    sched.step()
    yv, pv = predict(val_ds); val_acc = float((yv == pv).mean())
    print(json.dumps({"epoch": epoch + 1, "loss": round(loss.item(), 4), "val_acc": round(val_acc, 4)}))  # one JSON line per epoch: the web UI will read these later
    if val_acc > best:
        best = val_acc
        torch.save(model.state_dict(), ROOT / "models" / CFG["run"] / "best.pt")

model.load_state_dict(torch.load(ROOT / "models" / CFG["run"] / "best.pt"))
yt, pt = predict(test_ds)
metrics = {
    "run": CFG["run"], "config": CFG, "classes": classes,
    "val_acc": best, "test_acc": float((yt == pt).mean()),
    "confusion_matrix": confusion_matrix(yt, pt).tolist(),
    "per_class": classification_report(yt, pt, target_names=classes, output_dict=True),
}
json.dump(metrics, open(ROOT / "reports" / CFG["run"] / "metrics.json", "w"), indent=2)
print("test accuracy:", round(metrics["test_acc"], 3))
print(np.array(metrics["confusion_matrix"]))

import matplotlib.pyplot as plt
from PIL import Image
wrong = [i for i, (a, b) in enumerate(zip(yt, pt)) if a != b]
fig, axes = plt.subplots(1, max(1, min(6, len(wrong))), figsize=(14, 3), squeeze=False)
for ax, i in zip(axes[0], wrong[:6]):
    path, _ = test_ds.samples[i]
    ax.imshow(Image.open(path)); ax.set_title(f"true {classes[yt[i]]}\npred {classes[pt[i]]}"); ax.axis("off")
plt.show()