add generator type detection

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 import torch
+import base64
+import io
 from PIL import Image
-from src.models.inference import load_model, predict
+from src.models.inference import load_model, predict, predict_generator, load_generator_model
 from src.services.metadata_checker import get_metadata
 from src.services.gradcam import generate_gradcam
 import tempfile
@@ -9,8 +11,12 @@ import os
 
 # Load model once
 model = load_model()
+generator_model = load_generator_model()
 
 def analyze_image(image):
+    if image is None:
+        return "Please upload an image.", None
+
     # Save PIL image to temp file
     with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
         image.save(tmp.name)
@@ -21,17 +27,25 @@ def analyze_image(image):
         metadata = get_metadata(tmp_path)
         cam_array = generate_gradcam(tmp_path, model=model)
         cam_image = Image.fromarray(cam_array)
+        generator_result = predict_generator(tmp_path, model=generator_model)
     finally:
         os.remove(tmp_path)
 
     label = prediction["label"]
     confidence = prediction["confidence"]
-    raw_score = prediction["raw_score"]
 
     result_text = f"**{label}** — Confidence: {confidence}%\n\n"
-    result_text += f"Raw Score: {raw_score}\n\n"
     result_text += f"⚠️ This is a model-based estimate, not definitive proof.\n\n"
-    result_text += f"**Metadata:**\n"
+
+    # Generator type
+    gen_type = generator_result["generator_type"]
+    gen_conf = generator_result["confidence"]
+    result_text += f"**Generator Type:** {gen_type} ({gen_conf}%)\n\n"
+    result_text += "**Class Probabilities:**\n"
+    for cls, prob in generator_result["class_probabilities"].items():
+        result_text += f"- {cls}: {prob}%\n"
+
+    result_text += f"\n**Metadata:**\n"
     result_text += f"- Format: {metadata['format']}\n"
     result_text += f"- Dimensions: {metadata['dimensions']}\n"
     result_text += f"- File Size: {metadata['file_size_kb']} KB\n"

saved_models/generator_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5bda82f8d7eb4b25248e53aae79dedf7b344440fea165845cf04fded0d692770
+size 45314775

src/data/generator_loader.py

@@ -0,0 +1,122 @@
+import os
+import pandas as pd
+import random
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader, random_split
+from src.data.transforms import train_transforms, val_transforms
+
+DATASET_ROOT = "/Users/siemoncha/Desktop/CS/datasets/artifact-dataset"
+
+# Manual class mapping based on architecture knowledge
+SOURCE_CLASS_MAP = {
+    # Class 0 - Real
+    "coco": 0,
+    "ffhq": 0,
+    "lsun": 0,
+    "imagenet": 0,
+    "landscape": 0,
+    "afhq": 0,
+    "celebahq": 0,
+    "metfaces": 0,
+
+    # Class 1 - GAN
+    "stylegan1": 1,
+    "stylegan2": 1,
+    "stylegan3": 1,
+    "pro_gan": 1,
+    "big_gan": 1,
+    "star_gan": 1,
+    "cycle_gan": 1,
+    "gansformer": 1,
+    "generative_inpainting": 1,
+    "lama": 1,
+    "mat": 1,
+    "sfhq": 1,
+    "cips": 1,
+    "projected_gan": 1,
+    "gau_gan": 1,
+
+    # Class 2 - Diffusion
+    "stable_diffusion": 2,
+    "ddpm": 2,
+    "glide": 2,
+    "latent_diffusion": 2,
+    "vq_diffusion": 2,
+    "denoising_diffusion_gan": 2,
+    "diffusion_gan": 2,
+    "palette": 2,
+
+    # Class 3 - Other
+    "taming_transformer": 3,
+    "face_synthetics": 3,
+}
+
+CLASS_NAMES = {0: "Real", 1: "GAN", 2: "Diffusion", 3: "Other"}
+MAX_PER_CLASS = 10000
+
+
+class GeneratorDataset(Dataset):
+    def __init__(self, samples, transform=None):
+        self.samples = samples
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        img_path, label = self.samples[idx]
+        image = Image.open(img_path).convert("RGB")
+        if self.transform:
+            image = self.transform(image)
+        return image, label
+
+
+def load_generator_samples():
+    class_samples = {0: [], 1: [], 2: [], 3: []}
+
+    for source, cls in SOURCE_CLASS_MAP.items():
+        csv_path = os.path.join(DATASET_ROOT, source, "metadata.csv")
+        if not os.path.exists(csv_path):
+            print(f"Skipping {source} - no metadata.csv")
+            continue
+        df = pd.read_csv(csv_path)
+        for _, row in df.iterrows():
+            img_path = os.path.join(DATASET_ROOT, source, row["image_path"])
+            class_samples[cls].append((img_path, cls))
+
+    # Balance classes
+    for cls in class_samples:
+        class_samples[cls] = class_samples[cls][:MAX_PER_CLASS]
+        print(f"Class {cls} ({CLASS_NAMES[cls]}): {len(class_samples[cls])} samples")
+
+    all_samples = []
+    for cls in class_samples:
+        all_samples.extend(class_samples[cls])
+
+    print(f"Total: {len(all_samples)}")
+    return all_samples
+
+
+def get_generator_dataloaders(batch_size=32):
+    all_samples = load_generator_samples()
+
+    # Shuffle before splitting
+    random.shuffle(all_samples)
+
+    train_size = int(0.75 * len(all_samples))
+    val_size = int(0.125 * len(all_samples))
+    test_size = len(all_samples) - train_size - val_size
+
+    train_samples = all_samples[:train_size]
+    val_samples = all_samples[train_size:train_size + val_size]
+    test_samples = all_samples[train_size + val_size:]
+
+    train_set = GeneratorDataset(train_samples, transform=train_transforms)
+    val_set = GeneratorDataset(val_samples, transform=val_transforms)
+    test_set = GeneratorDataset(test_samples, transform=val_transforms)
+
+    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2)
+    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2)
+    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2)
+
+    return train_loader, val_loader, test_loader

