def calc_metadata(args): #i need this function cause the metadata is NOT stored inside cfg
if (args.meta_data is None):
metadata=None
else: #convert json to Metadata() format
imported_Metadata_as_JSON = json.load(open(args.meta_data)) #parses argument as json
imported_Metadata_as_Metadata = Metadata(name=imported_Metadata_as_JSON["name"]) #create Metadata type and init it with "name" attribute
imported_Metadata_as_Metadata.set(thing_classes = imported_Metadata_as_JSON["thing_classes"]) #put classes information into Metadata
metadata = imported_Metadata_as_Metadata
return metadata
def postprocess(self, inference_output):
"""
Return predict result in batch
"""
print("Starting model inference output postprocessing...")
results = []
for model_inference in inference_output:
image = model_inference[0]
outputs = model_inference[1]
predictions = outputs["instances"].to("cpu")
severstal_metadata = Metadata()
severstal_metadata.set(
thing_classes=["Type 1", "Type 2", "Type 3", "Type 4"])
visualizer_pred = Visualizer(image[:, :, ::-1],
metadata=severstal_metadata,
scale=0.5)
image_pred = visualizer_pred.draw_instance_predictions(predictions)
image_cv2 = cv2.cvtColor(image_pred.get_image()[:, :, ::-1],
cv2.COLOR_BGR2RGB)
image_string = base64.b64encode(
cv2.imencode(".jpg", image_cv2)[1].tobytes()).decode("utf-8")
image_b64 = "data:image/jpg;base64," + image_string
if predictions.has("pred_classes"):
classes = predictions.pred_classes.numpy().tolist()
else:
classes = None
if predictions.has("scores"):
scores = predictions.scores.numpy().tolist()
else:
scores = None
if predictions.has("pred_boxes"):
boxes = predictions.pred_boxes.tensor.numpy().tolist()
else:
boxes = None
if predictions.has("pred_masks"):
#/!\ For an unknown reason (lack of memory, timeout...?), this doesn't work with TorchServe:
#/!\ (it works perfectly in a Jupyter notebook!)
#masks = predictions.pred_masks.numpy().tolist()
masks = None
else:
masks = None
result = {
"data": image_b64,
"classes": classes,
"scores": scores,
"boxes": boxes,
"masks": masks
}
results.append(json.dumps(result))
print("Model inference output postprocessing done")
return results
def make_inference(image, model_weights, threshold, n=5, save=False):
"""
Makes inference on image (single image) using model_config, model_weights and threshold.
Returns image with n instance predictions drawn on.
Params:
-------
image (str) : file path to target image
model_weights (str) : file path to model weights
threshold (float) : confidence threshold for model prediction, default 0.5
n (int) : number of prediction instances to draw on, default 5
Note: some images may not have 5 instances to draw on depending on threshold,
n=5 means the top 5 instances above the threshold will be drawn on.
save (bool) : if True will save image with predicted instances to file, default False
"""
# Create predictor and model config
cfg, predictor = create_predictor(model_weights, threshold)
# Convert PIL image to array
image = np.asarray(image)
# Create metadata
metadata = Metadata()
metadata.set(thing_classes=subset)
# Create visualizer instance
visualizer = Visualizer(img_rgb=image, metadata=metadata, scale=0.3)
outputs = predictor(image)
# Get instance predictions from outputs
instances = outputs["instances"]
# Draw on predictions to image
vis = visualizer.draw_instance_predictions(instances[:n].to("cpu"))
return vis.get_image(), instances[:n]
class detector:
def __init__(self, rgb_image, depth_image, fx, fy, cx, cy):
self.set_up_faster_rcnn()
self.set_up_fpointnet()
self.detection(rgb_image, depth_image, fx, fy, cx, cy)
def set_up_faster_rcnn(self):
self.cfg = get_cfg()
self.cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
self.cfg.MODEL.WEIGHTS = "weights/model_final.pth"
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.65
self.cfg.MODEL.DEVICE = 'cpu'
self.predictor = DefaultPredictor(self.cfg)
self.metadata = Metadata()
self.metadata.set(thing_classes=['station', 'forklift'])
def set_up_fpointnet(self):
self.FrustumPointNet = FrustumPointNetv1(n_classes=6, n_channel=6)
self.pth = torch.load("weights/frustum_model.pth", map_location='cpu')
self.FrustumPointNet.load_state_dict(self.pth['model_state_dict'])
self.model = self.FrustumPointNet.eval()
def detection(self, rgb_image, depth_image, fx, fy, cx, cy):
print('start detection')
rgb_image = rgb_image
depth_image = np.nan_to_num(depth_image, nan=0)
outputs = self.predictor(rgb_image)
prob_list = outputs["instances"].scores
class_list = outputs["instances"].pred_classes
box2d_list = outputs["instances"].pred_boxes.tensor
pitch = 0.09557043068606919
rotation = np.array([[1, 0, 0], [0, np.cos(pitch), -np.sin(pitch)],
[0, np.sin(pitch),
np.cos(pitch)]])
count = 0
pose = np.zeros([1, 4])
pose_list = []
for idx in range(len(class_list)):
object_class = class_list[idx].numpy()
prob = prob_list[idx].numpy()
xmin, ymin, xmax, ymax = map(int, box2d_list[idx])
if (xmax - xmin) > 1.5 * (ymax - ymin):
continue
rgb = np.zeros_like(rgb_image)
depth = np.zeros_like(depth_image)
rgb[ymin:ymax, xmin:xmax] = rgb_image[ymin:ymax, xmin:xmax]
depth[ymin:ymax, xmin:xmax] = depth_image[ymin:ymax, xmin:xmax]
print("class: {} ,depth_mean: {}".format(
object_class, np.mean(depth[ymin:ymax, xmin:xmax])))
pcs = depth2pc(rgb, depth, fx, fy, cx, cy,
1).point_cloud_generator()
pcs[:, 0:3] = np.dot(pcs[:, 0:3].astype(np.float32), rotation)
mask = pcs[:, 2] != 0
pcs = pcs[mask, :]
box2d_center = np.array([(xmin + xmax) / 2.0, (ymin + ymax) / 2.0])
uvdepth = np.zeros((1, 3))
uvdepth[0, 0:2] = box2d_center
uvdepth[0, 2] = np.mean(pcs[:, 2]) #20 # some random depth
x = ((uvdepth[:, 0] - cx) * uvdepth[:, 2]) / fx
y = ((uvdepth[:, 1] - cy) * uvdepth[:, 2]) / fy
uvdepth[:, 0] = x
uvdepth[:, 1] = y
frustum_angle = -1 * np.arctan2(uvdepth[0, 2], uvdepth[
0, 0]) # angle as to positive x-axis as in the Zoox paper
# Pass objects that are too small
if len(pcs) < 5:
continue
if object_class == 0:
object_class = 'box'
data = provider.FrustumDataset(npoints=2048,
pcs=pcs,
object_class=object_class,
frustum_angle=frustum_angle,
prob=prob)
point_set, rot_angle, prob, one_hot_vec = data.data()
point_set = torch.unsqueeze(torch.tensor(point_set),
0).transpose(2, 1).float()
one_hot_vec = torch.unsqueeze(torch.tensor(one_hot_vec),
0).float()
# print('start fpointnets')
logits, mask, stage1_center, center_boxnet, object_pts, \
heading_scores, heading_residuals_normalized, heading_residuals, \
size_scores, size_residuals_normalized, size_residuals, center = \
self.model(point_set, one_hot_vec)
corners_3d = get_box3d_corners(center, heading_residuals,
size_residuals)
logits = logits.detach().numpy()
mask = mask.detach().numpy()
center_boxnet = center_boxnet.detach().numpy()
object_pts = object_pts.detach().squeeze().numpy().transpose(
1, 0)
stage1_center = stage1_center.detach().numpy()
center = center.detach().numpy()
heading_scores = heading_scores.detach().numpy()
# heading_residuals_normalized = heading_residuals_normalized.detach().numpy()
heading_residuals = heading_residuals.detach().numpy()
size_scores = size_scores.detach().numpy()
size_residuals = size_residuals.detach().numpy()
corners_3d = corners_3d.detach().numpy()
output = np.argmax(logits, 2)
heading_class = np.argmax(heading_scores)
size_class = np.argmax(size_scores)
corners_3d = corners_3d[0, heading_class, size_class]
pred_angle = provider.class2angle(
heading_class, heading_residuals[0, heading_class],
NUM_HEADING_BIN)
pred_size = provider.class2size(size_class,
size_residuals[0, size_class])
cloud = pcs[:, 0:3].astype(np.float32)
object_cloud = (object_pts - center_boxnet.repeat(
object_pts.shape[0], 0)).astype(np.float32)
station_size = (0.979, 0.969, 0.979)
cube = generate_station_model_with_normal(
np.array([[0, 0, 0]]), station_size, -pred_angle)
station_cloud = cube.generate_points().astype(np.float32)
cloud_icp = cicp(object_cloud,
station_cloud,
max_iterations=20)
T, R, t = cloud_icp.cicp()
cloud_t = np.tile(t, (station_cloud.shape[0], 1))
station_cloud_rect = station_cloud[:, :3] - cloud_t
station_cloud_rect = station_cloud_rect + center.repeat(
station_cloud_rect.shape[0], 0)
object_cloud = object_cloud + center.repeat(
object_cloud.shape[0], 0)
station_cloud[:, :3] = station_cloud[:, :3] + center.repeat(
station_cloud.shape[0], 0)
center = center - t[np.newaxis, :]
corners_3d_rect = box3d_corners(center, pred_angle,
station_size)
object_cloud = rotate_pc_along_y(object_cloud, -rot_angle)
station_cloud_rect = rotate_pc_along_y(station_cloud_rect,
-rot_angle)
station_cloud[:, :3] = rotate_pc_along_y(
station_cloud[:, :3], -rot_angle)
center = rotate_pc_along_y(center, -rot_angle)
corners_3d = rotate_pc_along_y(corners_3d, -rot_angle)
corners_3d_rect = rotate_pc_along_y(corners_3d_rect,
-rot_angle)
center[0, 1] = 0.815
pose[0, :3] = center
pose[0, 3] = pred_angle + rot_angle
pose_list.append(pose.copy())
count += 1
station_rect_pub = point_cloud_publisher(
'/points_station_rect%d' % (count), station_cloud_rect)
bbox_pub_rect = bbox_publisher('/bbox_rect%d' % (count),
corners_3d_rect,
color="green")
object_pub = point_cloud_publisher(
'/points_object%d' % (count), object_cloud)
station_rect_pub.point_cloud_publish()
bbox_pub_rect.bbox_publish()
object_pub.point_cloud_publish()
pose_pub = pose_publisher('station_pose', pose_list)
pose_pub.pose_publish()
print('once detection')
class detector:
def __init__(self, rgb_image, depth_image, fx, fy, cx, cy):
self.set_up_faster_rcnn()
self.set_up_fpointnet()
self.detection(rgb_image, depth_image, fx, fy, cx, cy)
def set_up_faster_rcnn(self):
self.cfg = get_cfg()
self.cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
self.cfg.MODEL.WEIGHTS = "weights/model_final.pth"
self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.65
self.predictor = DefaultPredictor(self.cfg)
self.metadata = Metadata()
self.metadata.set(thing_classes=['station', 'forklift'])
def set_up_fpointnet(self):
self.FrustumPointNet = FrustumPointNetv1(n_classes=6,
n_channel=6).cuda()
self.pth = torch.load("weights/frustum_model.pth")
self.FrustumPointNet.load_state_dict(self.pth['model_state_dict'])
self.model = self.FrustumPointNet.eval()
def detection(self, rgb_image, depth_image, fx, fy, cx, cy):
print('start detection')
rgb_image = rgb_image
depth_image = np.nan_to_num(depth_image, nan=0)
outputs = self.predictor(rgb_image)
prob_list = outputs["instances"].scores
class_list = outputs["instances"].pred_classes
box2d_list = outputs["instances"].pred_boxes.tensor
# # v = Visualizer(rgb_image[:, :, ::-1], MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]), scale=1.2)
# v = Visualizer(rgb_image[:, :, ::-1], metadata=self.metadata, scale=1.2)
# out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# cv2.namedWindow('test',0)
# cv2.imshow('test',out.get_image()[:, :, ::-1])
# cv2.waitKey(0)
# print("depth mean {}".format(np.mean(depth_image)))
pitch = 0.09557043068606919
rotation = np.array([[1, 0, 0], [0, np.cos(pitch), -np.sin(pitch)],
[0, np.sin(pitch),
np.cos(pitch)]])
count = 0
pose = np.zeros([1, 4])
pose_list = []
for idx in range(len(class_list)):
object_class = class_list[idx].cpu().numpy()
prob = prob_list[idx].cpu().numpy()
xmin, ymin, xmax, ymax = map(int, box2d_list[idx])
if (xmax - xmin) > 1.5 * (ymax - ymin):
continue
rgb = np.zeros_like(rgb_image)
depth = np.zeros_like(depth_image)
rgb[ymin:ymax, xmin:xmax] = rgb_image[ymin:ymax, xmin:xmax]
depth[ymin:ymax, xmin:xmax] = depth_image[ymin:ymax, xmin:xmax]
print("class: {} ,depth_mean: {}".format(
object_class, np.mean(depth[ymin:ymax, xmin:xmax])))
pcs = depth2pc(rgb, depth, fx, fy, cx, cy,
1).point_cloud_generator()
pcs[:, 0:3] = np.dot(pcs[:, 0:3].astype(np.float32), rotation)
mask = pcs[:, 2] != 0
pcs = pcs[mask, :]
box2d_center = np.array([(xmin + xmax) / 2.0, (ymin + ymax) / 2.0])
uvdepth = np.zeros((1, 3))
uvdepth[0, 0:2] = box2d_center
uvdepth[0, 2] = np.mean(pcs[:, 2]) #20 # some random depth
x = ((uvdepth[:, 0] - cx) * uvdepth[:, 2]) / fx
y = ((uvdepth[:, 1] - cy) * uvdepth[:, 2]) / fy
uvdepth[:, 0] = x
uvdepth[:, 1] = y
frustum_angle = -1 * np.arctan2(uvdepth[0, 2], uvdepth[
0, 0]) # angle as to positive x-axis as in the Zoox paper
# Pass objects that are too small
if len(pcs) < 5:
continue
if object_class == 0:
object_class = 'box'
data = provider.FrustumDataset(npoints=2048,
pcs=pcs,
object_class=object_class,
frustum_angle=frustum_angle,
prob=prob)
point_set, rot_angle, prob, one_hot_vec = data.data()
point_set = torch.unsqueeze(torch.tensor(point_set),
0).transpose(2, 1).float().cuda()
one_hot_vec = torch.unsqueeze(torch.tensor(one_hot_vec),
0).float().cuda()
# print('start fpointnets')
logits, mask, stage1_center, center_boxnet, object_pts, \
heading_scores, heading_residuals_normalized, heading_residuals, \
size_scores, size_residuals_normalized, size_residuals, center = \
self.model(point_set, one_hot_vec)
corners_3d = get_box3d_corners(center, heading_residuals,
size_residuals)
logits = logits.cpu().detach().numpy()
mask = mask.cpu().detach().numpy()
center_boxnet = center_boxnet.cpu().detach().numpy()
object_pts = object_pts.cpu().detach().squeeze().numpy(
).transpose(1, 0)
stage1_center = stage1_center.cpu().detach().numpy()
center = center.cpu().detach().numpy()
heading_scores = heading_scores.cpu().detach().numpy()
# heading_residuals_normalized = heading_residuals_normalized.cpu().detach().numpy()
heading_residuals = heading_residuals.cpu().detach().numpy()
size_scores = size_scores.cpu().detach().numpy()
size_residuals = size_residuals.cpu().detach().numpy()
corners_3d = corners_3d.cpu().detach().numpy()
output = np.argmax(logits, 2)
heading_class = np.argmax(heading_scores)
size_class = np.argmax(size_scores)
corners_3d = corners_3d[0, heading_class, size_class]
pred_angle = provider.class2angle(
heading_class, heading_residuals[0, heading_class],
NUM_HEADING_BIN)
pred_size = provider.class2size(size_class,
size_residuals[0, size_class])
cloud = pcs[:, 0:3].astype(np.float32)
object_cloud = (object_pts - center_boxnet.repeat(
object_pts.shape[0], 0)).astype(np.float32)
station_size = (0.979, 0.969, 0.979)
cube = generate_station_model_with_normal(
np.array([[0, 0, 0]]), station_size, -pred_angle)
# object_cloud = rotate_pc_along_y(object_cloud,pred_angle)
# cube = generate_station_model_with_normal(np.array([[0,0,0]]),station_size,0)
station_cloud = cube.generate_points().astype(np.float32)
# # object_cloud_crop = object_cloud[object_cloud[:,1]>(center[0][1]-0.48)]
# # station_cloud_crop = station_cloud[station_cloud[:,1]>(center[0][1]-0.48)]
# object_cloud_crop = object_cloud[object_cloud[:,1]>(-0.48)]
# station_cloud_crop = station_cloud[station_cloud[:,1]>(-0.48)]
# # station_cloud = rotate_pc_along_y(station_cloud, rot_angle)
# cloud_icp = cicp(object_cloud,station_cloud_crop,max_iterations=30)
cloud_icp = cicp(object_cloud,
station_cloud,
max_iterations=20)
# cloud_icp = icp(object_cloud,station_cloud,max_iterations=20)
T, R, t = cloud_icp.cicp()
# R_angle = np.arccos(R[0,0])
# print(R)
# print(t)
# if abs(t[0]) > abs(t[2]):
# t[2] = 0
# if abs(t[0]) < abs(t[2]):
# t[0] = 0
# print(t)
cloud_t = np.tile(t, (station_cloud.shape[0], 1))
# station_cloud_rect = np.dot(station_cloud[:,:3]-cloud_t, np.transpose(np.linalg.pinv(R)))
# station_cloud_rect = np.dot(station_cloud[:,:3]-cloud_t, np.transpose(np.linalg.pinv(R)))
# station_cloud[:,:3] = rotate_pc_along_y(station_cloud[:,:3],-R_angle)
station_cloud_rect = station_cloud[:, :3] - cloud_t
station_cloud_rect = station_cloud_rect + center.repeat(
station_cloud_rect.shape[0], 0)
object_cloud = object_cloud + center.repeat(
object_cloud.shape[0], 0)
station_cloud[:, :3] = station_cloud[:, :3] + center.repeat(
station_cloud.shape[0], 0)
center = center - t[np.newaxis, :]
corners_3d_rect = box3d_corners(center, pred_angle,
station_size)
object_cloud = rotate_pc_along_y(object_cloud, -rot_angle)
station_cloud_rect = rotate_pc_along_y(station_cloud_rect,
-rot_angle)
station_cloud[:, :3] = rotate_pc_along_y(
station_cloud[:, :3], -rot_angle)
center = rotate_pc_along_y(center, -rot_angle)
corners_3d = rotate_pc_along_y(corners_3d, -rot_angle)
corners_3d_rect = rotate_pc_along_y(corners_3d_rect,
-rot_angle)
center[0, 1] = 0.815
pose[0, :3] = center
pose[0, 3] = pred_angle + rot_angle
pose_list.append(pose.copy())
# # station_cloud_rect = rotate_pc_along_y(station_cloud[:,:3]-cloud_t, R_angle)
# # cloud = np.dot(cloud, np.transpose(R))+np.tile(t,(cloud.shape[0],1))
# object_pub = point_cloud_publisher('/points_object',object_cloud_crop)
# station_pub = point_cloud_publisher('/points_station',station_cloud_crop[:,:3])
count += 1
# station_pub = point_cloud_publisher('/points_station%d'%(count),station_cloud[:,:3])
station_rect_pub = point_cloud_publisher(
'/points_station_rect%d' % (count), station_cloud_rect)
# bbox_pub = bbox_publisher('/bbox%d'%(count),corners_3d)
bbox_pub_rect = bbox_publisher('/bbox_rect%d' % (count),
corners_3d_rect,
color="green")
object_pub = point_cloud_publisher(
'/points_object%d' % (count), object_cloud)
# cloud_pub = point_cloud_publisher('/points_cloud%d'%(count),cloud)
# cloud_pub1 = point_cloud_publisher('/points_cloud1',cloud1)
station_rect_pub.point_cloud_publish()
bbox_pub_rect.bbox_publish()
object_pub.point_cloud_publish()
# bbox_list.append(bbox_pub)
# bbox_list.append(bbox_pub_rect)
# point_cloud_list.append(station_pub)
# point_cloud_list.append(station_rect_pub)
# point_cloud_list.append(object_pub)
# point_cloud_list.append(cloud_pub)
# with open('results.txt','ab') as f:
# np.savetxt(f, box_info, delimiter=" ")
# rate = rospy.Rate(10)
# while not rospy.is_shutdown():
# # point_cloud_list[3].point_cloud_publish()
# # # object_pub.point_cloud_publish()
# station_pub.point_cloud_publish()
# station_rect_pub.point_cloud_publish()
# bbox_pub.bbox_publish()
# bbox_pub_rect.bbox_publish()
# object_pub.point_cloud_publish()
# cloud_pub.point_cloud_publish()
# # # cloud_pub1.point_cloud_publish()
# rate.sleep()
pose_pub = pose_publisher('station_pose', pose_list)
pose_pub.pose_publish()
print('once detection')