Spaces:
Sleeping
Sleeping
File size: 4,163 Bytes
8b6b2e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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
'''
# 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}!!!"
try:
shutil.rmtree('odometer_image.jpg')
except:
pass
else:
processed_image = image_np
# 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() |