TrueFrame / image_detect.py
Gaurav-Mhatre's picture
Update image_detect.py
a4f3cac verified
import torch
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
class ImageDeepfakeDetector:
def __init__(self):
print("⚡ Loading Image AI Model...")
try:
# Both model and processor now pull directly from Hugging Face!
self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
print("✅ Image Model Ready.")
except Exception as e:
print(f"❌ Error: {e}")
exit()
def predict(self, image_path):
try:
image = Image.open(image_path)
# --- THE "PHYSICS" FIX ---
# We are now detecting "Camera Physics" vs "Generative Patterns"
labels = [
# === REAL (Camera Physics) ===
# Real cameras leave specific traces: noise, focus fall-off, organic skin.
"photo taken with a canon or nikon dslr camera", # Specific camera brand hints
"real human skin with natural pores and imperfections", # Texture check
"authentic photograph with natural lighting shadows", # Lighting physics
"candid photo from a smartphone camera", # Lower quality real
"high resolution raw photography", # HD Real
# === FAKE (Generative Patterns) ===
# AI generators leave specific traces: overly smooth, distorted backgrounds.
"ai generated image by midjourney", # Specific Generator
"stable diffusion synthetic image", # Specific Generator
"deepfake face swap with artifacts", # Manipulation
"computer generated 3d character render", # CGI look
"perfectly smooth plastic skin texture", # The "AI Glow"
"unnatural distorted background details" # Background check
]
inputs = self.processor(
text=labels,
images=image,
return_tensors="pt",
padding=True
)
with torch.no_grad():
outputs = self.model(**inputs)
probs = outputs.logits_per_image.softmax(dim=1)
scores = probs.tolist()[0]
# --- SCORING ---
# Real = Indices 0 to 4 (5 labels)
real_score = sum(scores[:5])
# Fake = Indices 5 to 10 (6 labels)
fake_score = sum(scores[5:])
# --- "UNCERTAINTY" LOGIC ---
# If the scores are very close (e.g., 49% Real vs 51% Fake),
# it means the AI is guessing. In court, "Innocent until proven guilty".
# So, if the gap is small (< 5%), we assume it's REAL.
score_diff = abs(real_score - fake_score)
if fake_score > real_score:
# It thinks it's fake, but is it sure?
if score_diff < 0.05:
# Not sure enough -> Call it Real (but low confidence)
return "REAL", real_score
else:
return "DEEPFAKE DETECTED", fake_score
else:
return "REAL", real_score
except Exception as e:
print(f"Error predicting {image_path}: {e}")
return "ERROR", 0.0