JohnAlexander23 commited on
Commit
cf24a38
1 Parent(s): 7434ac2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import streamlit as st
2
  import torch
3
- from transformers import AutoProcessor, AutoModelForObjectDetection
4
  from PIL import Image
5
  import requests
6
  from io import BytesIO
@@ -15,9 +15,9 @@ THRESHOLD = 0.2
15
  # Load model and processor
16
  @st.cache_resource
17
  def load_model():
18
- model_id = 'onnx-community/yolov10m'
19
- processor = AutoProcessor.from_pretrained(model_id)
20
- model = AutoModelForObjectDetection.from_pretrained(model_id)
21
  return processor, model
22
 
23
  processor, model = load_model()
 
1
  import streamlit as st
2
  import torch
3
+ from transformers import YolosImageProcessor, YolosForObjectDetection
4
  from PIL import Image
5
  import requests
6
  from io import BytesIO
 
15
  # Load model and processor
16
  @st.cache_resource
17
  def load_model():
18
+ model_id = 'hustvl/yolos-tiny'
19
+ processor = YolosImageProcessor.from_pretrained(model_id)
20
+ model = YolosForObjectDetection.from_pretrained(model_id)
21
  return processor, model
22
 
23
  processor, model = load_model()