Update app.py

app.py

@@ -24,55 +24,51 @@ class MultiTaskModel(nn.Module):
         return obj_logits, bin_logits
 
 ########################################
-# 2. Reconstruct the Model and Load Weights
+# 2. Load the Label Mapping and Set num_obj_classes
 ########################################
-# IMPORTANT: The checkpoint was saved with a single object class,
-# so we set num_obj_classes to 1.
-num_obj_classes = 1
+# 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),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
 
+########################################
+# 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"
+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()
 
 ########################################
-# 3. Load Label Mapping and Define Transforms
-########################################
-# Attempt to load the mapping from JSON.
-# If the mapping contains more than one label, we override it with a single-label mapping
-try:
-    with open("obj_label_mapping.json", "r") as f:
-        obj_label_to_idx = json.load(f)
-    if len(obj_label_to_idx) != 1:
-        obj_label_to_idx = {"Detected Object": 0}
-except Exception as e:
-    print("Error loading mapping, using default mapping. Error:", e)
-    obj_label_to_idx = {"Detected Object": 0}
-
-idx_to_obj_label = {v: k for k, v in obj_label_to_idx.items()}
-
-bin_label_names = ["AI-Generated", "Real"]
-
-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])
-])
-
-########################################
-# 4. Define the Inference Function
+# 5. Define the Inference Function
 ########################################
 def predict_image(img: Image.Image) -> str:
     img = img.convert("RGB")
@@ -86,13 +82,13 @@ def predict_image(img: Image.Image) -> str:
     return f"Prediction: {obj_name} ({bin_name})"
 
 ########################################
-# 5. Create Gradio UI
+# 6. Create Gradio UI
 ########################################
 demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Multi-Task Image Classifier",
+    title="Multi-Task Image Classifier Trained by [Abdellahi El Moustapha](https://abmstpha.github.io/),
     description="Upload an image to receive two predictions:\n1) The primary object in the image,\n2) Whether the image is AI-generated or Real."
 )