import os import time import numpy as np import cv2 import base64 from flask import Flask, request, jsonify, render_template from PIL import Image from ultralytics import YOLO import torch import torch.nn as nn import torchvision.transforms as transforms import segmentation_models_pytorch as smp from torchvision import models app = Flask(__name__) # --- Configuration --- MODEL_PATHS = { "density": "models/best_density_efficientnet_gan.pth", "regression": "models/20251125_214552_regression_resnet18.onnx", "gate": "models/gate_resnet18.pth", "yolos": "models/yolov8s.pt", "yolox": "models/yolov8x.pt" } device = "cpu" # Spaces use CPU by default # --- Load Density Model (PyTorch) --- def get_density_model(): return smp.Unet(encoder_name="efficientnet-b0", encoder_weights=None, in_channels=3, classes=1, activation=None) print("Loading models...") try: density_model = get_density_model() state_dict = torch.load(MODEL_PATHS["density"], map_location=torch.device('cpu')) density_model.load_state_dict(state_dict) density_model.to(device) density_model.eval() print("Density model loaded.") except Exception as e: print(f"Error loading density: {e}"); density_model = None # --- Load Gatekeeper Model (PyTorch) --- try: gate_model = models.resnet18() # FIX: Must match training architecture EXACTLY (including Dropout) gate_model.fc = nn.Sequential( nn.Dropout(0.5), nn.Linear(512, 1), nn.Sigmoid() ) gate_model.load_state_dict(torch.load(MODEL_PATHS["gate"], map_location=torch.device('cpu'))) gate_model.to(device) gate_model.eval() print("Gatekeeper model loaded.") except Exception as e: print(f"Error loading gatekeeper: {e}"); gate_model = None # --- Load Regression Model (ONNX) --- try: regression_net = cv2.dnn.readNetFromONNX(MODEL_PATHS["regression"]) print("Regression model loaded.") except Exception as e: print(f"Error loading regression: {e}"); regression_net = None # --- Load YOLO Models --- yolo_small = YOLO(MODEL_PATHS["yolos"]) yolo_large = YOLO(MODEL_PATHS["yolox"]) print("All models loaded.") # --- Transforms --- norm_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) gate_transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) reg_transform = transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # --- Inference Functions --- def run_density_inference(image_pil): if density_model is None: return 0, 0, "", 0.0 start = time.time() image_np = np.array(image_pil) h, w = image_np.shape[:2] pad_h, pad_w = (32 - h % 32) % 32, (32 - w % 32) % 32 if pad_h or pad_w: image_np = cv2.copyMakeBorder(image_np, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0) img_tensor = norm_transform(Image.fromarray(image_np)).unsqueeze(0).to(device) with torch.no_grad(): output = density_model(img_tensor) raw_map = output[0, 0, :h, :w].cpu().numpy() count = np.sum(raw_map) / 100.0 heatmap_norm = cv2.normalize(raw_map, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U) heatmap_color = cv2.applyColorMap(heatmap_norm, cv2.COLORMAP_JET) _, buffer = cv2.imencode('.jpg', heatmap_color) return int(round(count)), round((time.time()-start)*1000, 1), base64.b64encode(buffer).decode('utf-8') def run_regression_inference(image_pil): if regression_net is None: return 0, 0 start = time.time() img_tensor = reg_transform(image_pil).unsqueeze(0).numpy() regression_net.setInput(img_tensor) count = regression_net.forward()[0][0] return int(max(0, round(count))), round((time.time()-start)*1000, 1) def run_yolo_inference(model, image_pil): start = time.time() img_cv2 = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR) results = model(image_pil, classes=[0], conf=0.15, iou=0.3, verbose=False) count = len(results[0].boxes) mask = np.zeros_like(img_cv2, dtype=np.float32) for box in results[0].boxes: x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() cx, cy = int((x1+x2)/2), int((y1+y2)/2) radius = int(max(int(x2-x1), int(y2-y1))/1.5) cv2.circle(mask, (cx, cy), radius, (0,0,255), -1) mask = cv2.GaussianBlur(mask, (99, 99), 0) mask = mask / (mask.max() + 1e-6) overlay = np.clip(img_cv2.astype(np.float32) + (mask * 255 * 0.6), 0, 255).astype(np.uint8) _, buffer = cv2.imencode('.jpg', overlay) return count, round((time.time()-start)*1000, 1), base64.b64encode(buffer).decode('utf-8') # --- Logic Chain --- def get_best_estimate(density_count, yolo_count, gate_prob): if gate_prob > 0.6: return density_count, "density", f"Scene Classified as Dense (Confidence: {gate_prob:.2f})" elif gate_prob < 0.4: return yolo_count, "yolox", f"Scene Classified as Sparse (Confidence: {1-gate_prob:.2f})" else: diff_ratio = density_count / (yolo_count + 1e-6) if diff_ratio > 2.0: return density_count, "density", "Uncertain Scene, but Density Map detected heavy crowding." elif diff_ratio < 0.5: return yolo_count, "yolox", "Uncertain Scene, but YOLO detected clear objects." else: avg = int((density_count + yolo_count) / 2) return avg, "hybrid", f"Scene Ambiguous (Conf: {gate_prob:.2f}). Using Ensemble Average." @app.route('/') def index(): return render_template('index.html') @app.route('/predict', methods=['POST']) def predict(): if 'file' not in request.files: return jsonify({'error': 'No file'}), 400 file = request.files['file'] try: image = Image.open(file.stream).convert('RGB') t_start = time.time() # 1. Run Gatekeeper First gate_prob = 0.5 # Default to uncertain if model fails if gate_model is not None: gate_input = gate_transform(image).unsqueeze(0).to(device) with torch.no_grad(): gate_prob = gate_model(gate_input).item() else: print("Warning: Gatekeeper model is None. Using default prob 0.5") # 2. Run All Models d_c, d_t, d_map = run_density_inference(image) r_c, r_t = run_regression_inference(image) ys_c, ys_t, ys_map = run_yolo_inference(yolo_small, image) yx_c, yx_t, yx_map = run_yolo_inference(yolo_large, image) # 3. Decision Logic best_count, best_model, reason = get_best_estimate(d_c, yx_c, gate_prob) return jsonify({ 'density': {'count': d_c, 'time': d_t, 'heatmap': d_map}, 'regression': {'count': r_c, 'time': r_t}, 'yolos': {'count': ys_c, 'time': ys_t, 'heatmap': ys_map}, 'yolox': {'count': yx_c, 'time': yx_t, 'heatmap': yx_map}, 'recommendation': { 'count': best_count, 'model': best_model, 'reason': reason }, 'gate_prob': gate_prob, 'total_time': round((time.time()-t_start)*1000, 1) }) except Exception as e: import traceback; traceback.print_exc() return jsonify({'error': str(e)}), 500 if __name__ == '__main__': os.makedirs('templates', exist_ok=True) app.run(host='0.0.0.0', port=7860, debug=False)