Spaces:
Runtime error
Runtime error
sshi
commited on
Commit
•
d78a0b4
1
Parent(s):
5efd297
App bug fix.
Browse files
app.py
CHANGED
@@ -75,10 +75,10 @@ feature_extractor = AutoFeatureExtractor.from_pretrained("hustvl/yolos-small", s
|
|
75 |
|
76 |
# Build model and load checkpoint
|
77 |
checkpoint = './checkpoints/epoch=1-step=2184.ckpt'
|
78 |
-
|
79 |
|
80 |
-
|
81 |
-
|
82 |
|
83 |
# colors for visualization
|
84 |
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
|
@@ -134,10 +134,10 @@ def visualize_preds(image, preds, threshold=0.9):
|
|
134 |
return plot_results(image, probas[keep], bboxes_scaled)
|
135 |
|
136 |
|
137 |
-
def detect(img
|
138 |
|
139 |
# Run inference
|
140 |
-
preds = generate_preds(feature_extractor,
|
141 |
|
142 |
return visualize_preds(img, preds)
|
143 |
|
|
|
75 |
|
76 |
# Build model and load checkpoint
|
77 |
checkpoint = './checkpoints/epoch=1-step=2184.ckpt'
|
78 |
+
model_yolos = Detr.load_from_checkpoint(checkpoint, lr=2.5e-5, weight_decay=1e-4)
|
79 |
|
80 |
+
model_yolos.to(device)
|
81 |
+
model_yolos.eval()
|
82 |
|
83 |
# colors for visualization
|
84 |
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
|
|
|
134 |
return plot_results(image, probas[keep], bboxes_scaled)
|
135 |
|
136 |
|
137 |
+
def detect(img):
|
138 |
|
139 |
# Run inference
|
140 |
+
preds = generate_preds(feature_extractor, model_yolos, img)
|
141 |
|
142 |
return visualize_preds(img, preds)
|
143 |
|