scripts/test_yolob8_tflite.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct  4 16:44:12 2023

@author: lin
"""
import glob
import sys
sys.path.append('../../..')
import os
import cv2
import json
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
# from utils.bbox_op import non_max_supression

def one_multiple_iou(box, boxes, box_area, boxes_area):
    """
    Compute the intersection over union. 1 to multiple
    Inputs:
        box:   numpy array with 1 box, ymin, xmin, ymax, xmax
        boxes: numpy array with shape [N, 4] holding N boxes

    Outputs:
        a numpy array with shape [N*1] representing box areas
    """

    # this is the iou of the box against all other boxes
    assert boxes.shape[0] == boxes_area.shape[0]

    ymin = np.maximum(box[0], boxes[:, 0])  # bottom
    xmin = np.maximum(box[1], boxes[:, 1])  # left
    ymax = np.minimum(box[2], boxes[:, 2])  # top
    xmax = np.minimum(box[3], boxes[:, 3])  # rifht

    # we ignore areas where the intersection side would be negative
    # this is done by using maxing the side length by 0
    intersections = np.maximum(ymax - ymin, 0) * np.maximum(xmax - xmin, 0)
    # each union is then the box area
    # added to each other box area minusing their intersection calculated above
    unions = box_area + boxes_area - intersections
    # element wise division
    # if the intersection is 0, then their ratio is 0
    ious = intersections / unions
    return ious
def select_non_overlapping_bboxes(boxes, scores, iou_th):
    ymin = boxes[:, 0]
    ymax = boxes[:, 2]
    xmin = boxes[:, 1]
    xmax = boxes[:, 3]

    # box coordinate ranges are inclusive-inclusive
    areas = (ymax - ymin) * (xmax - xmin)
    scores_indexes = list(np.argsort(scores))
    keep_idx = []
    while len(scores_indexes) > 0:
        index = scores_indexes.pop()
        keep_idx.append(index)

        ious = one_multiple_iou(
            boxes[index], boxes[scores_indexes], areas[index], areas[scores_indexes]
        )
        filtered_indexes = set((ious > iou_th).nonzero()[0])

        scores_indexes = [
            v for (i, v) in enumerate(scores_indexes) if i not in filtered_indexes
        ]
    return keep_idx
def non_max_supression(boxes, scores, classes, iou_th):
    """
    remover overlaped boundingboxes. Starting by the box with the highest score
    if the iou is greater than the threshold, remove it, else keep it.
    Inputs:
        boxes： numpy array with shape [N, 4] holding N boxes。 [ymin, xmin, ymax, xmax]
        scores： numpy array with shape [N, 1] holding the prediction score of each box
        classes: numpy array with shape [N, 1] holding the class that each box belongs
        iou_th: intersection over union threshold to consider the overlapping boxes have detect 2 objects
    Output:
        boxes, scores, classes with intersection over union ratio less than the threshold.

    """
    #    assert boxes.shape[0] == scores.shape[0]
    if len(scores) == 0:
        return boxes, scores, classes
    keep_idx = select_non_overlapping_bboxes(boxes, scores, iou_th)

    return boxes[keep_idx], scores[keep_idx], classes[keep_idx]

def preprocess(img_path):
    image_np = cv2.imread(img_path)

    image_np = center_crop(image_np)

    image_np = cv2.resize(image_np, (640, 640))
    #image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
    image_np = image_np.astype(float)
    image_np /= 255.0
    return image_np

def center_crop(img):
    width, height = img.shape[1], img.shape[0]

    crop_size = width if width < height else height

    mid_x, mid_y = int(width/2), int(height/2)

    cs2 = int(crop_size/2)
    crop_img = img[mid_y-cs2:mid_y+cs2, mid_x-cs2:mid_x+cs2]

    return crop_img

def postprocess_prediction(preds):

    bboxes = preds[0][:4]
    class_prob = preds[0, 4:]
    classes = np.argmax(class_prob, axis=0)
    scores = np.max(class_prob, axis=0)
    # filter by threshold
    valid_idx = np.where(scores>=min_th)[0]
    bboxes = bboxes[:, valid_idx]
    classes = classes[valid_idx]
    scores = scores[valid_idx]
    bboxes = bboxes.transpose()
    bboxes = bboxes*640
    xmin = bboxes[:,0]-bboxes[:,2]//2
    xmax = bboxes[:,0]+bboxes[:,2]//2
    ymin = bboxes[:,1]-bboxes[:,3]//2
    ymax = bboxes[:,1]+bboxes[:,3]//2
    xmin = np.clip(xmin, 0, 640)
    ymin = np.clip(ymin, 0, 640)

    bboxes = np.vstack([ymin, xmin, ymax, xmax])
    bboxes = bboxes.transpose()
    bboxes = bboxes.astype(int)

    bboxes, scores, classes = non_max_supression(bboxes, scores, classes, iou_th=0.5)
    idx = np.argsort(scores)[::-1]
    bboxes = bboxes[idx]
    classes = classes[idx]
    scores = scores[idx]
    return bboxes, classes, scores

def plot_prediction(image_np, bboxes, classes, scores, label_map):
    color=(255,0,0)
    thickness=5
    font_scale=3

    for i, box in enumerate(bboxes):
        box = bboxes[i, :]
        
        ymin, xmin, ymax, xmax = box
        
        image_np = cv2.rectangle(image_np, (xmin, ymin), (xmax, ymax), color=color, thickness=thickness)
        text_x = xmin - 10 if xmin > 20 else xmin + 10
        text_y = ymin - 10 if ymin > 20 else ymin + 10
        display_str = label_map[str(classes[i])]

        cv2.putText(
            image_np,
            display_str,
            (text_x, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            color,
            thickness,
        )
    plt.imshow(image_np)
    plt.show()

def predict_yolo_tflite(intenpreter, image_np):
    input_tensor = np.expand_dims(image_np, axis=0)
    input_tensor = tf.convert_to_tensor(input_tensor, dtype=tf.float32)

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

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

if __name__ == "__main__":
    min_th = 0.1
    labels_json = "coco_labels.json"
    with open(labels_json) as f:
        label_map = json.load(f)
    img_path = "test_images"
    saved_tflite = "tflite_model.tflite"
    # load model
    interpreter = tf.lite.Interpreter(model_path=saved_tflite)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    print(input_details)
    print(output_details)
    images = glob.glob(os.path.join(img_path, "*"))
    for img in images:
        image_np = preprocess(img)
        print(image_np.shape)
    
        # image_np = np.array(Image.open(image_path))
        preds = predict_yolo_tflite(interpreter, image_np)
        bboxes, classes, scores = postprocess_prediction(preds)
    
        plot_prediction(image_np, bboxes, classes, scores, label_map)