fox_detect_gradio.py

import os
import cv2
import gradio as gr
from tensorflow.lite.python.interpreter import Interpreter

from utils import (get_labels, 
                    parse_image_for_detection, 
                    resize_image, 
                    normalize_image, 
                    save_image, 
                    get_random_images)

current_dir = os.path.dirname(__file__)
MODEL_PATH = os.path.join(current_dir, 'model', 'model.tflite')
LABEL_PATH = os.path.join(current_dir, 'model', 'labels.txt')
IMAGES_DIRPATH = os.path.join(current_dir, 'images')
OUTPUT_DIR = os.path.join(current_dir, 'output')

example_image_list = [
    "image_0012.png", ## fox mid grass dark
    "image_0026.png", ## fox hidden
    "image_0010.png", ## fox dark path NICE
    "image_0023.png", ## costume mid garden facing
    "image_0022.png", ## costume mid garden bent over
    "image_0024.png", ## costume closeup
    "image_0027.png", ## fox color
    "image_0018.png", ## fox dark path far
    "image_0011.png", ## fox dark path
    "image_0013.png", ## fox dark terrace bright
    "image_0014.png", ## costume mid garden happy
    "image_0020.png", ## costume far away
    "image_0025.png", ## costume far away
    "image_0015.png", ## fox dark terrace left
    "image_0016.png", ## fox dark path
    "image_0021.png", ## fox dark bottom
    "image_0028.png", ## fox dark path
    "image_0019.png", ## person
    "image_0017.png", ## paddington and ball
]

def detect(modelpath, img, labels_filepath, min_conf=0.5, output_dir='/content/output'):
    # Load the Tensorflow Lite model into memory ----------------
    interpreter = Interpreter(model_path=modelpath)
    interpreter.resize_tensor_input(0, [1, 320, 320, 3])
    interpreter.allocate_tensors()
    # Get model details -----------------------------------------
    input_details = interpreter.get_input_details()
    detect_height = input_details[0]['shape'][1]
    detect_width = input_details[0]['shape'][2]
    # Get model details -----------------------------------------

    ## load the labels and parse the image for detection --------
    labels = get_labels(labels_filepath)
    img, image_width, image_height = parse_image_for_detection(img)
    np_image = resize_image(img, detect_width, detect_height)
    np_image = normalize_image(np_image, interpreter)
    # Perform the actual detection by running the model with the image as input
    tensor_index = input_details[0]['index']
    interpreter.set_tensor(tensor_index, np_image)
    interpreter.invoke()
    ## ----------------------------------------------------------

    ## --- now get the boxes, classes and scores from the detection
    boxes, classes, scores = distill_detections(interpreter)
    img = draw_detections(img, boxes, classes, scores, image_height, image_width, labels, min_conf)
    output_image = save_image(img, output_dir)

    return output_image

def distill_detections(interpreter):
    """ receives the already invoked interpreter and returns the boxes, classes and scores
    """
    output_details = interpreter.get_output_details()

    boxes_index = 1
    classes_index = 3 
    scores_index = 0 

    # Retrieve detection results
    boxes = interpreter.get_tensor(output_details[boxes_index]['index'])[0] # Bounding box coordinates of detected objects 1
    classes = interpreter.get_tensor(output_details[classes_index]['index'])[0] # Class index of detected objects 3
    scores = interpreter.get_tensor(output_details[scores_index]['index'])[0] # Confidence of detected objects 0

    return boxes, classes, scores

def draw_detections(image, boxes, classes, scores, image_height, image_width, labels, min_conf):
    """ receives the original image, the detected boxes, classes and scores. 
    and draws the bounding boxes with labels on the image.
    """
    # Loop over all detections and draw detection box if confidence is above minimum threshold
    for i in range(len(scores)):
        if ((scores[i] > min_conf) and (scores[i] <= 1.0)):
            # Get bounding box coordinates and draw box
            # Interpreter can return coordinates that are outside of image dimensions, need to force them to be within image using max() and min()
            ymin = int(max(1,(boxes[i][0] * image_height)))
            xmin = int(max(1,(boxes[i][1] * image_width)))
            ymax = int(min(image_height,(boxes[i][2] * image_height)))
            xmax = int(min(image_width,(boxes[i][3] * image_width)))
            
            ## draw a bounding box around the detected object
            cv2.rectangle(image, (xmin,ymin), (xmax,ymax), (10, 255, 0), 2)

            ## now lets draw the label above the bounding box. 
            object_name = labels[int(classes[i])]
            label = '%s: %d%%' % (object_name, int(scores[i]*100)) # Example: 'person: 72%'
            labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
            label_ymin_base = max(ymin, labelSize[1] + 10) 
            ## draw the rectangle
            label_xmin = xmin
            label_ymin = label_ymin_base-labelSize[1]-10
            label_xmax = xmin+labelSize[0]
            label_ymax = label_ymin_base+baseLine-10
            ## draw a white rectangle to put the label text into
            cv2.rectangle(image, ## image
                          (label_xmin, label_ymin), ## top left
                          (label_xmax, label_ymax), ## bottom right
                          (255, 255, 255), ## color
                          cv2.FILLED) 
            ## write the label text
            text_xmin = xmin
            text_ymin = label_ymin_base-7
            cv2.putText(image, ## image
                        label, ## str
                        (text_xmin, text_ymin), 
                        cv2.FONT_HERSHEY_SIMPLEX, 
                        0.7, ## font scale
                        (0, 0, 0),  ## color
                        2) ## thickness
    return image

def gradio_entry(image, confidence=0.1):
    """ entry point for the gradio interface to run the detection"""
    output_filepath = detect(MODEL_PATH, image, LABEL_PATH, min_conf=confidence, output_dir=OUTPUT_DIR)
    return output_filepath

def main(debug=False):
    if debug:
        img = get_random_images(IMAGES_DIRPATH, 10)[0]
        output_filepath = detect(MODEL_PATH, img, LABEL_PATH, min_conf=0.5, output_dir=OUTPUT_DIR)
        os.system(f'open {output_filepath}')
        return

    default_conf = 0.2
    examples_for_display = []
    examples_for_full = []
    for img in example_image_list:
        img_path = os.path.join(IMAGES_DIRPATH, img)
        examples_for_full.append([img_path, 0.2])
        examples_for_display.append([img_path])

    input_image = gr.Image(width=800, height=600,
                                label='Input Image')
    input_slider_conf = gr.Slider(value=default_conf,
                                minimum=0.0, maximum=1.0, step=0.01,
                                label="Confidence",
                                info="Minimum confidence threshold")
    output_image= gr.Image(width=800, height=600,
                                label="Output Image")
    input_widgets = [input_image,
                    input_slider_conf]
    interface = gr.Interface(fn=gradio_entry,
                            inputs=input_widgets,
                            outputs=output_image,
                            examples=examples_for_display,
                            examples_per_page=18)
    ## now add event handler, so whenever we set the slider value, the full example is selected
    interface.load_data = lambda i: examples_for_full[i] ## loads both the image and the confidence

    interface.launch()

if __name__ == '__main__':
    main(debug=False)