Update app.py

app.py

@@ -95,7 +95,10 @@ def create_classifier(checkpoint_path, model_id="facebook/dinov3-vits16-pretrain
     device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
     print(f"Loading Classifier from {checkpoint_path}...")
     checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
-    label_cols = checkpoint.get('label_cols', ["Class_1", "Class_2"])
+    label_cols = checkpoint.get('label_cols', [
+            "Cardiomegaly", "Consolidation", "Edema",
+            "Atelectasis", "Pleural Effusion", "No Findings"
+        ])
     
     base_model = AutoModel.from_pretrained(model_id).to(device)
     if 'base_model_state_dict' in checkpoint: