Spaces:

abdelkader9090
/

analysis

Sleeping

analysis / match.py

Initial release

e67ded9 3 months ago

18.7 kB

﻿"""
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.")

	"""
	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_OSNET100, W_TORSO100, 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.")