def test_draw_bounding_box_on_image(self):
test_image = self.create_colorful_test_image()
test_image = Image.fromarray(test_image)
width_original, height_original = test_image.size
ymin = 0.25
ymax = 0.75
xmin = 0.4
xmax = 0.6
visualization_utils.draw_bounding_box_on_image(test_image, ymin, xmin,
ymax, xmax)
width_final, height_final = test_image.size
self.assertEqual(width_original, width_final)
self.assertEqual(height_original, height_final)
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices(device_type="GPU")
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(device=gpu, enable=True)
# load model
yolo_v3 = tf.saved_model.load("saved_model/final")
files_list = list_files("test_data", validExts="jpg")
for file in files_list:
image = Image.open(file)
img_tensor = pre_process_image(image)
img_tensor = img_tensor / 255.0
boxes, scores, classes = yolo_v3(img_tensor, training=False)
final_boxes = post_process_image(boxes, image.size)
for box, scs, cls in zip(final_boxes.numpy(), scores.numpy(),
classes.numpy()):
display_str = f"{list(PASCAL_VOC_CLASSES.keys())[list(PASCAL_VOC_CLASSES.values()).index(cls + 1)]}:{tf.round(100 * scs)}%"
draw_bounding_box_on_image(image,
box[1],
box[0],
box[3],
box[2],
color=SELECTED_COLORS[cls],
display_str=display_str)
plt.imshow(image)
plt.show()
def drawBox(image, box): image_np = image.copy() vis_util.draw_bounding_box_on_image(image_np, box[0], box[1], box[2], box[3], thickness=2) return image_np