Spaces:
Running
Running
| import streamlit as st | |
| from ultralytics import YOLO | |
| from huggingface_hub import hf_hub_download | |
| import cv2 | |
| import numpy as np | |
| import logging | |
| 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) | |