sshi commited on
Commit
d78a0b4
1 Parent(s): 5efd297

App bug fix.

Browse files
Files changed (1) hide show
  1. app.py +5 -5
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
- model = Detr.load_from_checkpoint(checkpoint, lr=2.5e-5, weight_decay=1e-4)
79
 
80
- model.to(device)
81
- model.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,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, model):
138
 
139
  # Run inference
140
- preds = generate_preds(feature_extractor, model, img)
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