File size: 1,453 Bytes
e191a7a fd8e705 e191a7a |
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 |
import matplotlib.pyplot as plt
import streamlit as st
from PIL import Image
import numpy as np
import torch
import cv2
@st.cache_resource
def load_model():
model = torch.hub.load(
'ultralytics/yolov5',
'custom',
path='best.pt'
)
return model
model = load_model()
def make_prediction(img_path):
img = Image.open(img_path).convert('RGB')
open_cv_image = np.array(img)
open_cv_image = open_cv_image[:, :, ::-1].copy()
results = model([open_cv_image])
return open_cv_image, results
def show_bounding_boxes(img, results, cls):
df = results.pandas().xyxy[0]
df = df[df['name'] == cls]
list_of_rows = [list(row) for row in df.values]
for each_row in list_of_rows:
cv2.rectangle(img,
(int(each_row[0]), int(each_row[1])),
(int(each_row[2]), int(each_row[3])),
(255, 0, 0), 2
)
return img
st.title("RandomLayNet (Layout Detection) :tea: :coffee:")
upload = st.file_uploader(label= "Upload an image : ", type = ["png","jpg","jpeg"])
if upload:
print(upload)
img, results = make_prediction(upload)
bounded_img = show_bounding_boxes(img, results, 'Table')
fig = plt.figure(figsize=(12,12))
ax = fig.add_subplot(111)
plt.imshow(bounded_img)
plt.xticks([],[])
plt.yticks([],[])
ax.spines[["top","bottom","right","left"]].set_visible(False)
st.pyplot(fig,use_container_width=True)
|