Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save robertpyke/b1bb4f7d8ac86074a850284498cb0c38 to your computer and use it in GitHub Desktop.
Save robertpyke/b1bb4f7d8ac86074a850284498cb0c38 to your computer and use it in GitHub Desktop.
#*****************************************************
#
# 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