Update app.py

app.py

@@ -237,7 +237,24 @@ if st.session_state.screen == 2:
             st.markdown("<p style='text-align: center; color: black;'>This process takes around 30 seconds.</p>", unsafe_allow_html=True)
             st.session_state.results = st.session_state.model.predict(st.session_state.tiles, conf=st.session_state.confidence) # predict     
             st.rerun()
-        
+
+        # show character locations
+        class_names = {0: "Waldo", 1: "Wenda", 2: "Odlaw", 3: "Wizard", 4: "Woof"}
+        for idx, pred in enumerate(results):
+            #extract detected character names and their confidence scores
+            st.session_state.detected_characters = {}
+            for *_, cls_id, conf in pred:
+                name = class_names[int(cls_id)]  # Convert class ID to character name
+                if name in st.session_state.detected_characters:
+                    if conf >= st.session_state.confidence:
+                        st.session_state.detected_characters[name].append(idx)
+
+        # display locations
+        st.session_state.locations = ""
+        for name, idx in st.session_state.detected_characters.items():
+            st.session_state.locations += f" {name} page: {str(idx)}, )"
+        st.write(st.session_state.locations)
+            
         # plot predicted image
         st.markdown(f"<h4 style='text-align: center; color: black;'>Image Tile {str(st.session_state.image_counter + 1)}</h4>", unsafe_allow_html=True)
         im_bgr = st.session_state.results[st.session_state.image_counter].plot()  # BGR-order numpy array