Showing
1 changed file
with
148 additions
and
0 deletions
Code/hotdog_recognition_example.py
0 → 100644
1 | +#***************************************************** | ||
2 | +# * | ||
3 | +# Copyright 2018 Amazon.com, Inc. or its affiliates. * | ||
4 | +# All Rights Reserved. * | ||
5 | +# * | ||
6 | +#***************************************************** | ||
7 | +""" A sample lambda for hotdog detection""" | ||
8 | +from threading import Thread, Event | ||
9 | +import os | ||
10 | +import json | ||
11 | +import numpy as np | ||
12 | +import awscam | ||
13 | +import cv2 | ||
14 | +import greengrasssdk | ||
15 | + | ||
16 | +class LocalDisplay(Thread): | ||
17 | + """ Class for facilitating the local display of inference results | ||
18 | + (as images). The class is designed to run on its own thread. In | ||
19 | + particular the class dumps the inference results into a FIFO | ||
20 | + located in the tmp directory (which lambda has access to). The | ||
21 | + results can be rendered using mplayer by typing: | ||
22 | + mplayer -demuxer lavf -lavfdopts format=mjpeg:probesize=32 /tmp/results.mjpeg | ||
23 | + """ | ||
24 | + def __init__(self, resolution): | ||
25 | + """ resolution - Desired resolution of the project stream """ | ||
26 | + # Initialize the base class, so that the object can run on its own | ||
27 | + # thread. | ||
28 | + super(LocalDisplay, self).__init__() | ||
29 | + # List of valid resolutions | ||
30 | + RESOLUTION = {'1080p' : (1920, 1080), '720p' : (1280, 720), '480p' : (858, 480)} | ||
31 | + if resolution not in RESOLUTION: | ||
32 | + raise Exception("Invalid resolution") | ||
33 | + self.resolution = RESOLUTION[resolution] | ||
34 | + # Initialize the default image to be a white canvas. Clients | ||
35 | + # will update the image when ready. | ||
36 | + self.frame = cv2.imencode('.jpg', 255*np.ones([640, 480, 3]))[1] | ||
37 | + self.stop_request = Event() | ||
38 | + | ||
39 | + def run(self): | ||
40 | + """ Overridden method that continually dumps images to the desired | ||
41 | + FIFO file. | ||
42 | + """ | ||
43 | + # Path to the FIFO file. The lambda only has permissions to the tmp | ||
44 | + # directory. Pointing to a FIFO file in another directory | ||
45 | + # will cause the lambda to crash. | ||
46 | + result_path = '/tmp/results.mjpeg' | ||
47 | + # Create the FIFO file if it doesn't exist. | ||
48 | + if not os.path.exists(result_path): | ||
49 | + os.mkfifo(result_path) | ||
50 | + # This call will block until a consumer is available | ||
51 | + with open(result_path, 'w') as fifo_file: | ||
52 | + while not self.stop_request.isSet(): | ||
53 | + try: | ||
54 | + # Write the data to the FIFO file. This call will block | ||
55 | + # meaning the code will come to a halt here until a consumer | ||
56 | + # is available. | ||
57 | + fifo_file.write(self.frame.tobytes()) | ||
58 | + except IOError: | ||
59 | + continue | ||
60 | + | ||
61 | + def set_frame_data(self, frame): | ||
62 | + """ Method updates the image data. This currently encodes the | ||
63 | + numpy array to jpg but can be modified to support other encodings. | ||
64 | + frame - Numpy array containing the image data of the next frame | ||
65 | + in the project stream. | ||
66 | + """ | ||
67 | + ret, jpeg = cv2.imencode('.jpg', cv2.resize(frame, self.resolution)) | ||
68 | + if not ret: | ||
69 | + raise Exception('Failed to set frame data') | ||
70 | + self.frame = jpeg | ||
71 | + | ||
72 | + def join(self): | ||
73 | + self.stop_request.set() | ||
74 | + | ||
75 | +def greengrass_infinite_infer_run(): | ||
76 | + """ Entry point of the lambda function""" | ||
77 | + try: | ||
78 | + # Use a squeezenet model and pick out the hotdog label. The model type | ||
79 | + # is classification. | ||
80 | + model_type = 'classification' | ||
81 | + # Create an IoT client for sending to messages to the cloud. | ||
82 | + client = greengrasssdk.client('iot-data') | ||
83 | + iot_topic = '$aws/things/{}/infer'.format(os.environ['AWS_IOT_THING_NAME']) | ||
84 | + # Create a local display instance that will dump the image bytes to a FIFO | ||
85 | + # file that the image can be rendered locally. | ||
86 | + local_display = LocalDisplay('480p') | ||
87 | + local_display.start() | ||
88 | + # The sample projects come with optimized artifacts, hence only the artifact | ||
89 | + # path is required. | ||
90 | + model_path = '/opt/awscam/artifacts/mxnet_squeezenet.xml' | ||
91 | + # Load the model onto the GPU. | ||
92 | + client.publish(topic=iot_topic, payload='Loading hotdog model') | ||
93 | + model = awscam.Model(model_path, {'GPU': 1}) | ||
94 | + client.publish(topic=iot_topic, payload='Hotdog model loaded') | ||
95 | + # Since this is a binary classifier only retrieve 2 classes. | ||
96 | + num_top_k = 2 | ||
97 | + # The height and width of the training set images | ||
98 | + input_height = 224 | ||
99 | + input_width = 224 | ||
100 | + # Do inference until the lambda is killed. | ||
101 | + while True: | ||
102 | + # Get a frame from the video stream | ||
103 | + ret, frame = awscam.getLastFrame() | ||
104 | + if not ret: | ||
105 | + raise Exception('Failed to get frame from the stream') | ||
106 | + # Resize frame to the same size as the training set. | ||
107 | + frame_resize = cv2.resize(frame, (input_height, input_width)) | ||
108 | + # Run the images through the inference engine and parse the results using | ||
109 | + # the parser API, note it is possible to get the output of doInference | ||
110 | + # and do the parsing manually, but since it is a classification model, | ||
111 | + # a simple API is provided. | ||
112 | + parsed_inference_results = model.parseResult(model_type, | ||
113 | + model.doInference(frame_resize)) | ||
114 | + # Get top k results with highest probabilities | ||
115 | + top_k = parsed_inference_results[model_type][0:num_top_k-1] | ||
116 | + # Get the probability of 'hotdog' label, which corresponds to label 934 in SqueezeNet. | ||
117 | + prob_hotdog = 0.0 | ||
118 | + for obj in top_k: | ||
119 | + if obj['label'] == 934: | ||
120 | + prob_hotdog = obj['prob'] | ||
121 | + break | ||
122 | + # Compute the probability of a hotdog not being in the image. | ||
123 | + prob_not_hotdog = 1.0 - prob_hotdog | ||
124 | + # Add two bars to indicate the probability of a hotdog being present and | ||
125 | + # the probability of a hotdog not being present. | ||
126 | + # See https://docs.opencv.org/3.4.1/d6/d6e/group__imgproc__draw.html | ||
127 | + # for more information about the cv2.rectangle method. | ||
128 | + # Method signature: image, point1, point2, color, and tickness. | ||
129 | + cv2.rectangle(frame, (0, 0), (int(frame.shape[1] * 0.2 * prob_not_hotdog), 80), | ||
130 | + (0, 0, 255), -1) | ||
131 | + cv2.rectangle(frame, (0, 90), (int(frame.shape[1] * 0.2 * prob_hotdog), 170), (0, 255, 0), -1) | ||
132 | + font = cv2.FONT_HERSHEY_SIMPLEX | ||
133 | + # See https://docs.opencv.org/3.4.1/d6/d6e/group__imgproc__draw.html | ||
134 | + # for more information about the cv2.putText method. | ||
135 | + # Method signature: image, text, origin, font face, font scale, color, | ||
136 | + # and tickness | ||
137 | + cv2.putText(frame, 'Not hotdog', (10, 70), font, 3, (225, 225, 225), 8) | ||
138 | + cv2.putText(frame, 'Hotdog', (10, 160), font, 3, (225, 225, 225), 8) | ||
139 | + # Send the top k results to the IoT console via MQTT | ||
140 | + client.publish(topic=iot_topic, payload=json.dumps({'Hotdog': prob_hotdog, | ||
141 | + 'Not hotdog': prob_not_hotdog})) | ||
142 | + # Set the next frame in the local display stream. | ||
143 | + local_display.set_frame_data(frame) | ||
144 | + except Exception as ex: | ||
145 | + client.publish(topic=iot_topic, payload='Error in hotdog lambda: {}'.format(ex)) | ||
146 | + | ||
147 | +# Execute the function above | ||
148 | +greengrass_infinite_infer_run() | ||
... | \ No newline at end of file | ... | \ No newline at end of file |
-
Please register or login to post a comment