Skip to content

Instantly share code, notes, and snippets.

@aallan
Created March 8, 2019 19:54
Show Gist options
  • Star 4 You must be signed in to star a gist
  • Fork 3 You must be signed in to fork a gist
  • Save aallan/abcf53741d45b7d4cbfd69c7b65b07a0 to your computer and use it in GitHub Desktop.
Save aallan/abcf53741d45b7d4cbfd69c7b65b07a0 to your computer and use it in GitHub Desktop.
Object detection python demonstration code for use with Google's Edge TPU (now with labelling)
import argparse
import platform
import subprocess
from edgetpu.detection.engine import DetectionEngine
from PIL import Image
from PIL import ImageFont, ImageDraw
# Function to draw a rectangle with width > 1
def draw_rectangle(draw, coordinates, color, width=1):
for i in range(width):
rect_start = (coordinates[0] - i, coordinates[1] - i)
rect_end = (coordinates[2] + i, coordinates[3] + i)
draw.rectangle((rect_start, rect_end), outline = color)
# Function to read labels from text files.
def ReadLabelFile(file_path):
with open(file_path, 'r') as f:
lines = f.readlines()
ret = {}
for line in lines:
pair = line.strip().split(maxsplit=1)
ret[int(pair[0])] = pair[1].strip()
return ret
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--model', help='Path of the detection model.', required=True)
parser.add_argument(
'--label', help='Path of the labels file.')
parser.add_argument(
'--input', help='File path of the input image.', required=True)
parser.add_argument(
'--output', help='File path of the output image.')
args = parser.parse_args()
if not args.output:
output_name = 'object_detection_result.jpg'
else:
output_name = args.output
# Initialize engine.
engine = DetectionEngine(args.model)
labels = ReadLabelFile(args.label) if args.label else None
# Open image.
img = Image.open(args.input)
draw = ImageDraw.Draw(img)
helvetica=ImageFont.truetype("./Helvetica.ttf", size=72)
#defaultfont = ImageFont.load_default()
# Run inference.
ans = engine.DetectWithImage(img, threshold=0.05, keep_aspect_ratio=True,
relative_coord=False, top_k=10)
# Display result.
if ans:
print("\n")
for obj in ans:
if obj.score > 0.75:
if labels:
print(labels[obj.label_id], 'score = ', obj.score)
else:
print ('score = ', obj.score)
box = obj.bounding_box.flatten().tolist()
#print ('box = ', box)
# Draw a rectangle.
draw_rectangle(draw, box, 'red', width=5)
if labels:
draw.text((box[0] + 20, box[1] + 20), labels[obj.label_id], fill='red', font=helvetica)
img.save(output_name)
print ('Saved to ', output_name)
else:
print ('No object detected!')
if __name__ == '__main__':
main()
@frank3427
Copy link

frank3427 commented Sep 23, 2019

in the object detection what about using rtsp camera stream?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment