Пример #1
0
    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)
Пример #2
0
            t2 = time_synchronized()
            detections = post_processing(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_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)
Пример #3
0
    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
Пример #4
0
            img_detections = []  # Stores detections for each image index
            img_detections.extend(detections)
            bev_maps = torch.squeeze(bev_maps).numpy()
            RGB_Map = np.zeros((cnf.BEV_WIDTH, cnf.BEV_WIDTH, 3))
            RGB_Map[:, :, 2] = bev_maps[0, :, :]  # r_map
            RGB_Map[:, :, 1] = bev_maps[1, :, :]  # g_map
            RGB_Map[:, :, 0] = bev_maps[2, :, :]  # b_map

            RGB_Map = (255 * RGB_Map).astype(np.uint8)
            for detections in img_detections:
                if detections is None:
                    continue
                # Rescale boxes to original image
                detections = np.array(detections)
                print('detections shape: {}'.format(detections.shape))
                detections = rescale_boxes(detections, configs.img_size, RGB_Map.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(RGB_Map, x, y, w, l, yaw, cnf.colors[int(cls_pred)])

            img2d = 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, img2d.shape, configs.img_size)
            img2d = show_image_with_boxes(img2d, objects_pred, calib, False)

            cv2.imshow("bev img", RGB_Map)
            cv2.imshow("img2d", img2d)

            if cv2.waitKey(0) & 0xFF == 27:
                break