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)