Upload app.py

app.py

@@ -359,7 +359,12 @@ class AppState:
         print(f"Loading RF-DETR from: {checkpoint}")
         self.model = load_model(checkpoint, self.config.get('resolution'))
 
-        # Try to load class names
+        # Set default class names (can be overridden by results.json)
+        # Index corresponds to class_id from the model
+        default_class_names = ["Background", "Granuloma"]
+        
+        # Try to load class names from results.json
+        loaded_classes = None
         try:
             results_json = "/tmp/results/results.json"
             if os.path.isfile(results_json):
@@ -372,10 +377,15 @@ class AppState:
                             name = item.get("class")
                             if name and name != "all" and name not in classes:
                                 classes.append(name)
-                self.class_names = classes if classes else None
-        except Exception:
-            pass
-
+                if classes:
+                    loaded_classes = classes
+        except Exception as e:
+            print(f"Could not load class names from results.json: {e}")
+        
+        # Use loaded classes if available, otherwise use defaults
+        self.class_names = loaded_classes if loaded_classes else default_class_names
+        print(f"Using class names: {self.class_names}")
+        
         print("✓ RF-DETR model loaded")
 
     def preload_all_models(self):