Update app.py

app.py

@@ -27,8 +27,8 @@ class MultiTaskModel(nn.Module):
 ########################################
 # 2. Reconstruct the Model and Load Weights
 ########################################
-# Set the number of object classes (update this to match your training)
-num_obj_classes = 139  # Example value; update as needed
+# IMPORTANT: Use the same number of object classes as in training.
+num_obj_classes = 494  # Updated to match the state dict from training
 
 device = torch.device("cpu")
 
@@ -43,7 +43,7 @@ model.to(device)
 
 # Download the state dict from HF Hub.
 repo_id = "Abdu07/multitask-model"  # Your repo name
-filename = "Yolloplusclassproject_weights.pth"  # New weight file name
+filename = "Yolloplusclassproject_weights.pth"  # The state dict file you uploaded
 weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
 
 # Load the state dict and update the model.
@@ -55,12 +55,10 @@ model.eval()
 # 3. Define Label Mappings and Transforms
 ########################################
 # Update these with your actual label mappings.
-idx_to_obj_label = {
-    0: "cat",
-    1: "dog",
-    2: "car",
-    # ... add the rest of your object classes ...
-}
+# They should reflect the 494 unique pseudo-labels produced during training.
+# For this example, we assume that the mapping is stored somewhere.
+# Here we provide a dummy mapping for illustration. Replace it with your real mapping.
+idx_to_obj_label = {i: f"label_{i}" for i in range(num_obj_classes)}
 bin_label_names = ["AI-Generated", "Real"]
 
 # Define the validation transforms (must match those used during training)