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