Update app.py

app.py

@@ -118,15 +118,35 @@ else:
     print("Skipping TensorFlow model loading because prerequisite Hugging Face models failed to load.")
 
 # Load parts list from JSON
+# --- (Corrected Code) ---
 PARTS_LIST_FILE = 'cars117.json'
 all_parts = []
 if os.path.exists(PARTS_LIST_FILE):
     with open(PARTS_LIST_FILE, 'r', encoding='utf-8') as f:
         data = json.load(f)
+    # Get the unique, sorted list of parts from the JSON file
     all_parts = sorted(list(set(part for entry in data.values() for part in entry.get('replaced_parts', []))))
-    if dl_model and dl_model.output_shape[-1] != len(all_parts):
-        print(f"Warning: TensorFlow model output classes ({dl_model.output_shape[-1]}) "
-              f"does not match number of parts in JSON ({len(all_parts)}). Predictions may be misaligned.")
+
+    # FIX: Dynamically handle mismatch between model output and parts list
+    if dl_model is not None:
+        model_output_size = dl_model.output_shape[-1]
+        parts_list_size = len(all_parts)
+
+        if model_output_size != parts_list_size:
+            print(f"WARNING: Model output size ({model_output_size}) and parts list size ({parts_list_size}) do not match.")
+            
+            # If the model expects more outputs, pad the parts list with dummy values
+            if model_output_size > parts_list_size:
+                diff = model_output_size - parts_list_size
+                print(f"Padding the parts list with {diff} dummy entries to prevent a crash.")
+                for i in range(diff):
+                    all_parts.append(f"_dummy_part_{i+1}_")
+            
+            # If the model expects fewer outputs, truncate the parts list
+            else:
+                diff = parts_list_size - model_output_size
+                print(f"Truncating the parts list by {diff} entries to match the model's output.")
+                all_parts = all_parts[:model_output_size]
 else:
     print(f"Error: Parts list file '{PARTS_LIST_FILE}' not found. Predicted part names will be unavailable.")