Spaces:
Sleeping
Sleeping
File size: 1,696 Bytes
d246850 |
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 |
import torch
import timm
from torchvision import transforms
from PIL import Image
import requests
from io import BytesIO
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# 🔹 Preprocessing (same as training)
transform = transforms.Compose([
transforms.Resize((300, 300)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
def load_model(model_path: str):
"""Load EfficientNet-B3 model from file."""
if model_path.endswith(".pt"):
model = torch.load(model_path, map_location=DEVICE, weights_only=False)
else:
model = timm.create_model("efficientnet_b3", pretrained=False, num_classes=1)
model.load_state_dict(torch.load(model_path, map_location=DEVICE))
model.to(DEVICE)
model.eval()
return model
def load_image(input_source):
"""Load and preprocess image from file upload or URL."""
if isinstance(input_source, str): # URL
response = requests.get(input_source)
img = Image.open(BytesIO(response.content)).convert("RGB")
else: # Uploaded file
img = Image.open(input_source).convert("RGB")
return img
def preprocess_image(img):
"""Apply transform and return tensor."""
return transform(img).unsqueeze(0)
def predict(model, img_tensor):
"""Predict class and confidence score."""
with torch.no_grad():
outputs = model(img_tensor.to(DEVICE))
probs = torch.sigmoid(outputs)
prob = probs.item()
print(prob)
label = "🧠 AI-generated" if prob >= 0.00001 else "📸 Real"
return label, prob
|