Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| import imutils | |
| import easyocr | |
| import os | |
| import pathlib | |
| import platform | |
| from xyxy_converter import yolov5_to_image_coordinates | |
| import shutil | |
| system_platform = platform.system() | |
| if system_platform == 'Windows': pathlib.PosixPath = pathlib.WindowsPath | |
| CUR_DIR = os.getcwd() | |
| YOLO_PATH = f"{CUR_DIR}/yolov5" | |
| MODEL_PATH = "runs/train/exp/weights/best.pt" | |
| def main(): | |
| st.title("Odometer value extractor with Streamlit") | |
| # Use st.camera to capture images from the user's camera | |
| img_file_buffer = st.camera_input(label='Please, take a photo of odometer', key="odometer") | |
| # Check if an image is captured | |
| if img_file_buffer is not None: | |
| # Convert the image to a NumPy array | |
| image = Image.open(img_file_buffer) | |
| image_np = np.array(image) | |
| resized_image = cv2.resize(image_np, (640, 640)) | |
| resized_image = resized_image.astype(np.uint8) | |
| cv2.imwrite('odometer_image.jpg', resized_image) | |
| # detect( | |
| # weights='yolov5\runs\train\exp\weights\best.pt', | |
| # source='odometer_image.jpg', | |
| # img=640, | |
| # conf=0.4, | |
| # name='temp_exp', | |
| # hide_labels=True, | |
| # hide_conf=True, | |
| # save_txt=True, | |
| # exist_ok=True | |
| # ) | |
| # os.system('wandb disabled') | |
| os.chdir(YOLO_PATH) | |
| # try: | |
| # shutil.rmtree('runs/detect/temp_exp') | |
| # except: | |
| # pass | |
| image_path = "../odometer_image.jpg" | |
| # command = f"python detect.py --weights {MODEL_PATH} --source {image_path} --img 640 --conf 0.4 --name 'temp_exp' --hide-labels --hide-conf --save-txt --exist-ok" | |
| command = f''' | |
| python detect.py \ | |
| --weights {MODEL_PATH} \ | |
| --source {image_path} \ | |
| --img 640 \ | |
| --conf 0.4 \ | |
| --name temp_exp \ | |
| --hide-labels \ | |
| --hide-conf \ | |
| --save-txt \ | |
| --exist-ok \ | |
| --save-conf | |
| ''' | |
| # Run the command | |
| os.system(command) | |
| # st.write('The detection is completed!!!') | |
| os.chdir(CUR_DIR) | |
| # st.write(os.path.exists('yolov5/runs/detect/temp_exp')) | |
| if os.path.exists('yolov5/runs/detect/temp_exp'): | |
| processed_image = cv2.imread('yolov5/runs/detect/temp_exp/odometer_image.jpg') | |
| # st.write('Image boxed and loaded') | |
| text_files = os.listdir('yolov5/runs/detect/temp_exp/labels') | |
| original_img = cv2.imread('odometer_image.jpg') | |
| gray = cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY) | |
| if len(text_files) == 0: | |
| display_text = "An odometer is not detected in the image!!!" | |
| else: | |
| text_file_path = f'yolov5/runs/detect/temp_exp/labels/{text_files[0]}' | |
| x1, y1, x2, y2 = yolov5_to_image_coordinates(text_file_path) | |
| # st.write(x1, y1, x2, y2) | |
| cropped_image = gray[x1:x2, y1:y2] | |
| reader = easyocr.Reader(['en']) | |
| result = reader.readtext(cropped_image) | |
| if len(result) != 0: | |
| odometer_value = sorted(result, key=lambda x: x[2], reverse=True)[0][1] | |
| display_text = f"Odometer value: {odometer_value}" | |
| else: | |
| odometer_value = 'not detected' | |
| display_text = f"The odometer value is {odometer_value}!!!" | |
| else: | |
| display_text = "An odometer is not detected in the image!!!" | |
| processed_image = cv2.imread('odometer_image.jpg') | |
| try: | |
| shutil.rmtree('odometer_image.jpg') | |
| shutil.rmtree('runs/detect/temp_exp') | |
| except: | |
| pass | |
| # Resize or preprocess the image as needed for your model | |
| # For example, resizing to a specific input size | |
| # processed_image = cv2.resize(image_np, (224, 224)) | |
| # Process the image using your deep learning model | |
| # processed_image = process_image(image_np) | |
| # Display the processed image | |
| st.image(processed_image, caption=f"{display_text}", use_column_width=True) | |
| st.session_state.pop("odometer") | |
| if __name__ == "__main__": | |
| main() |