Update app.py

app.py

@@ -1,12 +1,12 @@
 import os
 import streamlit as st
+import base64
 # from transformers import pipeline
 from PIL import Image
-
+from st_clickable_images import clickable_images
 # pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
 
 
-
 st.title("Helmet Detection App: ⛑")
 confidence=st.slider("Confidence score (0-0.9)", 
             min_value=0.0, 
@@ -15,8 +15,12 @@ confidence=st.slider("Confidence score (0-0.9)",
             step=.1)
 file_name = st.file_uploader("Upload an image where workers may or may not wear helmets.")
 
-if file_name is not None:
+if (file_name is not None):
     col1, col2 = st.columns(2)
+    # img1 =""
+    # img2 =""
+    # col1.image(img1, use_column_width=True)
+    # col2.image(img2, use_column_width=True)
 
     image = Image.open(file_name)
     image.save("data/inputImg.jpg")
@@ -29,4 +33,31 @@ if file_name is not None:
     os.system("rm -rf yolov5/runs")
     col2.image(output, use_column_width=True)
     image.close()
-    output.close()
+    output.close()
+
+elif (st.checkbox('Select an example by clicking on the checkbox.')):
+    image_list=["https://media.istockphoto.com/id/1301722993/photo/group-of-contractors-celebrating-the-end-of-successful-construction-process.jpg?s=612x612&w=0&k=20&c=Xv7sO0AGHITZ6-SdTFrvXYnlWQ_Sc3wAcnGory4n1NA=","https://media.istockphoto.com/id/1472264590/photo/substation-maintenance-engineers.webp?b=1&s=170667a&w=0&k=20&c=dXhPK_sakEhOsAFb6InX4onVoaSgmUZ5A-V6eUlZf8E="]
+    # st.markdown("select a image")
+    clicked = clickable_images(
+        image_list,
+        titles=[f"Image #{str(i)}" for i in range(len(image_list))],
+        div_style={"display": "flex", "justify-content": "space-between", "flex-wrap": "wrap"},
+        img_style={"margin": "5px", "height": "200px"},
+        key=None,
+    )
+    st.markdown(f"Image # {clicked} selected" if clicked > -1 else "No image selected")
+    if clicked>-1:
+        col1, col2 = st.columns(2)
+        file_exp =f"sample/image{clicked+1}.jpg"
+        image = Image.open(file_exp)
+        image.save("data/inputImg.jpg")
+        col1.image(image, use_column_width=True)
+        os.system(f"cd yolov5/ && python detect.py --weights ../best.pt --img 416 --conf {confidence} --source ../data/inputImg.jpg")
+        im = Image.open("yolov5/runs/detect/exp/inputImg.jpg")
+        im.save("data/output.jpg")
+        im.close()
+        output =Image.open("data/output.jpg")
+        col2.image(output, use_column_width=True)
+        image.close()
+        output.close()
+        os.system("rm -rf yolov5/runs")