diff --git a/deploy/python/infer.py b/deploy/python/infer.py index 33f29a7b74bbdbcfb8b3ee0feb5cb59db78ba8f6..99d3b3f44d1eafd8a6560aa5c85c6531d6cdb649 100644 --- a/deploy/python/infer.py +++ b/deploy/python/infer.py @@ -71,10 +71,11 @@ class Detector(object): def preprocess(self, im): preprocess_ops = [] for op_info in self.config.preprocess_infos: - op_type = op_info.pop('type') + new_op_info = op_info.copy() + op_type = new_op_info.pop('type') if op_type == 'Resize': - op_info['arch'] = self.config.arch - preprocess_ops.append(eval(op_type)(**op_info)) + new_op_info['arch'] = self.config.arch + preprocess_ops.append(eval(op_type)(**new_op_info)) im, im_info = preprocess(im, preprocess_ops) inputs = create_inputs(im, im_info, self.config.arch) return inputs, im_info @@ -481,7 +482,8 @@ def predict_video(detector, camera_id): frame, results, detector.config.labels, - mask_resolution=detector.config.mask_resolution) + mask_resolution=detector.config.mask_resolution, + threshold=FLAGS.threshold) im = np.array(im) writer.write(im) if camera_id != -1: diff --git a/deploy/python/visualize.py b/deploy/python/visualize.py index b58a36926657f3246c10f3009952c9481fa3bc72..1c136be4dde9c97c01649b1dab62135eed8062d8 100644 --- a/deploy/python/visualize.py +++ b/deploy/python/visualize.py @@ -228,9 +228,22 @@ def draw_segm(im, color_mask = np.array(color_mask) im[idx[0], idx[1], :] *= 1.0 - alpha im[idx[0], idx[1], :] += alpha * color_mask - center_y, center_x = ndimage.measurements.center_of_mass(mask) - label_text = "{}".format(labels[clsid]) - vis_pos = (max(int(center_x) - 10, 0), int(center_y)) - cv2.putText(im, label_text, vis_pos, cv2.FONT_HERSHEY_COMPLEX, 0.3, - (255, 255, 255)) + sum_x = np.sum(mask, axis=0) + x = np.where(sum_x > 0.5)[0] + sum_y = np.sum(mask, axis=1) + y = np.where(sum_y > 0.5)[0] + x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1] + cv2.rectangle(im, (x0, y0), (x1, y1), + tuple(color_mask.astype('int32').tolist()), 1) + bbox_text = '%s %.2f' % (labels[clsid], score) + t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0] + cv2.rectangle(im, (x0, y0), (x0 + t_size[0], y0 - t_size[1] - 3), + tuple(color_mask.astype('int32').tolist()), -1) + cv2.putText( + im, + bbox_text, (x0, y0 - 2), + cv2.FONT_HERSHEY_SIMPLEX, + 0.3, (0, 0, 0), + 1, + lineType=cv2.LINE_AA) return Image.fromarray(im.astype('uint8'))