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 | |
| # Top down page rendering | |
| st.set_page_config(page_title='Hockey Breeds v2 - Objects', layout="wide", | |
| page_icon=":frame_with_picture:") | |
| st.title('Hockey Breeds v2 - Objects') | |
| intro = '''The first version of Hockey Breeds was fun and educational, but not useful for analyzing hockey videos. The second version is to a proof of concept | |
| with the ability to recognize individual "objects" within an image, which paves the way to ultimately tracking those objects through game play.''' | |
| st.markdown(intro) | |
| st.subheader('Try Hockey Breeds v2') | |
| st.markdown('''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) | |
| st.subheader('Object Detection Technical Details') | |
| desc = '''Hockey Breed detector v2 uses a state of the art (circa 2023) computer vision approach. | |
| I used the same training images as the first version of the Hockey Breeds model, but change the Neural Network to use YOLO object detection (YOLO v8). | |
| The output will be a set of hockey objects (defined by "bounding boxes") with labels for any hockey image uploaded. | |
| **Object List**: | |
| 1. net | |
| 1. stick | |
| 1. puck | |
| 1. skater | |
| 1. goalie | |
| 1. referee | |
| ''' | |
| st.markdown(desc) | |
| st.subheader("Sample") | |
| try: | |
| st.image('src/images/samples/v2/v2-sample1-090124.png', | |
| caption='Sample image with hockey objects detected') | |
| except Exception as e: | |
| st.error(f"Error loading image: {e}") | |
| st.subheader("Validation Results") | |
| st.markdown('''Validation of the model\'s performance was done using 15 images not included in the training set. The model had many issues; it did poorly with detecting *pucks* and *sticks* vs backgrounds and even goalies and skaters. It did very well on detecting referees.''') | |
| st.image("src/images/artifacts/confusion_matrix_v2.png", | |
| caption="Confusion Matrix for Hockey Breeds v2", ) | |