Update app.py

app.py

@@ -1,3 +1,4 @@
+import cv2
 import torch
 import streamlit as st
 import numpy as np
@@ -14,7 +15,8 @@ st.title('Upload an image file to detection')
 col1, col2, col3 = st.columns([0.5, 0.3, 0.2])
 
 with col1:
-        
+
+    st.header('flyer')
     uploaded_image_zero_file = st.file_uploader("Choose an image file (zero)")
     
     imageLocation = st.empty()
@@ -31,6 +33,9 @@ with col1:
     if st.button('detect'):
     
         draw = ImageDraw.Draw(image)
+
+        arrImg = np.array(image)
+        imageCv = cv2.cvtColor(arrImg, cv2.COLOR_RGB2BGR)
     
         for prediction in predictions:
             box = prediction["box"]
@@ -40,6 +45,11 @@ with col1:
             xmin, ymin, xmax, ymax = box.values()
             draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=5)
             draw.text((xmin, ymin), f"{label}: {round(score,5)}", fill="white")
+
+            crop_img = imageCv[ymin:ymax, xmin:xmax]
+            if crop_img.size > 0:
+                cv2.imwrite("data/cropped_image_{}.jpg".format(i), crop_img)
+                i += 1
     
         imageLocation.image(image)
         st.write(predictions)