fix multi run image reads and provide a sample for gringos

src/hockey_object_detection.py

@@ -9,46 +9,77 @@ 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)}
+    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.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}')
+                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 
+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')
+
+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 ML algorithm to use YOLO object detection (YOLO v8). 
+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**:
@@ -63,38 +94,51 @@ The output will be a set of hockey objects (defined by "bounding boxes") with la
 st.markdown(desc)
 
 st.subheader("Sample")
-st.image('src/images/samples/v2/v2-sample1-090124.png', caption='Sample image with hockey objects detected')
+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.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.")
+
 
-st.subheader("Try It Out")
+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)
 
-img = st.file_uploader("Upload an image for object detection", type=["jpg", "jpeg", "png"])
+    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', min_value=0.0, max_value=1.0, value=0.5, step=0.01)
+
+    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("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)
+        objects = run_inference(img, model_f, thresh)
+        
         st.dataframe(objects, column_config={
-            "0": "Object",
-            "1": "Confidence"
+                "0": "Object",
+                "1": "Confidence"
         })
 
     # check if the results list is empty
     if len(st.session_state.results) == 0:
-        st.image(img, caption='Uploaded Image')
+        st.write('**No hockey objects found in image!**')
+        st.image(img, caption='Uploaded Image had no hockey objects')
     else:
-        st.image(frame, caption='Uploaded Image')
+        st.image(img, caption='Image with hockey object bounding boxes')
+
+    st.button("Reset Image", on_click=reset_image)
 
-else:
-    st.session_state.results = []