Exemplo n.º 1
0
def record(r_led, g_led, event_manager, processing_complete, processed_frames):
    print('recording')
    g_led.Update(100, force=True)
    r_led.Update(1, force=True)
    processing_complete.wait()
    processing_complete.clear()
    event_manager.update(recording)

    errors = 0
    while True:  # wait here until all frames have been processed and sent to server
        try:
            # get the processed frames from queue
            n_frame, ball_candidates = processed_frames.get_nowait()
            # print(f"rec: {n_frame}")
            message = sf.MyMessage(c.TYPE_BALLS, (n_frame, ball_candidates))
            if not sf.send_message(client.client_socket, message, c.CLIENT):
                errors += 1
                print(f"error: {errors}")

        # if processing complete and no more data to send to server
        except queue.Empty:
            if processing_complete.is_set() and (processed_frames.qsize()
                                                 == 0):
                event_manager.clear()
                break
    # if there were no transmission errors send True, else send False
    if errors == 0:
        message = sf.MyMessage(c.TYPE_DONE, True)
    else:
        message = sf.MyMessage(c.TYPE_DONE, False)
    sf.send_message(client.client_socket, message, c.CLIENT)
    return
Exemplo n.º 2
0
def stream(r_led, g_led, event_manager, processing_complete, processed_frames):
    print('recording stream')
    message_list = []
    g_led.Update(1, force=True)
    r_led.Update(1, force=True)
    processing_complete.wait()
    processing_complete.clear()
    event_manager.update(record_stream)
    errors = 0

    while True:
        try:
            # get the frames from queue
            n_frame, frame_buf = processed_frames.get_nowait()

            if event_manager.record_stream.is_set():
                print(f"calib: {n_frame}")
                # y_data is a numpy array hxw with 8 bit greyscale brightness values
                y_data = np.frombuffer(frame_buf,
                                       dtype=np.float32,
                                       count=w * h).reshape(
                                           (h, w)).astype(np.uint8)
                message = sf.MyMessage(c.TYPE_IMG, (n_frame, y_data))
                if not sf.send_message(client.client_socket, message,
                                       c.CLIENT):
                    errors += 1
                    print(f"error: {errors}")

                message_list.extend(client.read_all_server_messages())
                for message in message_list:
                    if message['data'].type == c.TYPE_DONE:
                        event_manager.clear()

        except queue.Empty:
            if not event_manager.record_stream.is_set(
            ) and processed_frames.qsize(
            ) == 0:  # if the recording has finished
                print('stream done')
                processing_complete.set()
                break

    # if there were no transmission errors send True, else send False
    if errors == 0:
        message = sf.MyMessage(c.TYPE_DONE, True)
    else:
        message = sf.MyMessage(c.TYPE_DONE, False)

    sf.send_message(client.client_socket, message, c.CLIENT)
    def initialise_picamera(self):
        '''
        Sends message to both clients to start the picameras
        Waits for confirmation of camera startup

        Return: True if cameras started
        False: Cameras not started
        '''
        print('Initialising cameras')
        message = sf.MyMessage(c.TYPE_START_CAM, None)
        setup_complete = {}

        for client in self.client_names:
            self.send_to_client(client, message)
            setup_complete[client] = False

        while True:
            message_list = self.read_client_messages()
            for message in message_list:
                check_task_trigger(setup_complete, c.TYPE_START_CAM, message)

            if check_task_complete(setup_complete):
                print('Cameras initialised')
                return True

            time.sleep(0.001)
def shutdown():
    print('shutting down')
    shut_obj = sf.MyMessage(c.TYPE_SHUTDOWN, None)
    send_to_client(c.LEFT_CLIENT, shut_obj)
    send_to_client(c.RIGHT_CLIENT, shut_obj)
    while True:
        time.sleep(2)
        sys.exit()
