Antoine101 commited on
Commit
38a4f39
·
verified ·
1 Parent(s): c87829b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -26,8 +26,8 @@ def predict(image):
26
  obj_detector = pipeline(
27
  "object-detection", model="anindya64/detr-resnet-50-dc5-hardhat-finetuned"
28
  )
29
- results = obj_detector(train_dataset[0]["image"])
30
- return plot_results(image)
31
 
32
  title = "Do you have your helmet on?"
33
  description = """
 
26
  obj_detector = pipeline(
27
  "object-detection", model="anindya64/detr-resnet-50-dc5-hardhat-finetuned"
28
  )
29
+ results = obj_detector(image)
30
+ return plot_results(image, results)
31
 
32
  title = "Do you have your helmet on?"
33
  description = """