Update app.py

app.py

@@ -1,96 +1,107 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-import torchvision.transforms as T
+import torchvision.models as models
+import torchvision.transforms as transforms
 from PIL import Image
 from huggingface_hub import hf_hub_download
 
-#####################################
-# 1) Define the same custom class
-#####################################
+########################################
+# 1. Define Model Architecture
+########################################
 class MultiTaskModel(nn.Module):
     def __init__(self, backbone, feature_dim, num_obj_classes):
         super(MultiTaskModel, self).__init__()
         self.backbone = backbone
+        # Object recognition head
         self.obj_head = nn.Linear(feature_dim, num_obj_classes)
+        # Binary classification head (0: AI-generated, 1: Real)
         self.bin_head = nn.Linear(feature_dim, 2)
-
+    
     def forward(self, x):
         feats = self.backbone(x)
         obj_logits = self.obj_head(feats)
         bin_logits = self.bin_head(feats)
         return obj_logits, bin_logits
 
-#####################################
-# 2) Allowlist the class
-#####################################
-import torch.serialization
-torch.serialization.add_safe_globals([MultiTaskModel])
+########################################
+# 2. Reconstruct the Model and Load Weights
+########################################
+# Set the number of object classes (update this to match your training)
+num_obj_classes = 139  # Example; change as needed
+device = torch.device("cpu")
 
-#####################################
-# 3) Download & Load the full model
-#####################################
-repo_id = "Abdu07/multitask-model"  # or your actual repo
-filename = "multitask_model.pth"    # the file you uploaded
-model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+# Instantiate the backbone (ResNet-50 without its final layer)
+resnet = models.resnet50(pretrained=False)
+resnet.fc = nn.Identity()
+feature_dim = 2048
+
+# Build the model architecture
+model = MultiTaskModel(resnet, feature_dim, num_obj_classes)
+model.to(device)
 
-# Force PyTorch to load the full model object
-model = torch.load(model_path, map_location="cpu")  # default weights_only=True, but we added safe_globals
+# Download the state dict from HF Hub
+repo_id = "Abdu07/multitask-model"
+filename = "best_model.pt"  # Make sure this is the state dict file
+model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+state_dict = torch.load(model_path, map_location="cpu")
+model.load_state_dict(state_dict)
 model.eval()
 
-#####################################
-# 4) Label Mappings
-#####################################
+########################################
+# 3. Define Label Mappings and Transforms
+########################################
+# Update these mappings with your actual training labels.
 idx_to_obj_label = {
-    # Fill in with your actual object label indices
     0: "cat",
     1: "dog",
     2: "car",
-    # ...
+    # ... add your object classes here ...
 }
 bin_label_names = ["AI-Generated", "Real"]
 
-#####################################
-# 5) Validation Transforms
-#####################################
-val_transforms = T.Compose([
-    T.Resize(256),
-    T.CenterCrop(224),
-    T.ToTensor(),
-    T.Normalize(mean=[0.485, 0.456, 0.406],
-                std=[0.229, 0.224, 0.225])
+# Define the validation transforms (same as used during training/validation)
+val_transforms = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
 ])
 
-#####################################
-# 6) Inference Function
-#####################################
+########################################
+# 4. Define the Inference Function
+########################################
 def predict_image(img: Image.Image) -> str:
+    """
+    Takes an uploaded PIL image, processes it, and returns the model's prediction.
+    """
+    # Ensure image is in RGB
     img = img.convert("RGB")
-    img_t = val_transforms(img).unsqueeze(0)
+    # Apply validation transforms
+    img_tensor = val_transforms(img).unsqueeze(0).to(device)  # Shape: [1, 3, 224, 224]
     with torch.no_grad():
-        obj_logits, bin_logits = model(img_t)
-        obj_pred = torch.argmax(obj_logits, dim=1).item()
-        bin_pred = torch.argmax(bin_logits, dim=1).item()
+        obj_logits, bin_logits = model(img_tensor)
+    obj_pred = torch.argmax(obj_logits, dim=1).item()
+    bin_pred = torch.argmax(bin_logits, dim=1).item()
     obj_name = idx_to_obj_label.get(obj_pred, "Unknown")
     bin_name = bin_label_names[bin_pred]
-    return f"Object: {obj_name} | Authenticity: {bin_name}"
+    return f"Prediction: {obj_name} ({bin_name})"
 
-#####################################
-# 7) Gradio UI
-#####################################
+########################################
+# 5. Create Gradio UI
+########################################
 demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
     title="Multi-Task Image Classifier",
     description=(
-        "Upload an image to get two predictions: "
-        "1) The primary object, 2) Whether the image is AI-generated or real."
+        "Upload an image to receive two predictions:\n"
+        "1) The primary object in the image,\n"
+        "2) Whether the image is AI-generated or Real."
     )
 )
 
-def main():
-    demo.launch(server_name="0.0.0.0", enable_queue=True)
-
 if __name__ == "__main__":
-    main()
+    demo.launch(server_name="0.0.0.0", share=True)