digit extractor added

app.py

@@ -9,7 +9,7 @@ from PIL import Image, ImageDraw
 # import platform
 # from xyxy_converter import yolov5_to_image_coordinates
 # import shutil
-from models import get_odometer_xy
+from models import get_odometer_xy, get_digit
 
 # system_platform = platform.system()
 # if system_platform == 'Windows': pathlib.PosixPath = pathlib.WindowsPath
@@ -48,8 +48,14 @@ def main():
             st.image('odometer_image.jpg', caption=f"{display_text}", use_column_width=True)
         else:
             cropped_image = gray[y1:y2, x1:x2]
+            # cropped_image = resized_image[y1:y2, x1:x2]
             cv2.imwrite('odometer_number_image.jpg', cropped_image)
-            display_text = 'Here is the zoomed odometer value'
+            extracted_digit = get_digit(
+                model_path="digit_yolov8_best_float16.tflite", 
+                image_path='odometer_number_image.jpg', 
+                threshold=0.4
+            )
+            display_text = f'Here is the zoomed odometer value: {extracted_digit}.'
             st.image('odometer_number_image.jpg', caption=f"{display_text}", use_column_width=True)
 
             image = Image.open('odometer_image.jpg')

digit_yolov8_best_float16.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64673733adce09ce682b4527410ece79f62108e33cea42e8b14166e695e67d55
+size 6195116

models.py

@@ -1,6 +1,7 @@
 import tensorflow as tf
 import numpy as np
 from PIL import Image
+import cv2
 
 def get_odometer_xy(model_path, image_path):
     #model_path = 'odo_detector.tflite'
@@ -67,4 +68,59 @@ def get_odometer_xy(model_path, image_path):
 
     return x1, y1, x2, y2, final_score
 
+def get_digit(model_path, image_path, threshold=0.5):
+
+    interpreter = tf.lite.Interpreter(model_path=model_path)
+    interpreter.allocate_tensors()
+
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+
+    # Obtain the height and width of the corresponding image from the input tensor
+    image_height = input_details[0]['shape'][1] # 640
+    image_width = input_details[0]['shape'][2] # 640
+
+    # Image Preparation
+    # image_name = 'car.jpg'
+    # image = Image.open(image_path2)
+    # image_resized = image.resize((image_width, image_height)) # Resize the image to the corresponding size of the input tensor and store it in a new variable
+    image = cv2.imread(image_path)
+    image_resized = np.resize(image, (image_width, image_height, 3))
+
+    image_np = np.array(image_resized) #
+    image_np = np.true_divide(image_np, 255, dtype=np.float32) 
+    image_np = image_np[np.newaxis, :]
+
+    # inference
+    interpreter.set_tensor(input_details[0]['index'], image_np)
+    interpreter.invoke()
+
+    # Obtaining output results
+    output = interpreter.get_tensor(output_details[0]['index'])
+    output = output[0]
+    output = output.T
+
+    boxes_xywh = output[:, :4] #Get coordinates of bounding box, first 4 columns of output tensor
+    scores = output[:, 4]#np.max(output[..., 5:], axis=1) #Get score value, 5th column of output tensor
+    classes = np.argmax(output[:, 5:], axis=1) # Get the class value, get the 6th and subsequent columns of the output tensor, and store the largest value in the output tensor.
+
+    pred_list = []
+
+    for box, score, cls in zip(boxes_xywh, scores, classes):
+        x_center, y_center, width, height = box
+        x1 = int((x_center - width / 2) * image_width)
+        y1 = int((y_center - height / 2) * image_height)
+        x2 = int((x_center + width / 2) * image_width)
+        y2 = int((y_center + height / 2) * image_height)
+        pred_list.append(((x1, x2, cls, score)))
+
+    pred_list = [i for i in pred_list if i[-1] > threshold]
+    sorted_number_list = [i for i in sorted_number_list if i[0]>0 and i[1]>0]
+    sorted_number_list = sorted(pred_list, key=lambda x: x[0])
+    # sorted_number_list = sorted(sorted_number_list, reverse=True, key= lambda x: x[-1])
+    # output_digit = float(''.join([str(int(i[2])) if i[2]!=10 else '.' for i in sorted_number_list]))
+    output_digit = ''.join([str(int(i[2])) if i[2]!=10 else '.' for i in sorted_number_list])
+    # output_digit = ''.join([str(int(i[2])) if i[2]!=10 else '.' for i in sorted_number_list[:10]])
+
+    return output_digit