add v2 try it to Home page

src/Home.py

@@ -1,7 +1,72 @@
 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="🥅")
 
@@ -24,3 +89,42 @@ 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)
+

src/app.py

@@ -5,8 +5,10 @@ import logging
 app = st.navigation(
     {"App": [st.Page("Home.py", title="Home", icon=":material/home:"),
              st.Page("About.py", icon="🥅")],
-     "Models": [st.Page("Hockey_Breeds.py", title="v1 - Hockey Breeds", icon=":material/gradient:"),
-                st.Page("hockey_object_detection.py", title="v2 - Hockey Object Detection", icon=":material/filter_b_and_w:")]
+     "Models": [
+                st.Page("hockey_object_detection.py", title="v2 - Hockey Object Detection", icon=":material/filter_b_and_w:"),
+                st.Page("Hockey_Breeds.py", title="v1 - Hockey Breeds", icon=":material/gradient:")
+                ]
      }
 )
 app.run()

src/hockey_object_detection.py

@@ -1,69 +1,4 @@
 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
-
-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",
@@ -103,42 +38,3 @@ st.markdown('''Validation of the model\'s performance was done using 15 images n
 st.image("src/images/artifacts/confusion_matrix_v2.png",
          caption="Confusion Matrix for Hockey Breeds v2", )
 
-st.subheader("Try it out!")
-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)
-