Spaces:
Runtime error
Runtime error
File size: 4,733 Bytes
26ccc67 |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import streamlit as st
import torch
from detect import detect
from PIL import Image
from io import *
import glob
from datetime import datetime
import os
import wget
import time
def imageInput(device, src):
if src == 'Upload your own data.':
image_file = st.file_uploader("Upload An Image", type=['png', 'jpeg', 'jpg'])
col1, col2 = st.columns(2)
if image_file is not None:
img = Image.open(image_file)
with col1:
st.image(img, caption='Uploaded Image', use_column_width='always')
ts = datetime.timestamp(datetime.now())
imgpath = os.path.join('data/uploads', str(ts)+image_file.name)
outputpath = os.path.join('data/outputs', os.path.basename(imgpath))
with open(imgpath, mode="wb") as f:
f.write(image_file.getbuffer())
#call Model prediction--
model = torch.hub.load('ultralytics/yolov5', 'custom', path='models/YOLOv5m.pt', force_reload=True)
model.cuda() if device == 'cuda' else model.cpu()
pred = model(imgpath)
pred.render() # render bbox in image
for im in pred.ims:
im_base64 = Image.fromarray(im)
im_base64.save(outputpath)
#--Display predicton
img_ = Image.open(outputpath)
with col2:
st.image(img_, caption='Model Prediction(s)', use_column_width='always')
elif src == 'From test set.':
# Image selector slider
imgpath = glob.glob('data/images/*')
imgsel = st.slider('Select random images from test set.', min_value=1, max_value=len(imgpath), step=1)
image_file = imgpath[imgsel-1]
submit = st.button("Predict!")
col1, col2 = st.columns(2)
with col1:
img = Image.open(image_file)
st.image(img, caption='Selected Image', use_column_width='always')
with col2:
if image_file is not None and submit:
#call Model prediction--
model = torch.hub.load('ultralytics/yolov5', 'custom', path='models/YOLOv5m.pt', force_reload=True)
pred = model(image_file)
pred.render() # render bbox in image
for im in pred.ims:
im_base64 = Image.fromarray(im)
im_base64.save(os.path.join('data/outputs', os.path.basename(image_file)))
#--Display predicton
img_ = Image.open(os.path.join('data/outputs', os.path.basename(image_file)))
st.image(img_, caption='Model Prediction(s)')
def videoInput(device, src):
uploaded_video = st.file_uploader("Upload Video", type=['mp4', 'mpeg', 'mov'])
if uploaded_video != None:
ts = datetime.timestamp(datetime.now())
imgpath = os.path.join('data/uploads', str(ts)+uploaded_video.name)
outputpath = os.path.join('data/video_output', os.path.basename(imgpath))
with open(imgpath, mode='wb') as f:
f.write(uploaded_video.read()) # save video to disk
st_video = open(imgpath, 'rb')
video_bytes = st_video.read()
st.video(video_bytes)
st.write("Uploaded Video")
detect(weights="models/yoloTrained.pt", source=imgpath, device=0) if device == 'cuda' else detect(weights="models/YOLOv5m.pt", source=imgpath, device='cpu')
st_video2 = open(outputpath, 'rb')
video_bytes2 = st_video2.read()
st.video(video_bytes2)
st.write("Model Prediction")
def main():
# -- Sidebar
st.sidebar.title('⚙️Options')
datasrc = st.sidebar.radio("Select input source.", ['From test set.', 'Upload your own data.'])
option = st.sidebar.radio("Select input type.", ['Image', 'Video'], disabled = True)
if torch.cuda.is_available():
deviceoption = st.sidebar.radio("Select compute Device.", ['cpu', 'cuda'], disabled = False, index=1)
else:
deviceoption = st.sidebar.radio("Select compute Device.", ['cpu', 'cuda'], disabled = True, index=0)
# -- End of Sidebar
st.header('Unreal Engine 5 Tank Demo')
st.subheader('Select the options')
if option == "Image":
imageInput(deviceoption, datasrc)
elif option == "Video":
videoInput(deviceoption, datasrc)
if __name__ == '__main__':
main()
@st.cache
def loadModel():
start_dl = time.time()
model_file = wget.download('https://huggingface.co/sigil-ml/Unreal-Engine-5-Tanks-YOLOv5/blob/main/YOLOv5m.pt', out="models/")
finished_dl = time.time()
print(f"Model Downloaded, ETA:{finished_dl-start_dl}")
loadModel() |