Update app.py

app.py

@@ -6,6 +6,7 @@
 import streamlit as st
 import streamlit.components.v1 as components
 from PIL import ImageOps, Image
+from ultralytics import YOLO
 
 #####
 ###
@@ -106,10 +107,21 @@ with row3_col5:
 # Output and Interaction
 ###
 #####
-st.write("\n\n")
-components.html("""<hr style="height:2px;border:none;color:#333;background-color:#333;" /> """)
+if st.session_state.image is not None:
         
-_, col_1, col_2, _ = st.columns([1, 4, 4, 1], gap="medium")
-with col_2:
-    if st.session_state.image is not None:
-        st.image(st.session_state.image, use_column_width=True)
+    st.write("\n\n")
+    components.html("""<hr style="height:2px;border:none;color:#333;background-color:#333;" /> """)
+
+    # model
+    model = YOLO("/home/user/app/models/head_model/best.pt")  # load
+    pred = 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)
+    with col_2:
+        st.image(st.session_state.image, use_column_width=True)
+    
+