add a way to test this tflite

coco_labels.json

@@ -0,0 +1,82 @@
+{
+    "0": "person",
+    "1": "bicycle",
+    "2": "car",
+    "3": "motorcycle",
+    "4": "airplane",
+    "5": "bus",
+    "6": "train",
+    "7": "truck",
+    "8": "boat",
+    "9": "traffic light",
+    "10": "fire hydrant",
+    "11": "stop sign",
+    "12": "parking meter",
+    "13": "bench",
+    "14": "bird",
+    "15": "cat",
+    "16": "dog",
+    "17": "horse",
+    "18": "sheep",
+    "19": "cow",
+    "20": "elephant",
+    "21": "bear",
+    "22": "zebra",
+    "23": "giraffe",
+    "24": "backpack",
+    "25": "umbrella",
+    "26": "handbag",
+    "27": "tie",
+    "28": "suitcase",
+    "29": "frisbee",
+    "30": "skis",
+    "31": "snowboard",
+    "32": "sports ball",
+    "33": "kite",
+    "34": "baseball bat",
+    "35": "baseball glove",
+    "36": "skateboard",
+    "37": "surfboard",
+    "38": "tennis racket",
+    "39": "bottle",
+    "40": "wine glass",
+    "41": "cup",
+    "42": "fork",
+    "43": "knife",
+    "44": "spoon",
+    "45": "bowl",
+    "46": "banana",
+    "47": "apple",
+    "48": "sandwich",
+    "49": "orange",
+    "50": "broccoli",
+    "51": "carrot",
+    "52": "hot dog",
+    "53": "pizza",
+    "54": "donut",
+    "55": "cake",
+    "56": "chair",
+    "57": "couch",
+    "58": "potted plant",
+    "59": "bed",
+    "60": "dining table",
+    "61": "toilet",
+    "62": "tv",
+    "63": "laptop",
+    "64": "mouse",
+    "65": "remote",
+    "66": "keyboard",
+    "67": "cell phone",
+    "68": "microwave",
+    "69": "oven",
+    "70": "toaster",
+    "71": "sink",
+    "72": "refrigerator",
+    "73": "book",
+    "74": "clock",
+    "75": "vase",
+    "76": "scissors",
+    "77": "teddy bear",
+    "78": "hair drier",
+    "79": "toothbrush"
+}

test_yolob8_tflite.py

@@ -0,0 +1,205 @@
+#!/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)