pixelprotest/fox-robot

1import os2import cv23import gradio as gr4from tensorflow.lite.python.interpreter import Interpreter5 6from utils import (get_labels, 7                    parse_image_for_detection, 8                    resize_image, 9                    normalize_image, 10                    save_image, 11                    get_random_images)12 13current_dir = os.path.dirname(__file__)14MODEL_PATH = os.path.join(current_dir, 'model', 'model.tflite')15LABEL_PATH = os.path.join(current_dir, 'model', 'labels.txt')16IMAGES_DIRPATH = os.path.join(current_dir, 'images')17OUTPUT_DIR = os.path.join(current_dir, 'output')18 19example_image_list = [20    "image_0012.png", ## fox mid grass dark21    "image_0026.png", ## fox hidden22    "image_0010.png", ## fox dark path NICE23    "image_0023.png", ## costume mid garden facing24    "image_0022.png", ## costume mid garden bent over25    "image_0024.png", ## costume closeup26    "image_0027.png", ## fox color27    "image_0018.png", ## fox dark path far28    "image_0011.png", ## fox dark path29    "image_0013.png", ## fox dark terrace bright30    "image_0014.png", ## costume mid garden happy31    "image_0020.png", ## costume far away32    "image_0025.png", ## costume far away33    "image_0015.png", ## fox dark terrace left34    "image_0016.png", ## fox dark path35    "image_0021.png", ## fox dark bottom36    "image_0028.png", ## fox dark path37    "image_0019.png", ## person38    "image_0017.png", ## paddington and ball39]40 41def detect(modelpath, img, labels_filepath, min_conf=0.5, output_dir='/content/output'):42    # Load the Tensorflow Lite model into memory ----------------43    interpreter = Interpreter(model_path=modelpath)44    interpreter.resize_tensor_input(0, [1, 320, 320, 3])45    interpreter.allocate_tensors()46    # Get model details -----------------------------------------47    input_details = interpreter.get_input_details()48    detect_height = input_details[0]['shape'][1]49    detect_width = input_details[0]['shape'][2]50    # Get model details -----------------------------------------51 52    ## load the labels and parse the image for detection --------53    labels = get_labels(labels_filepath)54    img, image_width, image_height = parse_image_for_detection(img)55    np_image = resize_image(img, detect_width, detect_height)56    np_image = normalize_image(np_image, interpreter)57    # Perform the actual detection by running the model with the image as input58    tensor_index = input_details[0]['index']59    interpreter.set_tensor(tensor_index, np_image)60    interpreter.invoke()61    ## ----------------------------------------------------------62 63    ## --- now get the boxes, classes and scores from the detection64    boxes, classes, scores = distill_detections(interpreter)65    img = draw_detections(img, boxes, classes, scores, image_height, image_width, labels, min_conf)66    output_image = save_image(img, output_dir)67 68    return output_image69 70def distill_detections(interpreter):71    """ receives the already invoked interpreter and returns the boxes, classes and scores72    """73    output_details = interpreter.get_output_details()74 75    boxes_index = 176    classes_index = 3 77    scores_index = 0 78 79    # Retrieve detection results80    boxes = interpreter.get_tensor(output_details[boxes_index]['index'])[0] # Bounding box coordinates of detected objects 181    classes = interpreter.get_tensor(output_details[classes_index]['index'])[0] # Class index of detected objects 382    scores = interpreter.get_tensor(output_details[scores_index]['index'])[0] # Confidence of detected objects 083 84    return boxes, classes, scores85 86def draw_detections(image, boxes, classes, scores, image_height, image_width, labels, min_conf):87    """ receives the original image, the detected boxes, classes and scores. 88    and draws the bounding boxes with labels on the image.89    """90    # Loop over all detections and draw detection box if confidence is above minimum threshold91    for i in range(len(scores)):92        if ((scores[i] > min_conf) and (scores[i] <= 1.0)):93            # Get bounding box coordinates and draw box94            # Interpreter can return coordinates that are outside of image dimensions, need to force them to be within image using max() and min()95            ymin = int(max(1,(boxes[i][0] * image_height)))96            xmin = int(max(1,(boxes[i][1] * image_width)))97            ymax = int(min(image_height,(boxes[i][2] * image_height)))98            xmax = int(min(image_width,(boxes[i][3] * image_width)))99            100            ## draw a bounding box around the detected object101            cv2.rectangle(image, (xmin,ymin), (xmax,ymax), (10, 255, 0), 2)102 103            ## now lets draw the label above the bounding box. 104            object_name = labels[int(classes[i])]105            label = '%s: %d%%' % (object_name, int(scores[i]*100)) # Example: 'person: 72%'106            labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)107            label_ymin_base = max(ymin, labelSize[1] + 10) 108            ## draw the rectangle109            label_xmin = xmin110            label_ymin = label_ymin_base-labelSize[1]-10111            label_xmax = xmin+labelSize[0]112            label_ymax = label_ymin_base+baseLine-10113            ## draw a white rectangle to put the label text into114            cv2.rectangle(image, ## image115                          (label_xmin, label_ymin), ## top left116                          (label_xmax, label_ymax), ## bottom right117                          (255, 255, 255), ## color118                          cv2.FILLED) 119            ## write the label text120            text_xmin = xmin121            text_ymin = label_ymin_base-7122            cv2.putText(image, ## image123                        label, ## str124                        (text_xmin, text_ymin), 125                        cv2.FONT_HERSHEY_SIMPLEX, 126                        0.7, ## font scale127                        (0, 0, 0),  ## color128                        2) ## thickness129    return image130 131def gradio_entry(image, confidence=0.1):132    """ entry point for the gradio interface to run the detection"""133    output_filepath = detect(MODEL_PATH, image, LABEL_PATH, min_conf=confidence, output_dir=OUTPUT_DIR)134    return output_filepath135 136def main(debug=False):137    if debug:138        img = get_random_images(IMAGES_DIRPATH, 10)[0]139        output_filepath = detect(MODEL_PATH, img, LABEL_PATH, min_conf=0.5, output_dir=OUTPUT_DIR)140        os.system(f'open {output_filepath}')141        return142 143    default_conf = 0.2144    examples_for_display = []145    examples_for_full = []146    for img in example_image_list:147        img_path = os.path.join(IMAGES_DIRPATH, img)148        examples_for_full.append([img_path, 0.2])149        examples_for_display.append([img_path])150 151    input_image = gr.Image(width=800, height=600,152                                label='Input Image')153    input_slider_conf = gr.Slider(value=default_conf,154                                minimum=0.0, maximum=1.0, step=0.01,155                                label="Confidence",156                                info="Minimum confidence threshold")157    output_image= gr.Image(width=800, height=600,158                                label="Output Image")159    input_widgets = [input_image,160                    input_slider_conf]161    interface = gr.Interface(fn=gradio_entry,162                            inputs=input_widgets,163                            outputs=output_image,164                            examples=examples_for_display,165                            examples_per_page=18)166    ## now add event handler, so whenever we set the slider value, the full example is selected167    interface.load_data = lambda i: examples_for_full[i] ## loads both the image and the confidence168 169    interface.launch()170 171if __name__ == '__main__':172    main(debug=False)173    174 175