import streamlit as st
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import cv2
import numpy as np

@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


if 'results' not in st.session_state:
    st.session_state.results = []

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('Object Detection')

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 ML algorithm 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")
st.image('src/images/samples/v2/v2-sample1-090124.png', caption='Sample image with hockey objects detected')

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", )

st.subheader("Try It Out")

img = st.file_uploader("Upload an image for object detection", type=["jpg", "jpeg", "png"])

if img is not None:
    thresh =st.slider('Set the object confidence threshold', min_value=0.0, max_value=1.0, value=0.5, step=0.01)
    with st.status("Detecting hockey objects..."):
        st.write("Loading model...")
        model_f = get_model()

        st.write("Processing image...")
        frame = cv2.imdecode(np.frombuffer(img.read(), np.uint8), 1)
        
        st.write("Running inference on image...")
        objects = run_inference(frame, 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.image(img, caption='Uploaded Image')
    else:
        st.image(frame, caption='Uploaded Image')

else:
    st.session_state.results = []