import os import io import sys import subprocess # --------------------------------------------------------------------------- # MiVOLO must be installed at runtime rather than via requirements.txt. # Its setup.py needs pkg_resources/torch importable at build time, and pip's # isolated build environment for git installs (used automatically when # processing requirements.txt on HF Spaces) doesn't reliably have them. # Installing it here, after torch/setuptools are already present in this # same environment, with build isolation disabled, avoids that failure. # --------------------------------------------------------------------------- try: import mivolo # noqa: F401 except ImportError: subprocess.run( [ sys.executable, "-m", "pip", "install", "--no-cache-dir", "--no-build-isolation", "git+https://github.com/WildChlamydia/MiVOLO.git", ], check=True, ) import cv2 import torch import numpy as np import torch.nn as nn from PIL import Image from torchvision import transforms from huggingface_hub import hf_hub_download from ultralytics import YOLO from transformers import AutoModelForImageClassification, AutoConfig, AutoImageProcessor import gradio as gr import spaces import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt # --------------------------------------------------------------------------- # PyTorch 2.6+ defaults torch.load(weights_only=True), which blocks unpickling # full ultralytics model objects baked into older YOLO .pt checkpoints. # Older YOLO checkpoints pickle many internal ultralytics/torch classes, so # allowlisting them one at a time is fragile. Since these checkpoints come # from a known, trusted HF repo (arnabdhar/YOLOv8-Face-Detection), we patch # torch.load to default weights_only=False, restoring pre-2.6 behavior. # --------------------------------------------------------------------------- _original_torch_load = torch.load def _patched_torch_load(*args, **kwargs): kwargs.setdefault("weights_only", False) return _original_torch_load(*args, **kwargs) torch.load = _patched_torch_load # Models load on CPU here. They only move to the GPU inside functions # decorated with @spaces.GPU, since ZeroGPU grants GPU access per-call, # not as a permanently attached device. dtype = torch.float16 # --------------------------------------------------------------------------- # Pipeline A setup — face detection + gender classification (moderate crowds) # --------------------------------------------------------------------------- face_model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt") face_model = YOLO(face_model_path) mivolo_config = AutoConfig.from_pretrained("iitolstykh/mivolo_v2", trust_remote_code=True) mivolo_model = AutoModelForImageClassification.from_pretrained( "iitolstykh/mivolo_v2", trust_remote_code=True, torch_dtype=dtype ) mivolo_processor = AutoImageProcessor.from_pretrained("iitolstykh/mivolo_v2", trust_remote_code=True) id2label = mivolo_config.gender_id2label def classify_faces_mivolo(img_bgr, boxes, device, confidence_threshold=0.65): h, w = img_bgr.shape[:2] gender_counts = {"Man": 0, "Woman": 0, "Uncertain": 0} per_head_results = [] for box in boxes: x1, y1, x2, y2 = map(int, box) bw, bh = x2 - x1, y2 - y1 pad_x, pad_y = int(bw * 0.4), int(bh * 0.4) x1p, y1p = max(0, x1 - pad_x), max(0, y1 - pad_y) x2p, y2p = min(w, x2 + pad_x), min(h, y2 + pad_y) crop = img_bgr[y1p:y2p, x1p:x2p] try: face_input = mivolo_processor(images=[crop])["pixel_values"].to(dtype=dtype, device=device) body_input = mivolo_processor(images=[None])["pixel_values"].to(dtype=dtype, device=device) with torch.no_grad(): output = mivolo_model(faces_input=face_input, body_input=body_input) gender_idx = output.gender_class_idx[0].item() gender_prob = output.gender_probs[0].item() label_raw = id2label[gender_idx].lower() label = "Man" if label_raw.startswith("m") else "Woman" if gender_prob < confidence_threshold: label = "Uncertain" except Exception: label = "Uncertain" gender_counts[label] += 1 per_head_results.append({"box": (x1, y1, x2, y2), "label": label}) return gender_counts, per_head_results # --------------------------------------------------------------------------- # Shadow crowd reconciliation — ironclad subtraction rule, no extrapolation. # CSRNet's density regression is treated as ground truth for how many people # are actually in the frame. Anyone the face pipeline (YOLOv8-Face + MiVOLO) # didn't pick up (blurry, turned away, buried in shadow) is bucketed into # "Unrecognized / Shadow Crowd" instead of silently vanishing from the count. # --------------------------------------------------------------------------- def reconcile_crowd_counts(gender_counts, total_faces_detected, csrnet_total_count): men = gender_counts["Man"] women = gender_counts["Woman"] uncertain = gender_counts["Uncertain"] total_count = round(csrnet_total_count) # Unrecognized / Shadow Crowd = CSRNet Total Density Count - Total Faces Detected # Clamped at 0 so a noisy/underestimating density map can never go negative. shadow_crowd = max(0, total_count - total_faces_detected) return { "Total Count": total_count, "Men": men, "Women": women, "Uncertain": uncertain, "Unrecognized / Shadow Crowd": shadow_crowd, } # --------------------------------------------------------------------------- # Overlay — yellow = male, pink = female. Unrecognizable faces (Uncertain) are # deliberately left out of the color map, so no dot at all is drawn for them. # --------------------------------------------------------------------------- def draw_gender_overlay(img_bgr, per_head_results, total_count): overlay = img_bgr.copy() color_map = { "Man": (0, 255, 255), # yellow (BGR) "Woman": (180, 105, 255), # pink (BGR) } for head in per_head_results: label = head["label"] if label not in color_map: continue # unrecognizable gender -> no marker, no color x1, y1, x2, y2 = head["box"] cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 color = color_map[label] cv2.circle(overlay, (cx, cy), radius=6, color=color, thickness=-1) cv2.circle(overlay, (cx, cy), radius=6, color=(0, 0, 0), thickness=1) footer_height = 50 canvas = np.full((overlay.shape[0] + footer_height, overlay.shape[1], 3), 255, dtype=np.uint8) canvas[: overlay.shape[0], :, :] = overlay text = f"Total count: {total_count}" cv2.putText(canvas, text, (10, overlay.shape[0] + 35), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2) return canvas # --------------------------------------------------------------------------- # Pipeline B setup — CSRNet density estimation (extreme-density crowds, and # now also used inside Standard mode as the ground-truth total for shadow # crowd reconciliation) # --------------------------------------------------------------------------- def make_layers(cfg, in_channels=3, dilation=False): d_rate = 2 if dilation else 1 layers = [] for v in cfg: if v == "M": layers += [nn.MaxPool2d(kernel_size=2, stride=2)] else: layers += [nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate, dilation=d_rate), nn.ReLU(inplace=True)] in_channels = v return nn.Sequential(*layers) class CSRNet(nn.Module): def __init__(self): super(CSRNet, self).__init__() self.frontend = make_layers([64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512]) self.backend = make_layers([512, 512, 512, 256, 128, 64], in_channels=512, dilation=True) self.output_layer = nn.Conv2d(64, 1, kernel_size=1) def forward(self, x): x = self.frontend(x) x = self.backend(x) x = self.output_layer(x) return x csrnet_transform = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) CSRNET_WEIGHTS_PATH = "weights.pth" if not os.path.exists(CSRNET_WEIGHTS_PATH): csrnet_weights_file = hf_hub_download(repo_id="rootstrap-org/crowd-counting", filename="weights.pth") CSRNET_WEIGHTS_PATH = csrnet_weights_file csrnet_model = CSRNet() csrnet_model.load_state_dict(torch.load(CSRNET_WEIGHTS_PATH, map_location="cpu")) csrnet_model.eval() def run_csrnet(pil_img, device="cpu"): input_tensor = csrnet_transform(pil_img).unsqueeze(0).to(device) with torch.no_grad(): output = csrnet_model(input_tensor) density_map = output.squeeze().cpu().numpy() total_count = float(density_map.sum()) return total_count, density_map # --------------------------------------------------------------------------- # Main entry point used by the UI # --------------------------------------------------------------------------- @spaces.GPU def process_image(pil_img, mode): if pil_img is None: return None, "Please upload an image first." device = "cuda" if torch.cuda.is_available() else "cpu" pil_img = pil_img.convert("RGB") img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) if mode == "Standard (count + gender)": face_model.to(device) mivolo_model.to(device) csrnet_model.to(device) results = face_model(img_bgr, conf=0.25, imgsz=1280) boxes = results[0].boxes.xyxy.cpu().numpy() total_faces_detected = len(boxes) gender_counts, per_head_results = classify_faces_mivolo(img_bgr, boxes, device) csrnet_total_count, _ = run_csrnet(pil_img, device) final_counts = reconcile_crowd_counts(gender_counts, total_faces_detected, csrnet_total_count) overlay = draw_gender_overlay(img_bgr, per_head_results, final_counts["Total Count"]) overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB) unrecognizable_total = final_counts["Uncertain"] + final_counts["Unrecognized / Shadow Crowd"] summary = ( f"Total count: {final_counts['Total Count']}\n" f"Male: {final_counts['Men']}\n" f"Female: {final_counts['Women']}\n" f"Unrecognizable: {unrecognizable_total}\n\n" f"Yellow = Male, Pink = Female (unrecognizable faces are left unmarked)" ) return Image.fromarray(overlay_rgb), summary else: # Extreme Density mode csrnet_model.to(device) total_count, density_map = run_csrnet(pil_img, device) total_count = round(total_count) fig, ax = plt.subplots(figsize=(6, 6)) ax.imshow(density_map, cmap="jet") ax.set_title(f"Estimated count: {total_count}") ax.axis("off") buf = io.BytesIO() fig.savefig(buf, format="png", bbox_inches="tight") plt.close(fig) buf.seek(0) heatmap_img = Image.open(buf) summary = ( f"Estimated total count: {total_count}\n\n" f"Density-based estimate (CSRNet). Gender classification is not " f"performed in this mode — individual faces are too small/occluded " f"at this density to classify reliably." ) return heatmap_img, summary with gr.Blocks(title="Crowd Counting & Gender Classification") as demo: gr.Markdown("# Crowd Counting & Gender Classification") gr.Markdown( "Upload a crowd photo and choose a mode:\n\n" "- **Standard** — total count (reconciled against CSRNet's density estimate) + " "male/female breakdown, with a color-coded dot overlay (yellow = Male, pink = " "Female, unrecognizable faces left unmarked)\n" "- **Extreme Density** — total headcount only, via density estimation, for scenes " "too dense/occluded for reliable per-face gender classification" ) with gr.Row(): with gr.Column(): image_input = gr.Image(type="pil", label="Upload crowd photo") mode_input = gr.Radio( ["Standard (count + gender)", "Extreme Density (count only)"], value="Standard (count + gender)", label="Mode", ) submit_btn = gr.Button("Run", variant="primary") with gr.Column(): image_output = gr.Image(type="pil", label="Result") text_output = gr.Textbox(label="Summary", lines=7) submit_btn.click(fn=process_image, inputs=[image_input, mode_input], outputs=[image_output, text_output]) demo.launch()