Mohansai2004 commited on
Commit
747b70e
·
1 Parent(s): b6e5d1c

changed the model

Browse files
Files changed (1) hide show
  1. app/caption_model.py +1 -1
app/caption_model.py CHANGED
@@ -2,7 +2,7 @@ from transformers import pipeline
2
  from PIL import Image
3
 
4
  # Load object detection model
5
- MODEL_NAME = "hustvl/yolos-small"
6
  detector = pipeline("object-detection", model=MODEL_NAME)
7
 
8
  def caption_image(image: Image.Image):
 
2
  from PIL import Image
3
 
4
  # Load object detection model
5
+ MODEL_NAME = "facebook/detr-resnet-50"
6
  detector = pipeline("object-detection", model=MODEL_NAME)
7
 
8
  def caption_image(image: Image.Image):