manu02 commited on
Commit
97bc6d7
·
verified ·
1 Parent(s): 329f1e9

Update utils/complete_model.py

Browse files
Files changed (1) hide show
  1. utils/complete_model.py +4 -1
utils/complete_model.py CHANGED
@@ -409,7 +409,10 @@ def create_classifier(checkpoint_path, model_id="facebook/dinov3-vits16-pretrain
409
  try:
410
  # A. Load Checkpoint
411
  checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
412
- label_cols = checkpoint.get('label_cols', ["Class_1", "Class_2"]) # Fallback
 
 
 
413
 
414
  # B. Load Base Model (DINO)
415
  logger.info("🤖 Loading DINO backbone...")
 
409
  try:
410
  # A. Load Checkpoint
411
  checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
412
+ label_cols = checkpoint.get('label_cols', [
413
+ "Cardiomegaly", "Consolidation", "Edema",
414
+ "Atelectasis", "Pleural Effusion", "No Findings"
415
+ ])
416
 
417
  # B. Load Base Model (DINO)
418
  logger.info("🤖 Loading DINO backbone...")