""" Person Re-Identification — OSNet + Stripe Color + Soft Threshold ---------------------------------------------------------------- Key improvements over v2: 1. UPPER-BODY stripe histogram (avoids green grass noise) 2. Soft unmatched handling (unequal camera person counts) 3. Confidence-gated matching (low-confidence pairs = unmatched) 4. Per-channel torso color (BGR mean of torso strip, robust to lighting) 5. Detailed debug output (see exactly why each pair matched) Usage: python person_reid.py --cam1 path/cam1.jpg --cam2 path/cam2.jpg Requirements: pip install torch torchvision torchreid ultralytics scipy opencv-python Pillow """ import argparse import sys from pathlib import Path import cv2 import numpy as np import torch import torch.nn.functional as F import torchvision.transforms as T import torchvision.ops as ops import torchreid from PIL import Image from scipy.optimize import linear_sum_assignment # ───────────────────────────────────────────── # CONFIG # ───────────────────────────────────────────── OSNET_MODEL = "osnet_ain_x1_0" YOLO_MODEL = "yolo11x.pt" DETECT_CONF = 0.8 DETECT_IOU = 0.45 REID_INPUT_SIZE = (256, 128) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Match threshold — pairs below this are marked unmatched MATCH_THRESHOLD = 0.6 IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"} # Ensemble weights (must sum to 1.0) W_OSNET = 0.60 # deep embedding — primary signal W_TORSO = 0.20 # torso stripe color (shirt color, not background) W_SHAPE = 0.20 # relative body size prior — doubled to break same-outfit ties # ───────────────────────────────────────────── # MODELS # ───────────────────────────────────────────── def load_detector(): from ultralytics import YOLO m = YOLO(YOLO_MODEL) print(f"[✓] YOLO ({YOLO_MODEL})") return m def load_osnet(): m = torchreid.models.build_model( name=OSNET_MODEL, num_classes=1000, pretrained=True ) m = m.to(DEVICE).eval() print(f"[✓] OSNet ({OSNET_MODEL}, {DEVICE})") return m # ───────────────────────────────────────────── # DETECTION # ───────────────────────────────────────────── def detect_persons(detector, img_bgr): res = detector(img_bgr, classes=[0], conf=DETECT_CONF, iou=DETECT_IOU, max_det=50, verbose=False) boxes = res[0].boxes.xyxy scores = res[0].boxes.conf if boxes.numel() == 0: return [], np.zeros((0, 4), dtype=int) keep = ops.nms(boxes, scores, iou_threshold=DETECT_IOU) boxes = boxes[keep].cpu().numpy().astype(int) crops = [] H, W = img_bgr.shape[:2] for x1, y1, x2, y2 in boxes: x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(W, x2), min(H, y2) c = img_bgr[y1:y2, x1:x2] if c.size > 0: crops.append(c) return crops, boxes # ───────────────────────────────────────────── # SIGNAL 1 — OSNET EMBEDDING # ───────────────────────────────────────────── def _preprocess(crop_bgr): tf = T.Compose([ T.Resize(REID_INPUT_SIZE), T.ToTensor(), T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]), ]) rgb = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB) return tf(Image.fromarray(rgb)) def osnet_embeddings(osnet, crops): if not crops: return torch.zeros(0, 512) t = torch.stack([_preprocess(c) for c in crops]).to(DEVICE) with torch.no_grad(): f = osnet(t) f = F.normalize(f, dim=1) return f.cpu() # ───────────────────────────────────────────── # SIGNAL 2 — TORSO STRIPE COLOR # Strategy: take the middle-vertical 30% of the crop (torso band), # ignore the leftmost and rightmost 15% (background edges), # compute a hue histogram with 18 bins. # This is far less sensitive to green grass than full-body HSV. # ───────────────────────────────────────────── def torso_hue_hist(crop_bgr): h, w = crop_bgr.shape[:2] # torso band: rows 20%–55% of height r0, r1 = int(h * 0.20), int(h * 0.55) c0, c1 = int(w * 0.15), int(w * 0.85) band = crop_bgr[r0:r1, c0:c1] if band.size == 0: band = crop_bgr hsv = cv2.cvtColor(band, cv2.COLOR_BGR2HSV) # mask out low-saturation pixels (white/grey/black clothing handled separately) sat_mask = hsv[:,:,1] > 30 # if the shirt is mostly achromatic (white/black/grey), use value channel instead if sat_mask.sum() < band.size * 0.05: # use brightness histogram hist = cv2.calcHist([hsv], [2], None, [16], [0, 256]) else: hist = cv2.calcHist([hsv], [0], sat_mask.astype(np.uint8), [18], [0, 180]) cv2.normalize(hist, hist, alpha=1.0, norm_type=cv2.NORM_L1) return hist.flatten().astype(np.float32) def torso_sim_matrix(crops1, crops2): h1 = [torso_hue_hist(c) for c in crops1] h2 = [torso_hue_hist(c) for c in crops2] n, m = len(h1), len(h2) S = np.zeros((n, m), dtype=np.float32) for i in range(n): for j in range(m): # Bhattacharyya: 0=identical, 1=completely different → invert d = cv2.compareHist(h1[i], h2[j], cv2.HISTCMP_BHATTACHARYYA) S[i, j] = max(0.0, 1.0 - d) return S # ───────────────────────────────────────────── # SIGNAL 3 — RELATIVE SIZE (shape prior) # Persons that appear much larger in one cam vs the other # are less likely to be the same person at that moment. # Uses a soft Gaussian penalty on relative height difference. # ───────────────────────────────────────────── def shape_sim_matrix(boxes1, boxes2, H1, H2): def rel_h(boxes, H): return np.array([(b[3]-b[1])/H for b in boxes], dtype=np.float32) rh1 = rel_h(boxes1, H1) rh2 = rel_h(boxes2, H2) n, m = len(rh1), len(rh2) S = np.zeros((n, m), dtype=np.float32) for i in range(n): for j in range(m): diff = abs(rh1[i] - rh2[j]) S[i, j] = float(np.exp(-(diff**2) / 0.015)) return S # ───────────────────────────────────────────── # MATCHING # ───────────────────────────────────────────── def match_persons(feats1, feats2, crops1, crops2, boxes1, boxes2, H1, H2): n, m = len(feats1), len(feats2) empty = {"matches": [], "unmatched_cam1": list(range(n)), "unmatched_cam2": list(range(m)), "sim": np.zeros((n,m)), "osnet": np.zeros((n,m)), "torso": np.zeros((n,m)), "shape": np.zeros((n,m))} if n == 0 or m == 0: return empty osnet_S = torch.mm(feats1, feats2.T).numpy() torso_S = torso_sim_matrix(crops1, crops2) shape_S = shape_sim_matrix(boxes1, boxes2, H1, H2) ensemble = W_OSNET * osnet_S + W_TORSO * torso_S + W_SHAPE * shape_S # Hungarian on the ensemble cost row_ind, col_ind = linear_sum_assignment(1.0 - ensemble) matches = [] unmatched_c1 = set(range(n)) unmatched_c2 = set(range(m)) for r, c in zip(row_ind, col_ind): score = float(ensemble[r, c]) if score >= MATCH_THRESHOLD: matches.append((r, c, score)) unmatched_c1.discard(r) unmatched_c2.discard(c) return { "matches": matches, "unmatched_cam1": sorted(unmatched_c1), "unmatched_cam2": sorted(unmatched_c2), "sim": ensemble, "osnet": osnet_S, "torso": torso_S, "shape": shape_S, } # ───────────────────────────────────────────── # PRINT # ───────────────────────────────────────────── def print_results(R, n1, n2): print("\n" + "="*72) print(" PERSON RE-ID (OSNet {:.0f}% Torso-color {:.0f}% Shape {:.0f}%)".format( W_OSNET*100, W_TORSO*100, W_SHAPE*100)) print("="*72) print(f" Cam1: {n1} persons Cam2: {n2} persons Threshold: {MATCH_THRESHOLD}") print("-"*72) if not R["matches"]: print(" No matches above threshold.") else: print(f" {'Cam1':>6} {'Cam2':>6} {'Ensemble':>10} {'OSNet':>8} {'Torso':>8} {'Shape':>8} {'Grade':>8}") print(f" {'-'*6} {'-'*6} {'-'*10} {'-'*8} {'-'*8} {'-'*8} {'-'*8}") for r, c, s in sorted(R["matches"], key=lambda x: -x[2]): g = "HIGH" if s>0.72 else "MEDIUM" if s>0.58 else "LOW" print(f" {'P'+str(r+1):>6} {'P'+str(c+1):>6} {s:>10.4f}" f" {R['osnet'][r,c]:>8.4f} {R['torso'][r,c]:>8.4f}" f" {R['shape'][r,c]:>8.4f} {g:>8}") if R["unmatched_cam1"]: print(f"\n Unmatched Cam1: {[f'P{i+1}' for i in R['unmatched_cam1']]}") if R["unmatched_cam2"]: print(f" Unmatched Cam2: {[f'P{i+1}' for i in R['unmatched_cam2']]}") # Ensemble matrix print(f"\n Ensemble similarity matrix:") print(" " + "".join(f" Cam2-P{c+1:>2}" for c in range(R["sim"].shape[1]))) for r in range(R["sim"].shape[0]): print(f" Cam1-P{r+1}" + "".join(f" {v:.3f}" for v in R["sim"][r])) print("="*72+"\n") # ───────────────────────────────────────────── # VISUALIZATION # ───────────────────────────────────────────── PAIR_COLORS = [ (57, 255, 20), # neon green (0, 165, 255), # orange (255, 0, 255), # magenta (0, 255, 255), # yellow (255, 191, 0), # deep sky blue (147, 20, 255), # purple (0, 128, 255), # amber (0, 255, 128), # spring green ] UNMATCHED_COLOR = (60, 60, 220) # red-ish def _color(pair_idx): return PAIR_COLORS[pair_idx % len(PAIR_COLORS)] def annotate(img, boxes, person_map, scale): """ person_map: {person_idx: (pair_idx, score)} for matched person_idx not in map → unmatched """ out = img.copy() for idx, (x1,y1,x2,y2) in enumerate(boxes): x1,y1,x2,y2 = [int(v*scale) for v in (x1,y1,x2,y2)] x1,y1 = max(0,x1), max(0,y1) x2,y2 = min(out.shape[1]-1,x2), min(out.shape[0]-1,y2) if idx in person_map: pair_idx, score = person_map[idx] col = _color(pair_idx) label = f"P{idx+1} [{score:.2f}]" thick = 2 else: col = UNMATCHED_COLOR label = f"P{idx+1} [--]" thick = 2 cv2.rectangle(out, (x1,y1), (x2,y2), col, thick) (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2) ly = max(y1-4, th+6) cv2.rectangle(out, (x1, ly-th-4), (x1+tw+6, ly+2), col, -1) cv2.putText(out, label, (x1+3, ly-2), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,0), 2, cv2.LINE_AA) return out def show_visualization(img1, img2, boxes1, boxes2, R, output_path=None, show=True): map1 = {r: (pi, s) for pi, (r,c,s) in enumerate(R["matches"])} map2 = {c: (pi, s) for pi, (r,c,s) in enumerate(R["matches"])} max_w = 1800 scale = min(1.0, max_w / (img1.shape[1] + img2.shape[1])) def rs(img): return cv2.resize(img, (int(img.shape[1]*scale), int(img.shape[0]*scale)), interpolation=cv2.INTER_AREA) if scale < 1 else img a1 = annotate(rs(img1), boxes1, map1, scale) a2 = annotate(rs(img2), boxes2, map2, scale) # pad height H = max(a1.shape[0], a2.shape[0]) + 34 # extra room for legend def pad(img): p = np.zeros((H - img.shape[0], img.shape[1], 3), np.uint8) return np.vstack([img, p]) a1, a2 = pad(a1), pad(a2) div = np.full((H, 4, 3), 60, np.uint8) canvas = np.hstack([a1, div, a2]) # legend bottom strip lx, ly = 10, H - 26 for pi, (r, c, s) in enumerate(R["matches"]): col = _color(pi) txt = f" Cam1-P{r+1} <-> Cam2-P{c+1} {s:.3f} " (tw,th),_ = cv2.getTextSize(txt, cv2.FONT_HERSHEY_SIMPLEX, 0.52, 1) cv2.rectangle(canvas, (lx-2, ly-th-4), (lx+tw+2, ly+4), col, -1) cv2.putText(canvas, txt, (lx, ly), cv2.FONT_HERSHEY_SIMPLEX, 0.52, (0,0,0), 2, cv2.LINE_AA) lx += tw + 14 if lx > canvas.shape[1] - 200: # wrap to new line lx = 10 ly -= 28 output_path = Path(output_path) if output_path is not None else Path("reid_result.jpg") cv2.imwrite(str(output_path), canvas) print(f"[✓] Saved → {output_path}") if show: try: cv2.imshow("Person Re-ID", canvas) print("[→] Press any key to close.") cv2.waitKey(0) cv2.destroyAllWindows() except cv2.error: print("[!] No GUI available.") # ───────────────────────────────────────────── # MAIN # ───────────────────────────────────────────── def is_image_file(path: Path): return path.suffix.lower() in IMAGE_EXTS def sorted_image_files(folder_path): folder = Path(folder_path) if not folder.exists() or not folder.is_dir(): sys.exit(f"[ERROR] Not a folder: {folder_path}") files = sorted([p for p in folder.iterdir() if p.is_file() and is_image_file(p)]) if not files: sys.exit(f"[ERROR] No images found in: {folder_path}") return files def run(cam1_path, cam2_path, detector=None, osnet=None, output_path=None, show=True): img1 = cv2.imread(cam1_path) img2 = cv2.imread(cam2_path) if img1 is None: sys.exit(f"[ERROR] Cannot read: {cam1_path}") if img2 is None: sys.exit(f"[ERROR] Cannot read: {cam2_path}") print(f"\n[→] Cam1: {cam1_path} ({img1.shape[1]}×{img1.shape[0]})") print(f"[→] Cam2: {cam2_path} ({img2.shape[1]}×{img2.shape[0]})") if detector is None: detector = load_detector() if osnet is None: osnet = load_osnet() crops1, boxes1 = detect_persons(detector, img1) crops2, boxes2 = detect_persons(detector, img2) print(f"[✓] Detected → Cam1: {len(crops1)} Cam2: {len(crops2)}") if not crops1 or not crops2: print("[!] Need persons in both images.") return None f1 = osnet_embeddings(osnet, crops1) f2 = osnet_embeddings(osnet, crops2) R = match_persons(f1, f2, crops1, crops2, boxes1, boxes2, img1.shape[0], img2.shape[0]) print_results(R, len(crops1), len(crops2)) show_visualization(img1, img2, boxes1, boxes2, R, output_path=output_path, show=show) return R def run_folder_pairs(cam1_dir, cam2_dir, output_dir=None, show=True): files1 = sorted_image_files(cam1_dir) files2 = sorted_image_files(cam2_dir) count = min(len(files1), len(files2)) if len(files1) != len(files2): print(f"[!] Warning: cam1 has {len(files1)} images, cam2 has {len(files2)} images. Processing {count} synced pairs.") output_root = Path(output_dir) if output_dir else Path.cwd() output_root.mkdir(parents=True, exist_ok=True) detector = load_detector() osnet = load_osnet() for idx, (p1, p2) in enumerate(zip(files1[:count], files2[:count]), start=1): print(f"\n=== Pair {idx} ===") print(f" Cam1: {p1.name}") print(f" Cam2: {p2.name}") output_path = output_root / f"reid_result_{idx:04d}_{p1.stem}_{p2.stem}.jpg" run(str(p1), str(p2), detector=detector, osnet=osnet, output_path=str(output_path), show=show) if __name__ == "__main__": ap = argparse.ArgumentParser(description="Person Re-ID: OSNet + Torso Color + Shape") ap.add_argument("--cam1", required=True, help="Path to cam1 image or folder") ap.add_argument("--cam2", required=True, help="Path to cam2 image or folder") ap.add_argument("--threshold", type=float, default=MATCH_THRESHOLD) ap.add_argument("--w-osnet", type=float, default=W_OSNET) ap.add_argument("--w-torso", type=float, default=W_TORSO) ap.add_argument("--w-shape", type=float, default=W_SHAPE) ap.add_argument("--output-dir", default=None, help="Output directory for batch results") ap.add_argument("--no-show", action="store_true", help="Do not display GUI windows") args = ap.parse_args() MATCH_THRESHOLD = args.threshold W_OSNET = args.w_osnet W_TORSO = args.w_torso W_SHAPE = args.w_shape cam1_path = Path(args.cam1) cam2_path = Path(args.cam2) if cam1_path.is_dir() and cam2_path.is_dir(): run_folder_pairs(str(cam1_path), str(cam2_path), output_dir=args.output_dir, show=not args.no_show) elif cam1_path.is_file() and cam2_path.is_file(): output_path = None if args.output_dir: out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) output_path = out_dir / f"reid_result_{cam1_path.stem}_{cam2_path.stem}.jpg" run(str(cam1_path), str(cam2_path), output_path=str(output_path) if output_path else None, show=not args.no_show) else: sys.exit("[ERROR] Both --cam1 and --cam2 must be either both files or both folders.")