ABasiit/DCGANDiscriminator

1# streamlit_app.py2import streamlit as st3from PIL import Image4import torch5import torch.nn as nn6import torchvision.transforms as transforms7import io8 9# === Define your trained model architecture ===10class DCGANDiscriminator(nn.Module):11    def __init__(self):12        super(DCGANDiscriminator, self).__init__()13        self.model = nn.Sequential(14            nn.Conv2d(3, 64, 4, 2, 1),15            nn.LeakyReLU(0.2, inplace=True),16 17            nn.Conv2d(64, 128, 4, 2, 1),18            nn.BatchNorm2d(128),19            nn.LeakyReLU(0.2, inplace=True),20 21            nn.Conv2d(128, 256, 4, 2, 1),22            nn.BatchNorm2d(256),23            nn.LeakyReLU(0.2, inplace=True),24 25            nn.Conv2d(256, 512, 4, 2, 1),26            nn.BatchNorm2d(512),27            nn.LeakyReLU(0.2, inplace=True),28 29            nn.Conv2d(512, 1, 8),30            nn.Sigmoid()31        )32 33    def forward(self, x):34        return self.model(x).view(-1, 1)35 36# === Load model ===37device = torch.device("cuda" if torch.cuda.is_available() else "cpu")38model = DCGANDiscriminator().to(device)39model.load_state_dict(torch.load("dcgan_discriminator.pth", map_location=device))40model.eval()41 42# === Define transform ===43transform = transforms.Compose([44    transforms.Resize(128),45    transforms.CenterCrop(128),46    transforms.ToTensor(),47    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))48])49 50# === Streamlit UI ===51st.title("Deepfake Detector")52st.write("Upload a face image and the model will predict whether it's REAL or FAKE.")53 54uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])55 56if uploaded_file is not None:57    image = Image.open(uploaded_file).convert('RGB')58    st.image(image, caption='Uploaded Image', use_column_width=True)59 60    input_tensor = transform(image).unsqueeze(0).to(device)61    with torch.no_grad():62        output = model(input_tensor)63        confidence = output.item()64        label = "REAL" if confidence > 0.5 else "FAKE"65 66    st.markdown(f"### Prediction: **{label}**")67    st.write(f"Confidence: {confidence:.4f}")68