Search code examples
pythontensorflowopencvtensorflow-lite

How to Resize the Tensorflow Output Bounding Box to Original Image


So, I am working on a Model, that I have successfully trained, it is a .tflite model and uses Tensorflow Lite. I am using Python 3 as my interface with tensorflow and I am unable to add a bounding box from the returned image.

My Question is:

How to Resize the Tensorflow Output Bounding Box to Original Image

I am able to get the output for the 512x512 input file, but I am unable to go ahead and get a output for the original file, if I am correct. How to go ahead and re-size that output from my model in order to go ahead and crop that part from the original image and then save it.

Image for AI

Code

import tensorflow as tf
import numpy as np
import cv2
import pathlib
import os
from silence_tensorflow import silence_tensorflow
from PIL import Image

silence_tensorflow()
interpreter = tf.lite.Interpreter(model_path="model.tflite")

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

#print(input_details)
#print(output_details)

interpreter.allocate_tensors()

def draw_rect(image, box):
    h, w, c = image.shape
    y_min = int(max(1, (box[0] * h)))
    x_min = int(max(1, (box[1] * w)))
    y_max = int(min(h, (box[2] * h)))
    x_max = int(min(w, (box[3] * w)))

    # draw a rectangle on the image
    cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (13, 13, 13), 2)

for file in pathlib.Path('./').iterdir():

    if file.suffix != '.jpeg' and file.suffix != '.png':
        continue

    img = cv2.imread(r"{}".format(file.resolve()))
    print(f'[Converting] {file.resolve()}')
    new_img = cv2.resize(img, (512, 512))

    interpreter.set_tensor(input_details[0]['index'], [new_img])

    interpreter.invoke()
    rects = interpreter.get_tensor(
        output_details[0]['index'])

    scores = interpreter.get_tensor(
        output_details[2]['index'])

    for index, score in enumerate(scores[0]):
        if index == 0:
            print('[Quantity]')
            print(index,score)
        if index == 1:
            print('[Barcode]')
            print(index,score)
        if score > 0.2:
          draw_rect(new_img,rects[0][index])
          print(rects[0][index])

    cv2.imshow("image", new_img)
    cv2.waitKey(0)


Solution

  • Solution

    The Most easy solution for this problem is to go ahead and find the factor at which h, w, c of the downscaled image are compared to h, w, c of the original image.

    That can be understood using the following piece of code.

    h, w, c = image.shape # old image 512x512
    h1, w1, c1 = image2.shape # new image resize
    h = (h1/h) * h
    w = (w1/w) * w
    c = (c1/c) * c
    

    The new formed h, w, c values are belonging to the image2 or the Original Image.

    Draw Rectangle

    In order to go ahead and draw a rectangle we are going to go ahead and use the following approach on this problem in order to go ahead and make it viable.

    def draw_rect(image,image2, box):
        h, w, c = image.shape # old image 512x512
        h1, w1, c1 = image2.shape # new image resize
        h = (h1/h) * h
        w = (w1/w) * w
        c = (c1/c) * c
        y_min = int(max(1, (box[0] * h)))
        x_min = int(max(1, (box[1] * w)))
        y_max = int(min(h, (box[2] * h)))
        x_max = int(min(w, (box[3] * w)))
    
        # draw a rectangle on the image
        cv2.rectangle(image2, (x_min, y_min), (x_max, y_max), (13, 13, 13), 2)