4
4
5
5
#https://colab.research.google.com/github/tensorflow/tpu/blob/master/models/official/mask_rcnn/mask_rcnn_demo.ipynb#scrollTo=2oZWLz4xXsyQ
6
6
7
- class_mappings = {1 : 'person' , 3 : 'car' , 28 : 'umbrella' , 31 : 'handbag' }
7
+ class_mappings = {1 : 'person' , 3 : 'car' , 28 : 'umbrella' , 31 : 'handbag' }
8
8
session = tf .compat .v1 .Session ()
9
9
10
+
10
11
def load_model ():
11
12
saved_model_dir = 'gs://cloud-tpu-checkpoints/mask-rcnn/1555659850'
12
- _ = tf .compat .v1 .saved_model .loader .load (session , ['serve' ], saved_model_dir )
13
+ _ = tf .compat .v1 .saved_model .loader .load (session , ['serve' ],
14
+ saved_model_dir )
13
15
14
16
15
17
def predict (np_image_string , min_score , height , width ):
16
18
num_detections , detection_boxes , detection_classes , detection_scores , detection_masks , image_info = session .run (
17
- ['NumDetections:0' , 'DetectionBoxes:0' , 'DetectionClasses:0' , 'DetectionScores:0' , 'DetectionMasks:0' , 'ImageInfo:0' ],
18
- feed_dict = {'Placeholder:0' : np_image_string })
19
+ [
20
+ 'NumDetections:0' , 'DetectionBoxes:0' , 'DetectionClasses:0' ,
21
+ 'DetectionScores:0' , 'DetectionMasks:0' , 'ImageInfo:0'
22
+ ],
23
+ feed_dict = {'Placeholder:0' : np_image_string })
19
24
num_detections = np .squeeze (num_detections .astype (np .int32 ), axis = (0 ,))
20
25
detection_scores = np .squeeze (detection_scores , axis = (0 ,))[0 :num_detections ]
21
- response = {
22
- 'boxes' : np .squeeze (detection_boxes * image_info [0 , 2 ], axis = (0 ,))[0 :num_detections ],
23
- 'class_indices' : np .squeeze (detection_classes .astype (np .int32 ), axis = (0 ,))[0 :num_detections ],
26
+ response = {
27
+ 'boxes' :
28
+ np .squeeze (detection_boxes * image_info [0 , 2 ], axis = (0 ,))
29
+ [0 :num_detections ],
30
+ 'class_indices' :
31
+ np .squeeze (detection_classes .astype (np .int32 ), axis = (0 ,))
32
+ [0 :num_detections ],
24
33
}
25
34
ymin , xmin , ymax , xmax = np .split (response ['boxes' ], 4 , axis = - 1 )
26
35
instance_masks = np .squeeze (detection_masks , axis = (0 ,))[0 :num_detections ]
27
- processed_boxes = np .concatenate ([xmin , ymin , xmax - xmin , ymax - ymin ], axis = - 1 )
28
- response .update ({'seg_masks' : generate_segmentation_from_masks (instance_masks , processed_boxes , height , width )})
36
+ processed_boxes = np .concatenate ([xmin , ymin , xmax - xmin , ymax - ymin ],
37
+ axis = - 1 )
38
+ response .update ({
39
+ 'seg_masks' :
40
+ generate_segmentation_from_masks (instance_masks , processed_boxes ,
41
+ height , width )
42
+ })
29
43
keep_indices = detection_scores > min_score
30
- keep_indices = keep_indices & np .isin (response ['class_indices' ], list (class_mappings .keys ()))
44
+ keep_indices = keep_indices & np .isin (response ['class_indices' ],
45
+ list (class_mappings .keys ()))
31
46
for key in response :
32
47
response [key ] = response [key ][keep_indices ]
33
48
return response
@@ -54,6 +69,7 @@ def expand_boxes(boxes, scale):
54
69
55
70
return boxes_exp
56
71
72
+
57
73
def generate_segmentation_from_masks (masks ,
58
74
detected_boxes ,
59
75
image_height ,
@@ -74,7 +90,6 @@ def generate_segmentation_from_masks(masks,
74
90
the instance masks *pasted* on the image canvas.
75
91
"""
76
92
77
-
78
93
_ , mask_height , mask_width = masks .shape
79
94
scale = max ((mask_width + 2.0 ) / mask_width ,
80
95
(mask_height + 2.0 ) / mask_height )
@@ -84,37 +99,33 @@ def generate_segmentation_from_masks(masks,
84
99
padded_mask = np .zeros ((mask_height + 2 , mask_width + 2 ), dtype = np .float32 )
85
100
segms = []
86
101
for mask_ind , mask in enumerate (masks ):
87
- im_mask = np .zeros ((image_height , image_width ), dtype = np .uint8 )
88
- if is_image_mask :
89
- # Process whole-image masks.
90
- im_mask [:, :] = mask [:, :]
91
- else :
92
- # Process mask inside bounding boxes.
93
- padded_mask [1 :- 1 , 1 :- 1 ] = mask [:, :]
94
-
95
- ref_box = ref_boxes [mask_ind , :]
96
- w = ref_box [2 ] - ref_box [0 ] + 1
97
- h = ref_box [3 ] - ref_box [1 ] + 1
98
- w = np .maximum (w , 1 )
99
- h = np .maximum (h , 1 )
100
-
101
- mask = cv2 .resize (padded_mask , (w , h ))
102
- mask = np .array (mask > 0.5 , dtype = np .uint8 )
103
-
104
- x_0 = max (ref_box [0 ], 0 )
105
- x_1 = min (ref_box [2 ] + 1 , image_width )
106
- y_0 = max (ref_box [1 ], 0 )
107
- y_1 = min (ref_box [3 ] + 1 , image_height )
108
-
109
- im_mask [y_0 :y_1 , x_0 :x_1 ] = mask [(y_0 - ref_box [1 ]):(y_1 - ref_box [1 ]), (
110
- x_0 - ref_box [0 ]):(x_1 - ref_box [0 ])]
111
- segms .append (im_mask )
102
+ im_mask = np .zeros ((image_height , image_width ), dtype = np .uint8 )
103
+ if is_image_mask :
104
+ # Process whole-image masks.
105
+ im_mask [:, :] = mask [:, :]
106
+ else :
107
+ # Process mask inside bounding boxes.
108
+ padded_mask [1 :- 1 , 1 :- 1 ] = mask [:, :]
109
+
110
+ ref_box = ref_boxes [mask_ind , :]
111
+ w = ref_box [2 ] - ref_box [0 ] + 1
112
+ h = ref_box [3 ] - ref_box [1 ] + 1
113
+ w = np .maximum (w , 1 )
114
+ h = np .maximum (h , 1 )
115
+
116
+ mask = cv2 .resize (padded_mask , (w , h ))
117
+ mask = np .array (mask > 0.5 , dtype = np .uint8 )
118
+
119
+ x_0 = max (ref_box [0 ], 0 )
120
+ x_1 = min (ref_box [2 ] + 1 , image_width )
121
+ y_0 = max (ref_box [1 ], 0 )
122
+ y_1 = min (ref_box [3 ] + 1 , image_height )
123
+
124
+ im_mask [y_0 :y_1 ,
125
+ x_0 :x_1 ] = mask [(y_0 - ref_box [1 ]):(y_1 - ref_box [1 ]),
126
+ (x_0 - ref_box [0 ]):(x_1 - ref_box [0 ])]
127
+ segms .append (im_mask )
112
128
113
129
segms = np .array (segms )
114
130
assert masks .shape [0 ] == segms .shape [0 ]
115
131
return segms
116
-
117
-
118
-
119
-
120
-
0 commit comments