Update app.py

app.py

@@ -7,9 +7,9 @@ import numpy as np
 
 # Multiple specialized models for ensemble detection
 MODELS = [
-    "Ateeqq/ai-vs-human-image-detector",  # SigLIP-based
-    "umm-maybe/AI-image-detector",        # Vision Transformer
-    "Dafilab/ai-image-detector",          # EfficientNet-B4
+    "Ateeqq/ai-vs-human-image-detector",  # SigLIP-based - Best for DALL-E 3, Midjourney
+    "umm-maybe/AI-image-detector",        # Vision Transformer - Good for Stable Diffusion
+    "facebook/dinov2-small",              # Meta's DINOv2 - Excellent feature detector
 ]
 
 print("Loading models for ensemble detection...")
@@ -51,15 +51,31 @@ def predict(image):
         # Get predictions from all models
         for i, (processor, model) in enumerate(zip(processors_list, models_list)):
             try:
-                inputs = processor(images=image, return_tensors="pt").to(device)
-                
-                with torch.no_grad():
-                    outputs = model(**inputs)
-                    logits = outputs.logits
-                    probs = F.softmax(logits, dim=1)[0].cpu().numpy()
-                
-                real_prob = float(probs[0])
-                ai_prob = float(probs[1])
+                # Special handling for DINOv2 (feature extractor)
+                if i == 2:  # DINOv2
+                    from torchvision import transforms
+                    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])
+                    ])
+                    img_tensor = transform(image).unsqueeze(0).to(device)
+                    with torch.no_grad():
+                        features = model(img_tensor)
+                    # Use feature statistics for detection
+                    feature_mean = features.mean()
+                    feature_std = features.std()
+                    ai_prob = float((feature_std.cpu() / (feature_mean.cpu() + 1e-6)).clamp(0, 1))
+                    real_prob = 1.0 - ai_prob
+                else:
+                    inputs = processor(images=image, return_tensors="pt").to(device)
+                    with torch.no_grad():
+                        outputs = model(**inputs)
+                        logits = outputs.logits
+                        probs = F.softmax(logits, dim=1)[0].cpu().numpy()
+                    real_prob = float(probs[0])
+                    ai_prob = float(probs[1])
                 
                 all_real_probs.append(real_prob)
                 all_ai_probs.append(ai_prob)