Update app.py

app.py

@@ -26,7 +26,10 @@ class MultiTaskModel(nn.Module):
 ########################################
 # 2. Reconstruct the Model and Load Weights
 ########################################
-num_obj_classes = 494  # Make sure this matches your training
+# IMPORTANT: The checkpoint was saved with a single object class,
+# so we set num_obj_classes to 1.
+num_obj_classes = 1
+
 device = torch.device("cpu")
 
 resnet = models.resnet50(pretrained=False)
@@ -45,10 +48,17 @@ 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)
-# Create the inverse mapping
+# 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"]