Spaces:
Sleeping
Sleeping
File size: 2,580 Bytes
5482479 |
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 |
import tensorflow as tf
import numpy as np
from PIL import Image
def get_odometer_xy(model_path, image_path):
#model_path = 'odo_detector.tflite'
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_path)
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_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.zeros(len(scores))#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.
# Threshold Setting
# threshold = 0.7
final_score = 0
x_center, y_center, width, height = 0, 0, 0, 0
class_name = 'odometer'
# Bounding boxes, scores, and classes are drawn on the image
# draw = ImageDraw.Draw(image_resized)
for box, score, cls in zip(boxes_xywh, scores, classes):
if score >= final_score:
x_center, y_center, width, height = box
final_score = score
class_name = cls
else:
pass
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)
# draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
# text = f"Class: {class_name}, Score: {final_score:.2f}"
# draw.text((x1, y1), text, fill="red")
# Saving Images
# image_resized.save('test_img.jpg')
return x1, y1, x2, y2, final_score
|