def record(stereo_calib = None, record_time = c.REC_T):
    if stereo_calib is None:
        print('stereo cameras must first be calibrated')
        return False

    message_list = []
    pos = 0
    left_done = False
    right_done = False

    right_frames = 0
    left_frames = 0

    rec_obj = sf.MyMessage(c.TYPE_REC, record_time)
    send_to_client(c.LEFT_CLIENT, rec_obj)
    send_to_client(c.RIGHT_CLIENT, rec_obj)

    left_candidates = record_time*100*[None]
    right_candidates = record_time*100*[None]

    while True:
        message_list.extend(read_all_client_messages())

        while pos < len(message_list):
            if message_list[pos]['data'].type == c.TYPE_BALLS:

                if message_list[pos]['client'] == c.LEFT_CLIENT:
                    f = int(message_list[pos]['data'].message[0])
                    rectify_points(message_list[pos]['data'].message[1], stereo_calib.cameraMatrix1, stereo_calib.distCoeffs1, stereo_calib.R1, stereo_calib.P1)
                    left_candidates[f] = message_list[pos]['data'].message[1]
                    left_frames+=1

                elif message_list[pos]['client'] == c.RIGHT_CLIENT:
                    f = int(message_list[pos]['data'].message[0])
                    rectify_points(message_list[pos]['data'].message[1], stereo_calib.cameraMatrix2, stereo_calib.distCoeffs2, stereo_calib.R2, stereo_calib.P2)
                    right_candidates[f] = message_list[pos]['data'].message[1]
                    right_frames+=1
                
            elif message_list[pos]['data'].type == c.TYPE_DONE:
                if message_list[pos]['client'] == c.LEFT_CLIENT:
                    left_done = True
                elif message_list[pos]['client'] == c.RIGHT_CLIENT:
                    right_done = True
            else:
                print(f"unknown message format for recording: {message_list[pos]['data'].type}")
            pos+=1

        if left_done and right_done:
            print('recording finished')
           # -- Make both lists the same length -- #
            min_length = min(left_frames, right_frames)
            left_candidates = left_candidates[:min_length]
            right_candidates = right_candidates[:min_length]

            print(f"frames captured: {min_length}")

            return left_candidates, right_candidates
    def send_to_server(self, message_type, message_data):
        '''
        Send message to the server given the message type and data

        Return True if sent successfully
        '''
        message = sf.MyMessage(message_type, message_data)
        sf.send_message(self.socket, message, c.CLIENT)
        return True
    def stream(self, stream_t, show=False, save=False):
        '''
        Sends a message to both clients to start a stream for stream_t seconds
        '''
        stream_t = convert_string_to_float(stream_t)
        if stream_t is None: return False

        print(f"Streaming for {stream_t} seconds")
        message = sf.MyMessage(c.TYPE_STREAM, (stream_t, ))

        stream_complete = {}
        for client in self.client_names:
            self.send_to_client(client, message)
            stream_complete[client] = False

        img_queue = mp.Queue()
        img_play = mp.Event()

        image_process = mp.Process(target=image_viewer,
                                   args=(img_queue, img_play,
                                         c.STREAM_DELTA_T))
        image_process.start()
        img_play.set()

        img_dict = dict.fromkeys(self.client_names)
        for key in img_dict.keys():
            img_dict[key] = {}

        cur_frame = 0

        while True:
            message_list = self.read_client_messages()
            for message in message_list:
                if not check_task_trigger(stream_complete, c.TYPE_DONE, message) \
                                        and check_img(message):
                    if save:
                        save_img(message)

                    if show:
                        add_to_img_dict(img_dict, message)
                        img = None
                        img = combine_img(img_dict, cur_frame)

                        if img is not None:
                            show_img(img, img_queue)
                            cur_frame += c.STREAM_IMG_DELTA

            if check_task_complete(stream_complete) and img_queue.qsize() == 0:
                print('Stream complete')
                cv2.destroyAllWindows()
                image_process.terminate()

                return True

            time.sleep(0.001)
    def shutdown(self):
        '''
        Sends a shutdown command to both clients then shuts down the server

        Return: 
        '''
        print('Shutting down...')
        message = sf.MyMessage(c.TYPE_SHUTDOWN, None)
        for client_name in self.client_names:
            self.send_to_client(client_name, message)

        sys.exit()
    def record(self, record_t):
        '''
        Sends a message to both clients to start a recording for record_t seconds
        '''
        record_t = convert_string_to_float(record_t)

        if record_t is not None and isinstance(record_t, float):
            if record_t > 0 and record_t <= c.REC_T_MAX:
                record_t = record_t
            else:
                print(f"ERROR: {record_t}s is too long for recording")
                record_t = c.REC_T_MAX
        else:
            return False

        print(f"Recording for {record_t} seconds")
        message = sf.MyMessage(c.TYPE_RECORD, (record_t, ))

        recording_complete = {}
        for client in self.client_names:
            self.send_to_client(client, message)
            recording_complete[client] = False

        ball_candidate_dict = dict.fromkeys(self.client_names)
        for key in ball_candidate_dict.keys():
            ball_candidate_dict[key] = {}

        while True:
            message_list = self.read_client_messages()
            for message in message_list:
                if not check_task_trigger(recording_complete, c.TYPE_DONE,
                                          message):
                    check_ball_cand(message, ball_candidate_dict,
                                    self.stereo_calib)

            if check_task_complete(recording_complete):
                print('Recording complete')
                with open("ball_dict.pkl", "wb") as file:
                    pickle.dump(ball_candidate_dict, file)

                points_3d = tr.triangulate_points(ball_candidate_dict,
                                                  self.stereo_calib)

                plot_points_3d(points_3d)
                np.save("points_3d.npy", points_3d)

                trajectory = analyse_trajectory(points_3d)
                if trajectory is not None:
                    plot_trajectory(trajectory, save=False)

                return True

            time.sleep(0.001)
