示例#1
0
def stream(cfg,
           classes_file,
           weights,
           socket_ip,
           socket_port,
           image_size=128,
           confidence_threshold=0.6,
           nms_thres=0.5):
    print('+ Initializing model')
    model = Darknet(cfg, image_size)
    print('+ Loading model')
    load_darknet_weights(model, weights)
    print('+ Fusing model')
    model.fuse()
    print('+ Loading model to CPU')
    model.to('cpu').eval()
    print('+ Loading webcam')
    cap = LoadKinect(img_size=image_size)
    print('+ Loading classes')
    classes = load_classes(classes_file)
    colors = [[random.randint(0, 255) for _ in range(3)]
              for _ in range(len(classes))]
    print('+ Connecting to remote socket')
    global sock
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((socket_ip, socket_port))
    print('+ Enumerating cam')
    for counter, (path, img, im0, vid_cap) in enumerate(cap):
        t = time.time()

        print('+ Loading image to CPU')
        img = torch.from_numpy(img).unsqueeze(0).to('cpu')
        pred, _ = model(img)
        print('+ Detecting objects')
        det = non_max_suppression(pred, confidence_threshold, nms_thres)[0]

        if det is not None and len(det) > 0:
            detected_classes = []
            print('+ Rescaling model')
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
                                      im0.shape).round()

            print('+ Reading depth')

            depth = get_depth()
            depth_swap = np.swapaxes(depth, 0, 1)

            depth_strip1d = np.array([
                np.sort(stripe)[100] for stripe in depth_swap
            ]).astype(np.uint8)
            depth_strip2d_swap = np.array([
                np.ones(depth_swap.shape[1]) * depth for depth in depth_strip1d
            ]).astype(np.uint8)
            depth_strip2d = np.swapaxes(depth_strip2d_swap, 0, 1)

            depth_edge1d = np.zeros(depth_strip1d.shape)

            state = False
            for counter, _ in np.ndenumerate(depth_edge1d[:-1]):
                state = True if not state and depth_strip1d[
                    counter] < 230 else False
                depth_edge1d[counter[0]] = not state

            state = False
            state_cnt = 0
            for counter, _ in np.ndenumerate(depth_edge1d[:-1]):
                counter = counter[0]
                if depth_edge1d[counter] == state:
                    state_cnt += 1
                else:
                    if state_cnt < 10:
                        for r in range(max(0, counter - 10), counter):
                            depth_edge1d[counter] = state
                    state_cnt = 0
                    state = depth_edge1d[counter]

            depth_edge1d = depth_edge1d * 255

            depth_edge2d_swap = np.array([
                np.ones(100) * awddawd for awddawd in depth_edge1d
            ]).astype(np.uint8)
            depth_edge2d = np.swapaxes(depth_edge2d_swap, 0, 1)

            for *coordinates, conf, cls_conf, cls in det:
                if classes[int(cls)] in RISKY_CLASSES:
                    label = '%s %.2f' % (classes[int(cls)], conf)
                    plot_one_box(coordinates,
                                 im0,
                                 label=label,
                                 color=colors[int(cls)])
                    print(f"+ Detected {classes[int(cls)]}")
                    x_avg_depth = np.mean(depth[coordinates[0] -
                                                5:coordinates[0] + 5])
                    y_avg_depth = np.mean(depth[coordinates[1] -
                                                5:coordinates[1] + 5])
                    detected_classes.append({
                        classes[int(cls)]: {
                            'x': coordinates[0],
                            'y': coordinates[1],
                            'z':
                            np.average(np.array([x_avg_depth, y_avg_depth]))
                        }
                    })

            n = []
            for counter in detected_classes:
                width = im0.shape[1]
                x, y, z = counter[list(counter.keys())[0]].values()
                phi = (x / width * 2 - 1) * (CAMERA_FOV / 2)
                n.append(f"{list(counter.keys())[0]};{phi};{z}|")
            sock.send(''.join(str(x) for x in n)[:-1].encode('utf-8'))
        print('+ Cycle took %.3fs' % (time.time() - t))
        plt.imshow(bgr_to_rgb(im0))
        plt.show(block=False)
        plt.pause(.001)
def app():
    cfg = 'ml-data/yolov3.cfg'
    global image_size
    image_size = 320
    weights = 'ml-data/weights/yolov3.weights'
    classes_file = 'ml-data/classes.txt'
    socket_ip = '10.10.10.1'
    # socket_ip = '127.0.0.1'
    socket_port = 1337

    print('+ Initializing model')
    global model
    model = Darknet(cfg, image_size)
    print('+ Loading model')
    load_darknet_weights(model, weights)
    print('+ Fusing model')
    model.fuse()
    print('+ Loading model to CPU')
    model.to('cpu').eval()
    print('+ Loading classes')
    global classes
    classes = load_classes(classes_file)
    global colors
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(classes))]
    print('+ Connecting to remote socket')
    global sock
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((socket_ip, socket_port))

    while 1:
        #
        # Depth
        #
        depth_result = analyse_depth()
        depth_raw = depth_result["raw"]
        depth_done = depth_result["done"]
        depth_objects = depth_result["objects"]

        #
        # RGB
        #
        rgb_result = analyse_rgb()
        rgb_raw = rgb_result["raw"]
        rgb_done = rgb_result["done"]
        rgb_objects = rgb_result["objects"]

        print("FRAME [D]: " + depth_objects)
        print("FRAME [C]: " + rgb_objects)

        sock.send(bytes(f"{rgb_objects}|{depth_objects}".encode('utf-8')))

        time.sleep(0.01)

        # Plot
        vbar = np.zeros((depth_raw.shape[0], 5, 3)).astype(np.uint8)
        depthbar = np.concatenate((depth_raw, vbar, depth_done), axis=1)
        rgbbar = np.concatenate((rgb_raw, vbar, rgb_done), axis=1)
        hbar = np.zeros((5, depthbar.shape[1], 3)).astype(np.uint8)
        cv2.imshow('LineWarn', np.concatenate((depthbar, hbar, rgbbar), axis=0))

        if cv2.waitKey(10) == 27:
            break