File size: 1,190 Bytes
2f82ea0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
import torch
import streamlit as st
import requests
from pathlib import Path
# Define model path and URL
MODEL_PATH = "pytorch_model.pth"
MODEL_URL = "https://huggingface.co/zongzhuofan/co-detr-vit-large-coco/resolve/main/pytorch_model.pth"
# Download model if not exists
@st.cache_resource
def download_model():
if not Path(MODEL_PATH).exists():
with st.spinner("Downloading model... This might take a few minutes..."):
response = requests.get(MODEL_URL, stream=True)
with open(MODEL_PATH, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
return MODEL_PATH
# Load the model
model = YourModelClass()
model_path = download_model()
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()
st.title("Co-DETR Model")
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png"])
if uploaded_file is not None:
# Preprocess the image
input_data = preprocess_image(uploaded_file)
with torch.no_grad():
output = model(input_data)
# Postprocess output if necessary
st.write("Output:", output) |