Update app.py

app.py

@@ -114,13 +114,16 @@ if st.session_state.image is not None:
 
     # model
     model = YOLO("/home/user/app/models/head_model/best.pt")  # load
-    pred = model.predict([st.session_state.image]) # predict
+    results = model.predict([st.session_state.image]) # predict
     
     
     _, col_1, col_2, _ = st.columns([1, 4, 4, 1], gap="medium")
     with col_1:
         st.write('small image 1 here')
-        st.image(pred[0])
+        # Plot predicted image
+        im_bgr = results[0].plot()  # BGR-order numpy array
+        im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
+        st.image(im_rgb)
     with col_2:
         st.image(st.session_state.image, use_column_width=True)