Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 = ["
|
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
|
|
|
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
|