Spaces:
Running
Running
File size: 5,075 Bytes
6d4cc80 faeb28b 6d4cc80 faeb28b 6d4cc80 6445ea5 6d4cc80 faeb28b |
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 123 124 125 126 127 128 129 130 131 |
import streamlit as st
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import cv2
import numpy as np
import logging
@st.cache_resource
def get_model():
repo_id = "danbiagini/hockey_breeds_v2"
return hf_hub_download(repo_id=repo_id, filename="hockey_breeds-v2-101623.pt")
def run_inference(img, model, thresh=0.5):
model = YOLO(model_f)
st.session_state.results = model(img)
return draw_hockey_boxes(img, st.session_state.results, thresh)
def draw_hockey_boxes(frame, results, thresh=0.5):
colors = {0: (0, 255, 0), 1: (255, 0, 0), 2: (0, 0, 255), 3: (128, 0, 0), 4: (
0, 128, 0), 5: (0, 0, 128), 6: (0, 64, 0), 7: (64, 0, 0), 8: (0, 0, 64)}
font_scale = frame.shape[0] / 500
objects = []
for name in results:
for box in name.boxes.data.tolist():
x1, y1, x2, y2, score, class_id = box
objects.append((name.names[int(class_id)], score))
if score > thresh:
cv2.rectangle(frame, (int(x1), int(y1)),
(int(x2), int(y2)), colors[(class_id % 9)], 3)
cv2.putText(frame, f'{name.names[int(class_id)].upper()}: {score:.2f}', (int(x1), int(y1 - 10)),
cv2.FONT_HERSHEY_SIMPLEX, font_scale, colors[(class_id % 9)], 3, cv2.LINE_AA)
else:
print(
f'Found an object under confidence threshold {thresh} type: {name.names[class_id]}, score:{score}, x1, y2:{x1}, {y2}')
return objects
def reset_image():
st.session_state.img = None
def upload_img():
if st.session_state.upload_img is not None:
st.session_state.img = st.session_state.upload_img
def get_naked_image():
if st.session_state.img is not None:
img = st.session_state.img
img.seek(0)
return(cv2.imdecode(np.frombuffer(img.read(), np.uint8), 1))
return None
def use_sample_image():
st.session_state.img = open('src/images/samples/v2/net-chaos.jpg', 'rb')
# Init state
if 'results' not in st.session_state:
st.session_state.results = []
if 'thresh' not in st.session_state:
st.session_state.thresh = 0.5
if 'img' not in st.session_state:
st.session_state.img = None
st.set_page_config(page_title='TopShelf POC', layout="wide",
page_icon="🥅")
st.title('Welcome To Top Shelf :goal_net:',
help=':video_camera: + :ice_hockey_stick_and_puck: = :clipboard:')
st.subheader('Artificial Intelligence for Hockey Coaches and Players',
help='Proof of concept application')
overview = '''**Top Shelf** helps coaches and players analyze their gameplay, providing helpful suggestions & recommendations on areas for improvement.
We're starting with a focus on ice hockey, however this same technology could apply to other "invasion" games and sports, for example lacrosse, basketball, soccer, etc.
The special sauce behind **Top Shelf** is AI *Computer Vision* technology that recognizes various hockey related objects in videos.
The foundation of the technology is an AI model that can recognize players, nets, referees, pucks, and rink areas.
**Top Shelf** uses this technology to analyze game play and provide insightful suggestions on areas for improvement.
'''
st.markdown(overview)
st.subheader('Getting Started')
st.markdown('''We're currently in the training and testing phase of **Top Shelf** development. This is a proof of concept application that friends of **Top Shelf** can use to help in development.
To help us understand how our *Computer Vision* model is working you can upload hockey pictures and then the app will display what hockey objects were found. ''')
st.write("Upload an image file to try detecting hockey objects in your own hockey image, or use a sample image below.")
if st.session_state.img is None:
st.file_uploader("Upload an image and Hockey Breeds v2 will find the hockey objects in the image",
type=["jpg", "jpeg", "png"], key='upload_img', on_change=upload_img)
with st.expander("Sample Images"):
st.image('src/images/samples/v2/net-chaos.jpg')
st.button("Use Sample", on_click=use_sample_image)
img = get_naked_image()
if img is not None:
thresh = st.slider('Set the object confidence threshold', key='thresh',
min_value=0.0, max_value=1.0, value=0.5, step=0.05)
with st.status("Detecting hockey objects..."):
st.write("Loading model...")
model_f = get_model()
st.write("Running inference on image...")
objects = run_inference(img, model_f, thresh)
st.dataframe(objects, column_config={
"0": "Object",
"1": "Confidence"
})
# check if the results list is empty
if len(st.session_state.results) == 0:
st.write('**No hockey objects found in image!**')
st.image(img, caption='Uploaded Image had no hockey objects')
else:
st.image(img, caption='Image with hockey object bounding boxes')
st.button("Reset Image", on_click=reset_image)
|