def stream(run_time = c.CALIB_T, calibrate = False, display = False, timeout = False):
    message_list = []
    left_stream_imgs = []
    right_stream_imgs = []
    left_chessboards = []
    right_chessboards = []
    chessboards = mp.Queue()
    chessboard_searchers = []
    pos = 0
    left_done = False
    right_done = False
    disp_n_frame = 0
    cal_n_frame = 0
    img_size = None
    done = False
    stopframe = int(run_time*c.FRAMERATE)
    chessboards_found = mp.Value('i',0)

    rec_obj = sf.MyMessage(c.TYPE_STREAM, c.CALIB_IMG_DELAY)
    send_to_client(c.LEFT_CLIENT, rec_obj)
    send_to_client(c.RIGHT_CLIENT, rec_obj)

    if calibrate:
        # load camera intrinsic calibration data
        left_cal, right_cal = s_cal.load_calibs()
        
    while True:
        message_list.extend(read_all_client_messages())
        if len(message_list) > 0:
            while pos < len(message_list):
                # when the clients send an image during calibration
                if (message_list[pos]['data'].type == c.TYPE_IMG) and not done:
                    n_frame = message_list[pos]['data'].message[0]
                    print(n_frame)
                    y_data = message_list[pos]['data'].message[1]
                    if img_size is None:
                        (h,w) = y_data.shape[:2]
                        img_size = (w,h)
                    # add the img to the corresponding calibration img list
                    if message_list[pos]['client'] == c.LEFT_CLIENT:
                        left_stream_imgs.append((n_frame, y_data))
                    elif message_list[pos]['client'] == c.RIGHT_CLIENT:                    
                        right_stream_imgs.append((n_frame, y_data))
                    # cv2.imwrite(f"{message_list[pos]['client']}{n_frame}.png",y_data)

                    if display:
                        if (len(left_stream_imgs) > disp_n_frame) and (len(right_stream_imgs) > disp_n_frame): 
                            
                            ########## FOR TESTING ##############
                            os.chdir(c.IMG_P)
                            cv2.imwrite(f"l_{(disp_n_frame):04d}.png",left_stream_imgs[disp_n_frame][1])
                            cv2.imwrite(f"r_{(disp_n_frame):04d}.png",right_stream_imgs[disp_n_frame][1])
                            os.chdir(c.ROOT_P)
                            ########## FOR TESTING ##############

                            disp_frame = cv2.hconcat([left_stream_imgs[disp_n_frame][1],right_stream_imgs[disp_n_frame][1]])
                            cv2.imshow(f"stream", disp_frame)
                            print(disp_n_frame)
                            cv2.waitKey(100)
                            if left_stream_imgs[disp_n_frame][0] >=stopframe and timeout:
                                done_obj = sf.MyMessage(c.TYPE_DONE, 1)
                                send_to_client(c.LEFT_CLIENT, done_obj)
                                send_to_client(c.RIGHT_CLIENT, done_obj)
                                done = True
                                cv2.destroyAllWindows()
                            
                            disp_n_frame += 1

                    # look for chessboards
                    if calibrate:
                        if (len(left_stream_imgs) > cal_n_frame) and (len(right_stream_imgs) > cal_n_frame):
                            chessboard_search = mp.Process(target = search_for_chessboards, args=(chessboards_found, chessboards, [left_stream_imgs[cal_n_frame]], [right_stream_imgs[cal_n_frame]]))
                            chessboard_search.start()
                            chessboard_searchers.append(chessboard_search)
                            cal_n_frame += 1

                    if chessboards_found.value >= c.MIN_PATTERNS:
                        done_obj = sf.MyMessage(c.TYPE_DONE, 1)
                        send_to_client(c.LEFT_CLIENT, done_obj)
                        send_to_client(c.RIGHT_CLIENT, done_obj)
                        if display: 
                            done = True
                            cv2.destroyAllWindows()

                # when both clients send the done message, they are finished collecting frames
                elif (message_list[pos]['data'].type == c.TYPE_DONE):
                    if message_list[pos]['client'] == c.LEFT_CLIENT:
                        left_done = True
                    elif message_list[pos]['client'] == c.RIGHT_CLIENT:
                        right_done = True
                    if left_done and right_done:
                        if calibrate and chessboards_found.value >= c.MIN_PATTERNS:
                            # for searcher in chessboard_searchers:
                            #     searcher.join()
                            while True:
                                try:
                                    left_chessboard, right_chessboard = chessboards.get_nowait()
                                    left_chessboards.append(left_chessboard)
                                    right_chessboards.append(right_chessboard)
                                except queue.Empty:
                                    if chessboards.qsize() == 0:
                                        break

                            # # check all chessboards are valid in both images
                            s_cal.validate_chessboards(left_chessboards, right_chessboards)
                            # calibrated stereo cameras
                            RMS, cameraMatrix1, distCoeffs1, cameraMatrix2, distCoeffs2, R, T, E, F = s_cal.calibrate_stereo(
                                left_chessboards, right_chessboards, left_cal, right_cal, img_size)
    
                            # obtain stereo rectification projection matrices
                            R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(cameraMatrix1, distCoeffs1,
                                                        cameraMatrix2, distCoeffs2, img_size, R, T)

                            # save all calibration params to object
                            stereo_calib =  s_cal.StereoCal(RMS, cameraMatrix1, distCoeffs1, cameraMatrix2, distCoeffs2, R, T, E, F,
                                                        R1, R2, P1, P2, Q, validPixROI1, validPixROI2)

                            s_cal.save_stereo_calib(stereo_calib)
                        
                            print(f'calibration complete, rms: {RMS}')
                            return stereo_calib
                        else:
                            return None
                pos += 1