Update app.py

app.py

@@ -24,21 +24,35 @@ class MultiTaskModel(nn.Module):
         return obj_logits, bin_logits
 
 ########################################
-# 2. Load the Label Mapping and Set num_obj_classes
+# 2. Reconstruct the Model and Load Weights
+########################################
+num_obj_classes = 494  # Make sure this matches your training
+device = torch.device("cpu")
+
+resnet = models.resnet50(pretrained=False)
+resnet.fc = nn.Identity()
+feature_dim = 2048
+model = MultiTaskModel(resnet, feature_dim, num_obj_classes)
+model.to(device)
+
+repo_id = "Abdu07/multitask-model"
+filename = "Yolloplusclassproject_weights.pth"
+weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
+state_dict = torch.load(weights_path, map_location="cpu")
+model.load_state_dict(state_dict)
+model.eval()
+
+########################################
+# 3. Load Label Mapping and Define Transforms
 ########################################
 # Load the saved mapping from JSON
 with open("obj_label_mapping.json", "r") as f:
     obj_label_to_idx = json.load(f)
-# Use the mapping as-is; do not override it.
-num_obj_classes = len(obj_label_to_idx)
 # Create the inverse mapping
 idx_to_obj_label = {v: k for k, v in obj_label_to_idx.items()}
 
 bin_label_names = ["AI-Generated", "Real"]
 
-########################################
-# 3. Define Validation Transforms
-########################################
 val_transforms = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
@@ -48,27 +62,7 @@ val_transforms = transforms.Compose([
 ])
 
 ########################################
-# 4. Reconstruct the Model and Load Weights
-########################################
-device = torch.device("cpu")
-
-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)
-
-repo_id = "Abdu07/multitask-model"
-filename = "DualSight.pth"  # Ensure this checkpoint is from training with the same num_obj_classes
-weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
-state_dict = torch.load(weights_path, map_location="cpu")
-model.load_state_dict(state_dict)
-model.eval()
-
-########################################
-# 5. Define the Inference Function
+# 4. Define the Inference Function
 ########################################
 def predict_image(img: Image.Image) -> str:
     img = img.convert("RGB")
@@ -82,17 +76,15 @@ def predict_image(img: Image.Image) -> str:
     return f"Prediction: {obj_name} ({bin_name})"
 
 ########################################
-# 6. Create Gradio UI
+# 5. Create Gradio UI
 ########################################
 demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Multi-Task Image Classifier Trained by [Abdellahi El Moustapha](https://abmstpha.github.io/)",
+    title="Multi-Task Image Classifier",
     description="Upload an image to receive two predictions:\n1) The primary object in the image,\n2) Whether the image is AI-generated or Real."
 )
 
-
-
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", share=True)
+    demo.launch(server_name="0.0.0.0", share=True)