Spaces:
Running
Running
File size: 7,752 Bytes
a69b646 40e5433 a69b646 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | 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)
|