src/models/inference.py

@@ -2,6 +2,7 @@ import torch
 from PIL import Image
 from src.models.model import build_model
 from src.data.transforms import val_transforms
+from src.data.generator_loader import CLASS_NAMES
 
 MODEL_PATH = "saved_models/best_model.pth"
 
@@ -31,6 +32,38 @@ def predict(image_path: str, model=None):
         "raw_score": round(prob, 4)
     }
 
+GENERATOR_MODEL_PATH = "saved_models/generator_model.pth"
+
+def load_generator_model(model_path=GENERATOR_MODEL_PATH):
+    from src.models.train_generator import build_multiclass_model
+    model = build_multiclass_model(num_classes=4, pretrained=False)
+    model.load_state_dict(torch.load(model_path, map_location="cpu"))
+    model.eval()
+    return model
+
+def predict_generator(image_path: str, model=None):
+    if model is None:
+        model = load_generator_model()
+
+    image = Image.open(image_path).convert("RGB")
+    tensor = val_transforms(image).unsqueeze(0)
+
+    with torch.no_grad():
+        output = model(tensor)
+        probs = torch.softmax(output, dim=1)[0]
+        pred_class = probs.argmax().item()
+        confidence = probs[pred_class].item()
+
+    return {
+        "generator_type": CLASS_NAMES[pred_class],
+        "confidence": round(confidence * 100, 2),
+        "class_probabilities": {
+            CLASS_NAMES[i]: round(probs[i].item() * 100, 2)
+            for i in range(4)
+        }
+    }
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) < 2:

src/models/train_generator.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torchvision import models
+from sklearn.metrics import classification_report, confusion_matrix
+from src.data.generator_loader import get_generator_dataloaders, CLASS_NAMES
+
+
+def build_multiclass_model(num_classes=4, pretrained=True):
+    model = models.resnet18(weights="IMAGENET1K_V1" if pretrained else None)
+
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # Unfreeze layer4
+    for name, param in model.named_parameters():
+        if "layer4" in name or "fc" in name:
+            param.requires_grad = True
+
+    in_features = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Dropout(0.5),
+        nn.Linear(in_features, 256),
+        nn.ReLU(),
+        nn.Dropout(0.3),
+        nn.Linear(256, num_classes)
+    )
+    return model
+
+
+def evaluate(model, loader, device):
+    model.eval()
+    all_preds, all_labels = [], []
+
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            labels = labels.to(device)
+            outputs = model(images)
+            preds = outputs.argmax(dim=1)
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+
+    return all_preds, all_labels
+
+
+def train(epochs=15, batch_size=32, lr=1e-4):
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    train_loader, val_loader, test_loader = get_generator_dataloaders(batch_size=batch_size)
+    model = build_multiclass_model().to(device)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
+    scheduler = ReduceLROnPlateau(optimizer, patience=2)
+
+    best_val_acc = 0
+    early_stop_patience = 3
+    no_improve_count = 0
+
+    for epoch in range(epochs):
+        model.train()
+        train_loss, correct, total = 0, 0, 0
+
+        for images, labels in train_loader:
+            images = images.to(device)
+            labels = labels.to(device)
+
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+            preds = outputs.argmax(dim=1)
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+
+        train_acc = correct / total
+        avg_train_loss = train_loss / len(train_loader)
+
+        val_preds, val_labels = evaluate(model, val_loader, device)
+        val_acc = sum(p == l for p, l in zip(val_preds, val_labels)) / len(val_labels)
+        scheduler.step(1 - val_acc)
+
+        print(f"Epoch {epoch+1}/{epochs} | "
+              f"Train Loss: {avg_train_loss:.4f} | Train Acc: {train_acc:.4f} | "
+              f"Val Acc: {val_acc:.4f}")
+
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            no_improve_count = 0
+            torch.save(model.state_dict(), "saved_models/generator_model.pth")
+            print(f"  -> Best model saved")
+        else:
+            no_improve_count += 1
+            if no_improve_count >= early_stop_patience:
+                print(f"Early stopping at epoch {epoch+1}")
+                break
+
+    # Final evaluation
+    print("\n--- Final Evaluation ---")
+    test_preds, test_labels = evaluate(model, test_loader, device)
+    test_acc = sum(p == l for p, l in zip(test_preds, test_labels)) / len(test_labels)
+    print(f"Test Accuracy: {test_acc:.4f}")
+    print("\nClassification Report:")
+    print(classification_report(test_labels, test_preds,
+                                target_names=list(CLASS_NAMES.values())))
+    print("Confusion Matrix:")
+    print(confusion_matrix(test_labels, test_preds))
+
+
+if __name__ == "__main__":
+    train()