Created
June 20, 2018 23:53
-
-
Save robertpyke/b1bb4f7d8ac86074a850284498cb0c38 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#***************************************************** | |
# | |
# Adapted from Amazon's deep lens example project | |
# Cal256 object labelling (uses sagemaker pre-trained cal256 example model) | |
#***************************************************** | |
from threading import Thread, Event | |
import os | |
import json | |
import numpy as np | |
import awscam | |
import cv2 | |
import greengrasssdk | |
import mo | |
input_width = 224 #This should match the input width of the training set | |
input_height = 224 #This should match the input height of the training set | |
ret,model_path = mo.optimize('cal256-classification-model', input_width, input_height) | |
labels = [ | |
'ak47', | |
'american-flag', | |
'backpack', | |
'baseball-bat', | |
'baseball-glove', | |
'basketball-hoop', | |
'bat', | |
'bathtub', | |
'bear', | |
'beer-mug', | |
'billiards', | |
'binoculars', | |
'birdbath', | |
'blimp', | |
'bonsai-101', | |
'boom-box', | |
'bowling-ball', | |
'bowling-pin', | |
'boxing-glove', | |
'brain-101', | |
'breadmaker', | |
'buddha-101', | |
'bulldozer', | |
'butterfly', | |
'cactus', | |
'cake', | |
'calculator', | |
'camel', | |
'cannon', | |
'canoe', | |
'car-tire', | |
'cartman', | |
'cd', | |
'centipede', | |
'cereal-box', | |
'chandelier-101', | |
'chess-board', | |
'chimp', | |
'chopsticks', | |
'cockroach', | |
'coffee-mug', | |
'coffin', | |
'coin', | |
'comet', | |
'computer-keyboard', | |
'computer-monitor', | |
'computer-mouse', | |
'conch', | |
'cormorant', | |
'covered-wagon', | |
'cowboy-hat', | |
'crab-101', | |
'desk-globe', | |
'diamond-ring', | |
'dice', | |
'dog', | |
'dolphin-101', | |
'doorknob', | |
'drinking-straw', | |
'duck', | |
'dumb-bell', | |
'eiffel-tower', | |
'electric-guitar-101', | |
'elephant-101', | |
'elk', | |
'ewer-101', | |
'eyeglasses', | |
'fern', | |
'fighter-jet', | |
'fire-extinguisher', | |
'fire-hydrant', | |
'fire-truck', | |
'fireworks', | |
'flashlight', | |
'floppy-disk', | |
'football-helmet', | |
'french-horn', | |
'fried-egg', | |
'frisbee', | |
'frog', | |
'frying-pan', | |
'galaxy', | |
'gas-pump', | |
'giraffe', | |
'goat', | |
'golden-gate-bridge', | |
'goldfish', | |
'golf-ball', | |
'goose', | |
'gorilla', | |
'grand-piano-101', | |
'grapes', | |
'grasshopper', | |
'guitar-pick', | |
'hamburger', | |
'hammock', | |
'harmonica', | |
'harp', | |
'harpsichord', | |
'hawksbill-101', | |
'head-phones', | |
'helicopter-101', | |
'hibiscus', | |
'homer-simpson', | |
'horse', | |
'horseshoe-crab', | |
'hot-air-balloon', | |
'hot-dog', | |
'hot-tub', | |
'hourglass', | |
'house-fly', | |
'human-skeleton', | |
'hummingbird', | |
'ibis-101', | |
'ice-cream-cone', | |
'iguana', | |
'ipod', | |
'iris', | |
'jesus-christ', | |
'joy-stick', | |
'kangaroo-101', | |
'kayak', | |
'ketch-101', | |
'killer-whale', | |
'knife', | |
'ladder', | |
'laptop-101', | |
'lathe', | |
'leopards-101', | |
'license-plate', | |
'lightbulb', | |
'light-house', | |
'lightning', | |
'llama-101', | |
'mailbox', | |
'mandolin', | |
'mars', | |
'mattress', | |
'megaphone', | |
'menorah-101', | |
'microscope', | |
'microwave', | |
'minaret', | |
'minotaur', | |
'motorbikes-101', | |
'mountain-bike', | |
'mushroom', | |
'mussels', | |
'necktie', | |
'octopus', | |
'ostrich', | |
'owl', | |
'palm-pilot', | |
'palm-tree', | |
'paperclip', | |
'paper-shredder', | |
'pci-card', | |
'penguin', | |
'people', | |
'pez-dispenser', | |
'photocopier', | |
'picnic-table', | |
'playing-card', | |
'porcupine', | |
'pram', | |
'praying-mantis', | |
'pyramid', | |
'raccoon', | |
'radio-telescope', | |
'rainbow', | |
'refrigerator', | |
'revolver-101', | |
'rifle', | |
'rotary-phone', | |
'roulette-wheel', | |
'saddle', | |
'saturn', | |
'school-bus', | |
'scorpion-101', | |
'screwdriver', | |
'segway', | |
'self-propelled-lawn-mower', | |
'sextant', | |
'sheet-music', | |
'skateboard', | |
'skunk', | |
'skyscraper', | |
'smokestack', | |
'snail', | |
'snake', | |
'sneaker', | |
'snowmobile', | |
'soccer-ball', | |
'socks', | |
'soda-can', | |
'spaghetti', | |
'speed-boat', | |
'spider', | |
'spoon', | |
'stained-glass', | |
'starfish-101', | |
'steering-wheel', | |
'stirrups', | |
'sunflower-101', | |
'superman', | |
'sushi', | |
'swan', | |
'swiss-army-knife', | |
'sword', | |
'syringe', | |
'tambourine', | |
'teapot', | |
'teddy-bear', | |
'teepee', | |
'telephone-box', | |
'tennis-ball', | |
'tennis-court', | |
'tennis-racket', | |
'theodolite', | |
'toaster', | |
'tomato', | |
'tombstone', | |
'top-hat', | |
'touring-bike', | |
'tower-pisa', | |
'traffic-light', | |
'treadmill', | |
'triceratops', | |
'tricycle', | |
'trilobite-101', | |
'tripod', | |
't-shirt', | |
'tuning-fork', | |
'tweezer', | |
'umbrella-101', | |
'unicorn', | |
'vcr', | |
'video-projector', | |
'washing-machine', | |
'watch-101', | |
'waterfall', | |
'watermelon', | |
'welding-mask', | |
'wheelbarrow', | |
'windmill', | |
'wine-bottle', | |
'xylophone', | |
'yarmulke', | |
'yo-yo', | |
'zebra', | |
'airplanes-101', | |
'car-side-101', | |
'faces-easy-101', | |
'greyhound', | |
'tennis-shoes', | |
'toad', | |
'clutter' | |
] | |
class LocalDisplay(Thread): | |
""" Class for facilitating the local display of inference results | |
(as images). The class is designed to run on its own thread. In | |
particular the class dumps the inference results into a FIFO | |
located in the tmp directory (which lambda has access to). The | |
results can be rendered using mplayer by typing: | |
mplayer -demuxer lavf -lavfdopts format=mjpeg:probesize=32 /tmp/results.mjpeg | |
""" | |
def __init__(self, resolution): | |
""" resolution - Desired resolution of the project stream """ | |
# Initialize the base class, so that the object can run on its own | |
# thread. | |
super(LocalDisplay, self).__init__() | |
# List of valid resolutions | |
RESOLUTION = {'1080p' : (1920, 1080), '720p' : (1280, 720), '480p' : (858, 480)} | |
if resolution not in RESOLUTION: | |
raise Exception("Invalid resolution") | |
self.resolution = RESOLUTION[resolution] | |
# Initialize the default image to be a white canvas. Clients | |
# will update the image when ready. | |
self.frame = cv2.imencode('.jpg', 255*np.ones([640, 480, 3]))[1] | |
self.stop_request = Event() | |
def run(self): | |
""" Overridden method that continually dumps images to the desired | |
FIFO file. | |
""" | |
# Path to the FIFO file. The lambda only has permissions to the tmp | |
# directory. Pointing to a FIFO file in another directory | |
# will cause the lambda to crash. | |
result_path = '/tmp/results.mjpeg' | |
# Create the FIFO file if it doesn't exist. | |
if not os.path.exists(result_path): | |
os.mkfifo(result_path) | |
# This call will block until a consumer is available | |
with open(result_path, 'w') as fifo_file: | |
while not self.stop_request.isSet(): | |
try: | |
# Write the data to the FIFO file. This call will block | |
# meaning the code will come to a halt here until a consumer | |
# is available. | |
fifo_file.write(self.frame.tobytes()) | |
except IOError: | |
continue | |
def set_frame_data(self, frame): | |
""" Method updates the image data. This currently encodes the | |
numpy array to jpg but can be modified to support other encodings. | |
frame - Numpy array containing the image data of the next frame | |
in the project stream. | |
""" | |
ret, jpeg = cv2.imencode('.jpg', cv2.resize(frame, self.resolution)) | |
if not ret: | |
raise Exception('Failed to set frame data') | |
self.frame = jpeg | |
def join(self): | |
self.stop_request.set() | |
def greengrass_infinite_infer_run(): | |
""" Entry point of the lambda function""" | |
try: | |
# This action recognition model is implemented as multi classifier. Since the number | |
# of labels is small, create a dictionary that converts the machine labels to human | |
# readable labels. | |
model_type = 'classification' | |
# Create an IoT client for sending to messages to the cloud. | |
client = greengrasssdk.client('iot-data') | |
iot_topic = '$aws/things/{}/infer'.format(os.environ['AWS_IOT_THING_NAME']) | |
# Create a local display instance that will dump the image bytes to a FIFO | |
# file that the image can be rendered locally. | |
local_display = LocalDisplay('480p') | |
local_display.start() | |
# The sample projects come with optimized artifacts, hence only the artifact | |
# path is required. | |
# model_path now generated above... | |
# Load the model onto the GPU. | |
client.publish(topic=iot_topic, payload='Loading cal256 recognition model') | |
model = awscam.Model(model_path, {'GPU': 1}) | |
client.publish(topic=iot_topic, payload='cal256 recognition model loaded') | |
# The number of top results to stream to IoT. | |
num_top_k = 5 | |
# The height and width of the training set images | |
input_height = 224 | |
input_width = 224 | |
# Do inference until the lambda is killed. | |
while True: | |
# Get a frame from the video stream | |
ret, frame = awscam.getLastFrame() | |
if not ret: | |
raise Exception('Failed to get frame from the stream') | |
# Resize frame to the same size as the training set. | |
frame_resize = cv2.resize(frame, (input_height, input_width)) | |
# Model was trained in RGB format but getLastFrame returns image | |
# in BGR format so need to switch. | |
frame_resize = cv2.cvtColor(frame_resize, cv2.COLOR_BGR2RGB) | |
# Run the images through the inference engine and parse the results using | |
# the parser API, note it is possible to get the output of doInference | |
# and do the parsing manually, but since it is a classification model, | |
# a simple API is provided. | |
parsed_inference_results = model.parseResult(model_type, | |
model.doInference(frame_resize)) | |
# Get top k results with highest probabilities | |
top_k = parsed_inference_results[model_type][0:num_top_k] | |
# Add the label of the top result to the frame used by local display. | |
# See https://docs.opencv.org/3.4.1/d6/d6e/group__imgproc__draw.html | |
# for more information about the cv2.putText method. | |
# Method signature: image, text, origin, font face, font scale, color, and tickness | |
# labels here: http://www.vision.caltech.edu/Image_Datasets/Caltech256/images/ (we are base 0, so add 1) | |
client.publish(topic=iot_topic, payload='{}'.format(top_k)) | |
cv2.putText(frame, labels[top_k[0]['label']], (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 165, 20), 8) | |
# Set the next frame in the local display stream. | |
local_display.set_frame_data(frame) | |
# Send the top k results to the IoT console via MQTT | |
cloud_output = {} | |
for obj in top_k: | |
cloud_output['label'] = obj['prob'] | |
#client.publish(topic=iot_topic, payload=json.dumps(cloud_output)) | |
except Exception as ex: | |
client.publish(topic=iot_topic, payload='Error in cal256 recognition lambda: {}'.format(ex)) | |
greengrass_infinite_infer_run() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment