app.py

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()