the digit extractor model changed

app.py

@@ -49,12 +49,15 @@ def main():
         else:
             cropped_image = gray[y1:y2, x1:x2]
             # cropped_image = resized_image[y1:y2, x1:x2]
+            cropped_image = cv2.resize(cropped_image, (640, 640))
             cv2.imwrite('odometer_number_image.jpg', cropped_image)
+            
             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)
 

digit_yolov8_best_float16.tflite

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

models.py

@@ -85,9 +85,9 @@ def get_digit(model_path, image_path, threshold=0.5):
     # 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_resized = np.resize(image, (image_width, image_height, 3))
 
-    image_np = np.array(image_resized) #
+    image_np = np.array(image) #
     image_np = np.true_divide(image_np, 255, dtype=np.float32) 
     image_np = image_np[np.newaxis, :]
 
@@ -101,25 +101,58 @@ def get_digit(model_path, image_path, threshold=0.5):
     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.
+    scores = np.max(output[:, 4:], axis=1) #Get score value, 5th column of output tensor
+    classes = np.argmax(output[:, 4:], 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 = []
 
+    prob_threshold = threshold
+
     for box, score, cls in zip(boxes_xywh, scores, classes):
+
+        if score < prob_threshold:
+            continue
+
         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]
-    pred_list = [i for i in pred_list if i[0]>0 and i[1]>0]
-    sorted_number_list = sorted(pred_list, key=lambda x: x[0])
+        pred_list.append((x1, x2, cls, score))
+
+    pred_list = sorted(pred_list, key=lambda x: x[0])
+
+    num_list = []
+
+    temp_pred_list =[]
+
+    x_prev = 0
+
+    x_diff = min([elem[1] - elem[0] for elem in pred_list]) - 10
+
+    for idx, pred in enumerate(pred_list):
+    
+        if idx == 0:
+            temp_pred_list.append(pred)
+            x_prev = pred[0]
+        elif idx == len(pred_list) - 1:
+            temp_final_num = sorted(temp_pred_list, key=lambda x: x[-1], reverse=True)[0]
+            num_list.append(temp_final_num)
+        elif pred[0] - x_prev < x_diff:
+            temp_pred_list.append(pred)
+            x_prev = pred[0]
+        else:
+            temp_final_num = sorted(temp_pred_list, key=lambda x: x[-1], reverse=True)[0]
+            num_list.append(temp_final_num)
+            temp_pred_list = []
+            x_prev = pred[0]
+            temp_pred_list.append(pred)
+
+    sorted_number_list = sorted(num_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 = 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[:10]])
 
     return output_digit