Upload 3 files

fox_detect_gradio.py

@@ -0,0 +1,174 @@
+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, 'publish_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)
+    
+

requirements.txt

@@ -0,0 +1,7 @@
+jupyterlab
+jupyterlab-vim
+ipywidgets
+tensorflow
+opencv-python
+matplotlib
+gradio

utils.py

@@ -0,0 +1,108 @@
+import os
+import cv2
+import numpy as np
+# import sys
+import glob
+import gradio as gr
+import random
+# import importlib.util
+import datetime
+# from tensorflow.lite.python.interpreter import Interpreter
+
+# import matplotlib
+import matplotlib.pyplot as plt
+
+### ---------------------------- image utils ---------------------------------
+def parse_image_for_detection(img):
+    """
+    if img comes from gradio, it makes sure its a numpy array
+    if img is a file path, it reads it and converts from BGR to RGB
+    it also returns the width and height of the image
+    """
+    if isinstance(img, str):
+        ## if its a file path, we read it and convert from BGR to RGB
+        image = cv2.imread(img)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    else:
+        ## otherwise assume its a numpy array from Gradio UI.
+        ## but make sure that it actually is.
+        if not isinstance(img, np.ndarray):
+            img = np.array(img)
+        image = img
+
+    ## lets also get the width and height of the original image 
+    image_height, image_width, _ = image.shape
+
+    return image, image_width, image_height
+
+def resize_image(np_image, width, height):
+    image_resized = cv2.resize(np_image, (width, height))
+    np_image = np.expand_dims(image_resized, axis=0)
+    return np_image
+
+def normalize_image(np_image, interpreter):
+    ## check if the model expects a floating point input 
+    is_model_float = (interpreter.get_input_details()[0]['dtype'] == np.float32)
+
+    ## Normalize pixel values if using a floating model (i.e. if model is non-quantized)
+    if is_model_float:
+        input_mean = 256.0 / 2.0
+        input_std = 256.0 / 2.0
+        np_image = (np.float32(np_image) - input_mean) / input_std
+    return np_image
+
+def save_image(image, output_dir, output_width=1600, output_height=1200, dpi=80):
+    """ 
+    saves the image in the output dir, as a matplotlib figure 
+    """
+    ## make sure output directory exists
+    os.makedirs(output_dir, exist_ok=True)
+
+    ## first get the figsize in inches based on pixel output width, height
+    figsize = get_figsize_from_pixels(output_width, output_height, dpi=dpi)
+
+    ## now plot the image
+    plt.figure(figsize=figsize)
+    plt.imshow(image)
+    plt.tight_layout(pad=3)
+
+    ## generate an output filename with a timestamp
+    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')
+    output_filename = os.path.join(output_dir, f'img_{timestamp}.png')
+    ## save the figure with the output filename
+    plt.savefig(output_filename, dpi=dpi)
+    return output_filename
+### ---------------------------- image utils ---------------------------------
+
+
+### ---------------------------- basic utils ---------------------------------
+def get_labels(labels_filepath):
+    with open(labels_filepath, 'r') as f:
+        labels = [line.strip() for line in f.readlines()]
+    return labels
+
+def get_random_images(dirpath, image_count=10):
+    """ returns a list of random image filepaths from the dirpath """
+    images = glob.glob(dirpath + '/*.jpg') + \
+                glob.glob(dirpath + '/*.JPG') + \
+                glob.glob(dirpath + '/*.png') + \
+                glob.glob(dirpath + '/*.bmp')
+
+    # img_filepaths = random.sample(images, image_count)
+    # return img_filepaths
+    return sorted(images)
+
+def get_figsize_from_pixels(width, height, dpi=80):
+    """ returns the width and height in inches based on the dpi
+    used for matplotlib figures
+    """
+    width_in = width / dpi
+    height_in = height / dpi
+    return (width_in, height_in)
+### ---------------------------- basic utils ---------------------------------
+
+
+
+
+
+