Spaces:
Runtime error
Runtime error
File size: 4,765 Bytes
879d39c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import torch
import torch.nn as nn
from torchvision import transforms, models
import open_clip
from PIL import Image, ImageFilter
import numpy as np
import os
# --- 1. SETUP & DEVICE ---
DEVICE = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")
# --- 2. LOAD MODELS ---
# A. Load openclip (ViT-L-14)
print("Loading openclip...")
openclip_model, _, openclip_preprocess = open_clip.create_model_and_transforms(
'ViT-L-14', pretrained='datacomp_xl_s13b_b90k'
)
openclip_model.to(DEVICE)
# Define your openclip Forensic Head Architecture (matches your training)
class openclipHead(nn.Module):
def __init__(self, input_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, 1)
)
def forward(self, x): return self.net(x)
# Load openclip Weights
openclip_head = openclipHead(input_dim=768).to(DEVICE)
openclip_head.load_state_dict(torch.load('models/openclip_forensic_head.pth', map_location=DEVICE))
openclip_head.eval()
# B. Load ConvNeXt-Base
print("Loading ConvNeXt...")
cn_backbone = models.convnext_base(weights=None) # Architecture only
cn_backbone.to(DEVICE)
cn_backbone.eval()
class ConvNextHead(nn.Module):
def __init__(self, input_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, 1)
)
def forward(self, x): return self.net(x)
cn_head = ConvNextHead(input_dim=1024).to(DEVICE)
cn_head.load_state_dict(torch.load('models/convnext_forensic_head.pth', map_location=DEVICE))
cn_head.eval()
# ConvNext Preprocessing (Standard ImageNet)
cn_preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# --- 3. FEATURE EXTRACTION (Heuristics) ---
def extract_simple_features(image_path):
img = Image.open(image_path).convert('RGB')
img_array = np.array(img) / 255.0
edges = np.abs(np.diff(np.mean(img_array, axis=2), axis=0)).mean() + \
np.abs(np.diff(np.mean(img_array, axis=2), axis=1)).mean()
img_smooth = np.array(img.filter(ImageFilter.GaussianBlur(2))) / 255.0
noise = np.mean((img_array - img_smooth) ** 2) * 1000
return {
'noise_level': noise,
'edge_density': edges,
'is_too_clean': (noise < 0.05 and edges < 0.12) # Adjusted thresholds
}
# --- 4. THE ENSEMBLE INFERENCE ---
def run_ensemble(image_path):
img = Image.open(image_path).convert('RGB')
# openclip Score
img_openclip = openclip_preprocess(img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
sig_feat = openclip_model.encode_image(img_openclip)
sig_feat /= sig_feat.norm(dim=-1, keepdim=True)
sig_logit = openclip_head(sig_feat)
prob_openclip = torch.sigmoid(sig_logit).item()
# ConvNeXt Score
img_cn = cn_preprocess(img).unsqueeze(0).to(DEVICE)
with torch.no_grad():
feat = cn_backbone.features(img_cn)
feat = cn_backbone.avgpool(feat)
feat = torch.flatten(feat, 1)
cn_logit = cn_head(feat)
prob_cn = torch.sigmoid(cn_logit).item()
# Average the two for the "Raw Ensemble Score"
raw_ensemble_score = (prob_openclip + prob_cn) / 2
# Calibration
features = extract_simple_features(image_path)
if features['is_too_clean']:
calibrated_score = raw_ensemble_score * 0.55 # 45% discount for product shots
reason = "Clean product-shot detected. Reducing probability."
else:
calibrated_score = raw_ensemble_score
reason = "Standard analysis applied."
return {
'openclip_score': prob_openclip,
'convnext_score': prob_cn,
'raw_ensemble': raw_ensemble_score,
'calibrated': min(calibrated_score, 0.95),
'reason': reason,
'features': features
}
# --- 5. TEST IT ---
test_image = "/Users/rishitbaitule/Downloads/b.jpg" # Update this path!
if os.path.exists(test_image):
results = run_ensemble(test_image)
print("-" * 30)
print(f"Individual openclip: {results['openclip_score']:.2%}")
print(f"Individual ConvNeXt: {results['convnext_score']:.2%}")
print("-" * 30)
print(f"ENSEMBLE RAW SCORE: {results['raw_ensemble']:.2%}")
print(f"CALIBRATED SCORE: {results['calibrated']:.2%}")
print(f"REASON: {results['reason']}")
print("-" * 30)
else:
print("Image not found. Please check test_image path.") |