def draw_predictions(img, detections, num_classes=3):
for j in range(num_classes):
if len(detections[j]) > 0:
for det in detections[j]:
# (scores-0:1, x-1:2, y-2:3, z-3:4, dim-4:7, yaw-7:8)
_score, _x, _y, _z, _h, _w, _l, _yaw = det
drawRotatedBox(img, _x, _y, _w, _l, _yaw, cnf.colors[int(j)])
return img
def callback(self,data):
rospy.loginfo("detection")
with torch.no_grad():
gen = point_cloud2.read_points(data)
for idx, p in enumerate(gen):
print(idx)
print(p)
b = kitti_bev_utils.removePoints(gen, cnf.boundary)
imgs_bev = kitti_bev_utils.makeBVFeature(b, cnf.DISCRETIZATION, cnf.boundary)
input_imgs = imgs_bev.to(device=configs.device).float()
t1 = time_synchronized()
outputs = self.model(input_imgs)
t2 = time_synchronized()
detections = post_processing_v2(outputs, conf_thresh=configs.conf_thresh, nms_thresh=configs.nms_thresh)
img_detections = [] # Stores detections for each image index
img_detections.extend(detections)
img_bev = imgs_bev.squeeze() * 255
img_bev = img_bev.permute(1, 2, 0).numpy().astype(np.uint8)
img_bev = cv2.resize(img_bev, (configs.img_size, configs.img_size))
for detections in img_detections:
if detections is None:
continue
# Rescale boxes to original image
detections = rescale_boxes(detections, configs.img_size, img_bev.shape[:2])
for x, y, w, l, im, re, *_, cls_pred in detections:
yaw = np.arctan2(im, re)
# Draw rotated box
kitti_bev_utils.drawRotatedBox(img_bev, x, y, w, l, yaw, cnf.colors[int(cls_pred)])
img_bev = cv2.flip(cv2.flip(img_bev, 0), 1)
out_img = img_bev
cv2.imshow('test-img', out_img)
cv2.waitKey(1)
RGB_Map=None)
img_rgb = show_image_with_boxes(img_rgb, objects_pred, calib,
False)
# Rescale target
targets[:, 2:6] *= configs.img_size
# Get yaw angle
targets[:, 6] = torch.atan2(targets[:, 6], targets[:, 7])
img_bev = imgs.squeeze() * 255
img_bev = img_bev.permute(1, 2, 0).numpy().astype(np.uint8)
img_bev = cv2.resize(img_bev, (configs.img_size, configs.img_size))
for c, x, y, w, l, yaw in targets[:, 1:7].numpy():
# Draw rotated box
bev_utils.drawRotatedBox(img_bev, x, y, w, l, yaw,
cnf.colors[int(c)])
img_bev = cv2.rotate(img_bev, cv2.ROTATE_180)
if configs.mosaic and configs.show_train_data:
cv2.imshow('mosaic_sample', img_bev)
else:
out_img = merge_rgb_to_bev(img_rgb,
img_bev,
output_width=configs.output_width)
cv2.imshow('single_sample', out_img)
if cv2.waitKey(0) & 0xff == 27:
break
for idx in range(len(dataset)):
bev_map, labels, img_rgb, img_path = dataset.draw_img_with_label(idx)
calib = Calibration(
img_path.replace(".png", ".txt").replace("image_2", "calib"))
bev_map = (bev_map.transpose(1, 2, 0) * 255).astype(np.uint8)
bev_map = cv2.resize(bev_map, (cnf.BEV_HEIGHT, cnf.BEV_WIDTH))
for box_idx, (cls_id, x, y, z, h, w, l, yaw) in enumerate(labels):
# Draw rotated box
yaw = -yaw
y1 = int((x - cnf.boundary['minX']) / cnf.DISCRETIZATION)
x1 = int((y - cnf.boundary['minY']) / cnf.DISCRETIZATION)
w1 = int(w / cnf.DISCRETIZATION)
l1 = int(l / cnf.DISCRETIZATION)
drawRotatedBox(bev_map, x1, y1, w1, l1, yaw,
cnf.colors[int(cls_id)])
# Rotate the bev_map
bev_map = cv2.rotate(bev_map, cv2.ROTATE_180)
labels[:, 1:] = lidar_to_camera_box(labels[:, 1:], calib.V2C, calib.R0,
calib.P2)
img_rgb = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
img_rgb = show_rgb_image_with_boxes(img_rgb, labels, calib)
out_img = merge_rgb_to_bev(img_rgb,
bev_map,
output_width=configs.output_width)
cv2.imshow('bev_map', out_img)
if cv2.waitKey(0) & 0xff == 27:
break
img_detections = [] # Stores detections for each image index
img_detections.extend(detections)
img_bev = imgs_bev.squeeze() * 255
img_bev = img_bev.permute(1, 2, 0).numpy().astype(np.uint8)
img_bev = cv2.resize(img_bev, (configs.img_size, configs.img_size))
for detections in img_detections:
if detections is None:
continue
# Rescale boxes to original image
detections = rescale_boxes(detections, configs.img_size,
img_bev.shape[:2])
for x, y, w, l, im, re, cls_conf, cls_pred in detections:
yaw = np.arctan2(im, re)
# Draw rotated box
kitti_bev_utils.drawRotatedBox(img_bev, x, y, w, l, yaw,
cnf.colors[int(cls_pred)])
img_rgb = cv2.imread(img_paths[0])
calib = kitti_data_utils.Calibration(img_paths[0].replace(
".png", ".txt").replace("image_2", "calib"))
objects_pred = predictions_to_kitti_format(img_detections, calib,
img_rgb.shape,
configs.img_size)
img_rgb = show_image_with_boxes(img_rgb, objects_pred, calib,
False)
img_bev = cv2.flip(cv2.flip(img_bev, 0), 1)
out_img = merge_rgb_to_bev(img_rgb, img_bev, output_width=608)
print(
def client_recv(self, client, address):
while True:
# read message from socket
# client_msg_0\x00\x00\x00\x00\x00...
msg = client.recv(1024).decode("utf-8")
msg = msg.rstrip("\x00")
if msg == '':
return
if msg == "EOF":
return
elif msg == "quit_client":
client.close()
# self.sock.close()
print("> client exit...")
return
elif msg == "quit_server":
client.close()
self.sock.close()
print("> server exit...")
sys.exit(0)
else:
print("> -------", time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time())), "-------")
print("> receive the msg from client : {0}".format(msg))
print('> inference for {0}'.format(msg))
if(self.need_create_window):
# NOTE ObjSLAM
cv2.namedWindow("YOLO", flags=cv2.WINDOW_GUI_NORMAL)
self.need_create_window = False
# Inference
with torch.no_grad():
# img_paths, imgs_bev = self.test_dataloader_iter.next()
img_paths, imgs_bev = self.test_dataset[int(msg)]
img_paths = [img_paths]
imgs_bev = torch.from_numpy(
np.expand_dims(imgs_bev, axis=0))
input_imgs = imgs_bev.to(
device=self.configs.device).float()
outputs = self.model(input_imgs)
detections = post_processing_v2(
outputs, conf_thresh=self.configs.conf_thresh, nms_thresh=self.configs.nms_thresh)
img_detections = [] # Stores detections for each image index
img_detections.extend(detections)
img_bev = imgs_bev.squeeze() * 255
img_bev = img_bev.permute(1, 2, 0).numpy().astype(np.uint8)
img_bev = cv2.resize(
img_bev, (self.configs.img_size, self.configs.img_size))
for detections in img_detections:
if detections is None:
continue
# Rescale boxes to original image
detections = rescale_boxes(
detections, self.configs.img_size, img_bev.shape[:2])
for x, y, w, l, im, re, *_, cls_pred in detections:
yaw = np.arctan2(im, re)
# Draw rotated box
kitti_bev_utils.drawRotatedBox(
img_bev, x, y, w, l, yaw, cnf.colors[int(cls_pred)])
img_rgb = cv2.imread(img_paths[0])
calib = kitti_data_utils.Calibration(img_paths[0].replace(
".png", ".txt").replace("image_2", "calib"))
objects_pred = predictions_to_kitti_format(
img_detections, calib, img_rgb.shape, self.configs.img_size)
# NOTE: 输出json的代码
frame_object_list = []
for i in objects_pred:
frame_object = dict()
frame_object['type'] = i.type
frame_object['center'] = i.t
frame_object['length'] = i.l
frame_object['width'] = i.w
frame_object['height'] = i.h
frame_object['theta'] = i.ry
box3d_pts_2d, _ = kitti_data_utils.compute_box_3d(
i, calib.P)
if box3d_pts_2d is None:
frame_object['box3d_pts_2d'] = box3d_pts_2d
elif box3d_pts_2d.size == 16:
frame_object['box3d_pts_2d'] = box3d_pts_2d
else:
frame_object['box3d_pts_2d'] = box3d_pts_2d[:8, :]
frame_object_list.append(frame_object)
result = json.dumps(frame_object_list, cls=NumpyEncoder)
img_bev = cv2.flip(cv2.flip(img_bev, 0), 1)
scale = 1.5
cv2.resizeWindow("YOLO",
width=int(img_bev.shape[1] * scale),
height=int(img_bev.shape[0] * scale))
cv2.imshow('YOLO', img_bev)
cv2.waitKey(10)
self.batch_idx += 1
if len(result) > self.configs.max_length:
print("> WARNING: STRING IS TOO LONG! (MAX_LENGTH {0})".format(
self.configs.max_length))
client.send(result.encode(encoding='utf-8'))
print("> send the responce back to client, string length: {0}".format(
len(result)))
return
