forked from multimodallearning/pytorch-mask-rcnn
-
Notifications
You must be signed in to change notification settings - Fork 2
/
flask_helpers.py
231 lines (189 loc) · 7.54 KB
/
flask_helpers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
import os
# For Mask RCNN
import torch
import coco
import utils
import model as modellib
import visualize
from config import Config
# For processing images
import numpy as np
from PIL import Image
from io import BytesIO
import base64
import wget
# Global vars
ENCODING = 'utf-8'
def image_from_request(request):
"""Reads the attached image file from incoming request
Input:
request: Request object, incoming flask request
Output:
image: PIL image, the image attached to the request
"""
for key in request.files.keys():
attached_file = request.files[key]
try:
image = Image.open(attached_file.stream)
return image
except:
print('Error loading image: {}'.format(attached_file.filename))
return None
def extract_bounding_boxes(image, results, apply_mask=True):
"""Slices regions defined by bounding box into separate images
Inputs:
image: np.array, original image
results: results object, results of maskrcnn
apply_mask: boolean, toggles if the mask should be applied
Output:
output_images: list of np.arrays, image segments
"""
bounding_boxes = results['rois']
output_images = []
try:
num_images = bounding_boxes.shape[0]
for ix in range(num_images):
cur_image = image.copy()
if apply_mask:
cur_mask = np.squeeze(results['masks'][:,:,ix])
cur_image[cur_mask==0,:] = 255
#bbox array [num_instances, (y1, x1, y2, x2)].
bb = bounding_boxes[ix,:]
cur_out = cur_image[bb[0]:bb[2], bb[1]:bb[3]]
output_images.append(cur_out)
except Exception as e:
print(e)
return output_images
def save_images_locally(output_images):
"""Saves the results of extract_bounding_boxes locally
Inputs:
output_images: list of np.arrays, image segments
"""
for ix in range(len(output_images)):
cur_img = Image.fromarray(output_images[ix])
save_name = "result_{}.jpg".format(ix)
cur_img.save(save_name)
def outputs_to_base64(output_images, img_format):
"""Converts the output images to a list of base64 strings
Input:
output_images: list of np.arrays, image segments
img_format: string, format of original image
Output:
output_strings: list of string, base64 encoded images
"""
output_strings = []
# Encode back to base64 to send back to ImageProcessing function
for ix in range(len(output_images)):
cur_image = Image.fromarray(output_images[ix])
cropped_img_bytes = BytesIO()
cur_image.save(cropped_img_bytes, format=img_format)
cropped_img_bytes = cropped_img_bytes.getvalue()
base64_bytes = base64.b64encode(cropped_img_bytes)
base64_string = base64_bytes.decode(ENCODING)
output_strings.append(base64_string)
return output_strings
def image_to_array(image):
"""Converts a pil image to a numpy array of RGB uint8's
Input:
image: PIL or skimage image, image to be converted
Output:
im_array: numpy array, image ready to be handed to Mask R-CNN
"""
# Restricts to RGB
im_array = np.array(image)[:,:,:3]
if np.max(im_array) <= 1.0:
im_array = np.floor(im_array * 255).astype(np.uint8)
return im_array
################# Flask functions ########################
def get_default_model():
"""Loads the coco classifier as the default model
Input:
none
Outputs:
model: mask rcnn model, the coco dataset trained model for predictions
class_names: list of strings, the class names for the model prediction classes
"""
# Root directory of the project
ROOT_DIR = os.getcwd()
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.pth")
class InferenceConfig(coco.CocoConfig):
# Set batch size to 1 since we'll be running inference on
# one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
# GPU_COUNT = 0 for CPU
GPU_COUNT = 1
IMAGES_PER_GPU = 1
config = InferenceConfig()
if not torch.cuda.is_available():
config.GPU_COUNT = 0
config.display()
# Create model object.
model = modellib.MaskRCNN(model_dir=MODEL_DIR, config=config)
if config.GPU_COUNT:
model = model.cuda()
# Load weights trained on MS-COCO
model.load_state_dict(torch.load(COCO_MODEL_PATH))
# COCO Class names
# Index of the class in the list is its ID. For example, to get ID of
# the teddy bear class, use: class_names.index('teddy bear')
class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
'teddy bear', 'hair drier', 'toothbrush']
return model, class_names
def set_model(model_name, model_url, class_names):
"""Changes that model being used for the endpoint
Inputs:
model_name: string, name of the model to use
model_url: string, url that the model can be downloaded from
class_names: list of strings, the names for the classes
Outputs:
model: Mask RCNN model, model that predictions are made with
class_names: list of strings, the names for the classes
"""
# Directory name setting
root_dir = os.getcwd()
model_dir = os.path.join(root_dir, "logs")
if not os.path.exists(model_name):
model_path = wget.download(model_url)
else:
model_path = model_name
config = CustomInferenceConfig()
config.COCO_MODEL_PATH = model_path
config.NUM_CLASSES = len(class_names)
config.NAME = model_name
if not torch.cuda.is_available():
config.GPU_COUNT = 0
config.display()
model = modellib.MaskRCNN(config=config, model_dir=model_dir)
if config.GPU_COUNT:
print('Has Cuda')
model = model.cuda()
# Load weights trained on MS-COCO
model.load_weights(model_path)
return model, class_names
class CustomInferenceConfig(coco.CocoConfig):
"""Derives from the base Config class and overrides some values."""
# Train on 1 GPU and 8 images per GPU. We can put multiple images on each
# GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
GPU_COUNT = 1
IMAGES_PER_GPU = 1
# Number of training steps per epoch
STEPS_PER_EPOCH = 4
# Skip detections with < 90% confidence
DETECTION_MIN_CONFIDENCE = 0.90
# Necessary for docker image to optimize memory usage best
NUM_WORKERS = 0