-
Notifications
You must be signed in to change notification settings - Fork 0
/
streamlit_app.py
82 lines (72 loc) · 2.6 KB
/
streamlit_app.py
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
import streamlit as st
import torch
from datetime import datetime
import os
import requests
import time
import cv2
from PIL import Image
import numpy as np
from matplotlib import pyplot as pyplot
from git import Repo
@st.experimental_singleton
def loadModel():
start_dl = time.time()
Repo.clone_from("clone yolov7 github repository')
os.chdir('yolov7')
yolo_model=requests.get("custom yolo model weights.pt file")
with open("best.pt", 'wb')as file:
file.write(yolo_model.content)
finished_dl = time.time()
print(f"Model Downloaded, ETA:{finished_dl-start_dl}")
def page1():
image_file=st.sidebar.file_uploader("choose image file",type=['png','jpg','jpeg'])
if image_file is not None:
img=Image.open(image_file)
#img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
#img_arr=np.array(img)
#img = cv2.imdecode(img, 1)
col1, col2 = st.columns(2)
with col1:
st.image(img, caption='Uploaded Image', use_column_width='always',channels='RGB')
outpath = os.path.join(os.getcwd(), f"out_{os.path.basename(image_file.name)}")
model=torch.hub.load(".",'custom','best.pt',source='local')
if st.button('Detect Pothole'):
pred = model(img)
pred.render() # render bbox in image
for im in pred.imgs:
im_base64 = Image.fromarray(im)
im_base64.save(outpath)
img_ = Image.open(outpath)
with col2:
st.image(img_, caption='Model Prediction(s)', use_column_width='always',channels='RGB')
with st.expander("View Annotation Data"):
tab1, tab2, tab3 = st.tabs(['Pascal VOC', 'COCO','YOLO'])
pred = model(img)
with tab1:
df1=pred.pandas().xyxy[0]
st.dataframe(df1)
st.download_button(
label="Download Annotation Data as CSV",
data=df1.to_csv(),
file_name=f"Annotation(Pascal VOC) Data For {image_file.name}.csv",
mime='text/csv')
with tab2:
df1=pred.pandas().xywh[0]
st.dataframe(df1)
st.download_button(
label="Download Annotation Data as CSV",
data=df1.to_csv(),
file_name=f"Annotation(COCO) Data For {image_file.name}.csv",
mime='text/csv')
with tab3:
df1=pred.pandas().xywhn[0]
st.dataframe(df1)
st.download_button(
label="Download Annotation Data as CSV",
data=df1.to_csv(),
file_name=f"Annotation(YOLO) Data For {image_file.name}.csv",
mime='text/csv')
if __name__ == '__main__':
page1()
loadModel()