Update main.py

main.py

@@ -1,5 +1,9 @@
 import streamlit as st
 import tensorflow as tf
+import streamlit as st
+import argparse
+import io
+from io import BytesIO
 from PIL import ImageOps, Image
 import os                        # To work with operation system commands
 import cv2                       # To process images
@@ -13,6 +17,24 @@ from ultralytics import YOLO     # To Create Yolo model
 from termcolor import colored    # To colorfull outputs
 import matplotlib.pyplot as plt  # To visualizations
 
+st.set_page_config(
+    page_title="Auto NPR",
+    page_icon="✨",
+    layout="centered",
+    initial_sidebar_state="expanded",
+)
+
+top_image = Image.open('banner_top.png')
+bottom_image = Image.open('banner_bottom.png')
+main_image = Image.open('main_banner.png')
+
+st.image(main_image,use_column_width='auto')
+st.title(' Automatic Number Plate Recognition 🚘🚙')
+st.sidebar.image(top_image,use_column_width='auto')
+st.sidebar.header('Input 🛠')
+selected_type = st.sidebar.selectbox('Please select an activity type 🚀', ["Upload Image", "Live Video Feed"])
+st.sidebar.image(bottom_image,use_column_width='auto')
+
 #load weights
 best_model_plate = YOLO('best.pt')
 best_model_digits = YOLO('best2.pt')
@@ -71,9 +93,8 @@ def FINAL(img) :
     result is digits and char on car plate.
     '''
     # Read car image ( STEP-1 )
-    #img = cv2.imread(img)
-    #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    img = Image.open(img).convert('RGB')
+    img = cv2.imread(img)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
     # First prediction -> Detect car-plate ( STEP-2 )
     result1, _ = Detect_Plate(img)
@@ -95,19 +116,94 @@ def FINAL(img) :
 
     Plot_Result(img, raw_plate, digits)
 
-# set title
-st.title('Persian Plates Digits Detection')
+if selected_type == "Upload Image":
+    st.info('✨ Supports all popular image formats 📷 - PNG, JPG, BMP 😉')
+    uploaded_file = st.file_uploader("Upload Image of car's number plate 🚓", type=["png","jpg","bmp","jpeg"])
 
-# set header
-st.header('Please upload image')
+    if uploaded_file is not None:
+        with open(os.path.join(upload_path,uploaded_file.name),"wb") as f:
+            f.write((uploaded_file).getbuffer())
+        with st.spinner(f"Working... 💫"):
+            uploaded_image = os.path.abspath(os.path.join(upload_path,uploaded_file.name))
+            downloaded_image = os.path.abspath(os.path.join(download_path,str("output_"+uploaded_file.name)))
 
-# upload file
-img = st.file_uploader('', type=['jpeg', 'jpg', 'png'])
+            with open(uploaded_image,'rb') as imge:
+                img_bytes = imge.read()
 
-# display image
-if img is not None:
-    image = Image.open(img).convert('RGB')
-    st.image(image, use_column_width=True)
-
-
-st.write(FINAL(img))
+            img = Image.open(io.BytesIO(img_bytes))
+            results = FINAL(img)
+    
+            for img in results.imgs:
+                img_base64 = Image.fromarray(img)
+                img_base64.save(downloaded_image, format="JPEG")
+
+            final_image = Image.open(downloaded_image)
+            print("Opening ",final_image)
+            st.markdown("---")
+            st.image(final_image, caption='This is how your final image looks like 😉')
+            with open(downloaded_image, "rb") as file:
+                if uploaded_file.name.endswith('.jpg') or uploaded_file.name.endswith('.JPG'):
+                    if st.download_button(
+                                            label="Download Output Image 📷",
+                                            data=file,
+                                            file_name=str("output_"+uploaded_file.name),
+                                            mime='image/jpg'
+                                         ):
+                        download_success()
+                if uploaded_file.name.endswith('.jpeg') or uploaded_file.name.endswith('.JPEG'):
+                    if st.download_button(
+                                            label="Download Output Image 📷",
+                                            data=file,
+                                            file_name=str("output_"+uploaded_file.name),
+                                            mime='image/jpeg'
+                                         ):
+                        download_success()
+
+                if uploaded_file.name.endswith('.png') or uploaded_file.name.endswith('.PNG'):
+                    if st.download_button(
+                                            label="Download Output Image 📷",
+                                            data=file,
+                                            file_name=str("output_"+uploaded_file.name),
+                                            mime='image/png'
+                                         ):
+                        download_success()
+
+                if uploaded_file.name.endswith('.bmp') or uploaded_file.name.endswith('.BMP'):
+                    if st.download_button(
+                                            label="Download Output Image 📷",
+                                            data=file,
+                                            file_name=str("output_"+uploaded_file.name),
+                                            mime='image/bmp'
+                                         ):
+                        download_success()
+    else:
+        st.warning('⚠ Please upload your Image 😯')
+
+else:
+    st.info('✨ The Live Feed from Web-Camera will take some time to load up 🎦')
+    live_feed = st.checkbox('Start Web-Camera ✅')
+    FRAME_WINDOW = st.image([])
+    cap = cv2.VideoCapture(0)
+
+    if live_feed:
+        while(cap.isOpened()):
+            success, frame = cap.read()
+            if success == True:
+                ret,buffer=cv2.imencode('.jpg',frame)
+                frame=buffer.tobytes()
+                img = Image.open(io.BytesIO(frame))
+                model = instantiate_model()
+                results = model(img, size=640)
+                results.print()
+                img = np.squeeze(results.render())
+                img_BGR = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+            else:
+                break
+            frame = cv2.imencode('.jpg', img_BGR)[1].tobytes()
+            FRAME_WINDOW.image(frame)
+    else:
+        cap.release()
+        cv2.destroyAllWindows()
+        st.warning('⚠ The Web-Camera is currently disabled. 😯')
+
+st.markdown("<br><hr><center>Made with ❤️ by <a href='mailto:[email protected]?subject=Automatic Number Plate Recognition WebApp!&body=Please specify the issue you are facing with the app.'><strong>Prateek Ralhan</strong></a></center><hr>", unsafe_allow_html=True)