Coral Edge TPU semantic_segmentation

改造版のプログラムです。
examples/semantic_segmentation_2.py
# Lint as: python3 # Copyright 2019 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. r"""An example of semantic segmentation. The following command runs this script and saves a new image showing the segmented pixels at the location specified by `output`: ``` bash examples/install_requirements.sh semantic_segmentation.py python3 examples/semantic_segmentation_2.py \ --model test_data/deeplabv3_mnv2_pascal_quant_edgetpu.tflite \ --input test_data/bird.bmp \ --keep_aspect_ratio \ --output ${HOME}/segmentation_result.jpg >cd examples >python semantic_segmentation_2.py --keep_aspect_ratio ``` """ import argparse import numpy as np from PIL import Image from pycoral.adapters import common from pycoral.adapters import segment from pycoral.utils.edgetpu import make_interpreter import time import cv2 from timeit import default_timer as timer def create_pascal_label_colormap(): """Creates a label colormap used in PASCAL VOC segmentation benchmark. Returns: A Colormap for visualizing segmentation results. """ colormap = np.zeros((256, 3), dtype=int) indices = np.arange(256, dtype=int) for shift in reversed(range(8)): for channel in range(3): colormap[:, channel] |= ((indices >> channel) & 1) << shift indices >>= 3 return colormap def label_to_color_image(label): """Adds color defined by the dataset colormap to the label. Args: label: A 2D array with integer type, storing the segmentation label. Returns: result: A 2D array with floating type. The element of the array is the color indexed by the corresponding element in the input label to the PASCAL color map. Raises: ValueError: If label is not of rank 2 or its value is larger than color map maximum entry. """ if label.ndim != 2: raise ValueError('Expect 2-D input label') colormap = create_pascal_label_colormap() if np.max(label) >= len(colormap): raise ValueError('label value too large.') return colormap[label] def one_shot(img): #img = Image.open(args.input) if args.keep_aspect_ratio: resized_img, _ = common.set_resized_input( interpreter, img.size, lambda size: img.resize(size, Image.ANTIALIAS)) else: resized_img = img.resize((width, height), Image.ANTIALIAS) common.set_input(interpreter, resized_img) interpreter.invoke() result = segment.get_output(interpreter) if len(result.shape) == 3: result = np.argmax(result, axis=-1) if False: print('resized_img.size=',resized_img.size) # If keep_aspect_ratio, we need to remove the padding area. new_width, new_height = resized_img.size result = result[:new_height, :new_width] mask_img = Image.fromarray(label_to_color_image(result).astype(np.uint8)) # Concat resized input image and processed segmentation results. output_img = Image.new('RGB', (2 * new_width, new_height)) output_img.paste(resized_img, (0, 0)) #output_img.paste(mask_img, (width, 0)) output_img.paste(mask_img, (new_width, 0)) return output_img def camera_stream(): # Initialize video stream video = cv2.VideoCapture(args.video) time.sleep(1) accum_time = 0 curr_fps = 0 fps = "FPS: ??" prev_time = timer() first_f=True while(video.isOpened()): try: _, frame = video.read() if first_f: print('frame.shape=',frame.shape) first_f=False img = Image.fromarray(frame) out_img = one_shot(img) result = np.asarray(out_img) curr_time = timer() exec_time = curr_time - prev_time prev_time = curr_time accum_time = accum_time + exec_time curr_fps = curr_fps + 1 if accum_time > 1: accum_time = accum_time - 1 fps = "FPS: " + str(curr_fps) curr_fps = 0 cv2.putText(result, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.50, color=(255, 0, 0), thickness=2) cv2.imshow("Image", result) if( cv2.waitKey( 5 ) & 0xFF == ord( 'q' ) ): fps.stop() break except KeyboardInterrupt: break cv2.destroyAllWindows() video.release() time.sleep(2) if __name__ == '__main__': default_model = '../test_data/deeplabv3_mnv2_pascal_quant_edgetpu.tflite' parser = argparse.ArgumentParser() #parser.add_argument('--model', required=True, # help='Path of the segmentation model.') parser.add_argument('--model', help='Path of the segmentation model.', default=default_model) #parser.add_argument('--input', required=True, # help='File path of the input image.') parser.add_argument('--input', help='File path of the input image.', default='../test_data/bird.bmp') parser.add_argument('--output', default='semantic_segmentation_result.jpg', help='File path of the output image.') parser.add_argument( '--keep_aspect_ratio',action='store_true',default=False, help=( 'keep the image aspect ratio when down-sampling the image by adding ' 'black pixel padding (zeros) on bottom or right. ' 'By default the image is resized and reshaped without cropping. This ' 'option should be the same as what is applied on input images during ' 'model training. Otherwise the accuracy may be affected and the ' 'bounding box of detection result may be stretched.')) parser.add_argument('-v', '--video', type=str, default=0, required=False, help='Path to video') args = parser.parse_args() interpreter = make_interpreter(args.model, device=':0') interpreter.allocate_tensors() NET_w, NET_h = common.input_size(interpreter) width, height = NET_w, NET_h print("NET_w=%d , NET_h=%d" %(NET_w, NET_h)) ONE_SHOT=False if ONE_SHOT: img = Image.open(args.input) output_img=one_shot(img) output_img.show() #output_img.save(args.output) #print('Done. Results saved at', args.output) else: camera_stream()

実行は、
> cd examples
> python semantic_segmentation_2.py --keep_aspect_ratio

FPS:6 - 7 位です。
注) Coral は、省電力モードです。