denizaybey commited on
Commit
9f08bc7
·
verified ·
1 Parent(s): 4a4a14b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -72,11 +72,11 @@ def detect_objects(
72
  device = "cuda" if torch.cuda.is_available() else "cpu"
73
  model, image_processor = get_model_and_processor(checkpoint)
74
  model = model.to(device)
75
- classes = ["Airplane", "Drone", "Helicopter", "Satellite", "Quadcopter", "Vehicle"]
76
  if classes is not None:
77
  wrong_classes = [cls for cls in classes if cls not in model.config.label2id]
78
  if wrong_classes:
79
- gr.Warning(f"Classes not found in model config: {wrong_classes}")
80
  keep_ids = [model.config.label2id[cls] for cls in classes if cls in model.config.label2id]
81
  else:
82
  keep_ids = None
 
72
  device = "cuda" if torch.cuda.is_available() else "cpu"
73
  model, image_processor = get_model_and_processor(checkpoint)
74
  model = model.to(device)
75
+ classes = ["person","aeroplane","bicycle","car","motorbike","bus","train","truck","boat"]
76
  if classes is not None:
77
  wrong_classes = [cls for cls in classes if cls not in model.config.label2id]
78
  if wrong_classes:
79
+ gr.Warning(f"Classes not found in model config")
80
  keep_ids = [model.config.label2id[cls] for cls in classes if cls in model.config.label2id]
81
  else:
82
  keep_ids = None