Exemplo n.º 1
0
 def get_coords():
     if traj is None:
         return False, (0, 0)
     else:
         idx, landing_point = tr.get_landing_point(traj)
         if landing_point is not None:
             return True, perspective_transform(landing_point,
                                                perspective_matrix)
     return False, (0, 0)
Exemplo n.º 2
0
def test1(input_path, out_prefix=None):
    logger.info('test1(\'{}\') started'.format(input_path))
    cap = cv2.VideoCapture(input_path)
    # Define the codec and create VideoWriter object
    fourcc = cv2.cv.CV_FOURCC(*'XVID')
    fps = cap.get(cv2.cv.CV_CAP_PROP_FPS)
    width = cap.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)

    if out_prefix is None:
        out_prefix = os.path.join('output', os.path.basename(input_path))

    def write_video(path, frames):
        if frames:
            video_io.write_video(path, frames, fps,
                                 frame_size=(int(width), int(height)))
            logger.info("Video written to file {}".format(path))


    frames, mask_ball = \
        video_util.extract_ball_from_capture(cap,
                                             skip_count=10,
                                             get_mask=video_util.get_ball_mask,
                                             carpet_lowerb=(50, 0, 0),
                                             carpet_upperb=(70, 255, 255),
                                             ball_lowerb=(18, 50, 0),
                                             ball_upperb=(35, 255, 255)
                                             )
    cap.release()

    logger.info("Frames captured: {}".format(len(frames)))

    # write_video(out_prefix + '_frames.avi', frames)
    write_video(out_prefix + '_ball.avi', mask_ball)

    traj = tr.get_ball_trajectory(mask_ball)
    if traj:
        logger.info('trajectory found')
    else:
        logger.info('trajectory not found')

    for x, y in traj:
        center = (int(round(x)), int(round(y)))
        for frame in frames:
            cv2.circle(frame, center, 3, (0, 0, 255), thickness=-1)

    landing_point = tr.get_landing_point(traj)
    if landing_point:
        logger.info('landing point: {}'.format(landing_point))
        for frame in frames:
            x, y = landing_point
            center = (int(round(x)), int(round(y)))
            cv2.circle(frame, center, radius=10, color=(255, 255, 255),
                       thickness=1)

    landing_point = tr.get_landing_point_by_acceleration(traj)
    if landing_point:
        logger.info('landing point by acceleration: {}'.format(landing_point))
        for frame in frames:
            x, y = landing_point
            center = (int(round(x)), int(round(y)))
            cv2.circle(frame, center, radius=4, color=(255, 0, 0),
                       thickness=-1)
    else:
        logger.info('landing point not found')

    landing_point = tr.get_landing_point_by_angle(traj)
    if landing_point:
        logger.info('landing point by angle: {}'.format(landing_point))
        for frame in frames:
            x, y = landing_point
            center = (int(round(x)), int(round(y)))
            cv2.circle(frame, center, radius=4, color=(0, 255, 0),
                       thickness=-1)
    else:
        logger.info('landing point by angle not found')

    write_video(out_prefix + '_frames_ball.avi', frames)
    # write_video(out_prefix + '_mask_frames.avi', mask_ball)

    acc = tr.get_acceleration(traj)

    with io.open(out_prefix + '_traj.dat', mode='w') as handle:
        for i, (x, y) in enumerate(traj):
            ax, ay = acc[i]
            handle.write(u'{} {} {} {} {}\n'.format(i, x, y, ax, ay))
Exemplo n.º 3
0
def program(out_dir):
    cap = cv2.VideoCapture(0)
    cap.set(cv2.cv.CV_CAP_PROP_FPS, 30.0)
    # fps = cap.get(cv2.cv.CV_CAP_PROP_FPS)
    # cap.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, 1920)
    # cap.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, 1080)
    # cap.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, 1280)
    # cap.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, 720)
    cap.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, 864)
    cap.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, 480)

    perspective_matrix = get_perspective_transform(cap)
    #
    logger.info('matrix of perspective transform: {}'.format(
        perspective_matrix))

    img = read_from_capture(cap)
    cv2.imwrite(os.path.join(out_dir, 'img.jpg'), img)

    # Find carpet
    lowerb, upperb = gui.get_hls_range(img,
                                       winname='Choose HLS range for carpet')

    logger.info('Carpet HLS range: {} - {}'.format(lowerb, upperb))

    carpet_mask = img_util.green_carpet_mask(img, lowerb, upperb)
    cv2.imwrite(os.path.join(out_dir, 'carpet_mask.jpg'), carpet_mask)

    # Find target
    img = read_from_capture(cap)
    lowerb, upperb = gui.get_hls_range(img,
                                       winname='Choose HLS range for target')
    logger.info('Target HLS range: {} - {}'.format(lowerb, upperb))

    ret, target_contour = img_util.target_contour(img, lowerb, upperb,
                                                  carpet_mask)
    logging.debug('target contour: {}'.format(target_contour))
    if not ret:
        raise ValueError('Cannot find target')

    target_contour_plane = perspective_transform_contour(target_contour,
                                                         perspective_matrix)
    m = cv2.moments(target_contour_plane)
    mass = m['m00']
    # x,y coordinates on plane
    if mass > 0:
        target_coords = m['m10'] / mass, m['m01'] / mass
        logger.info(
            'target coordinates on plane: (x,y) = {}'.format(target_coords))
    else:
        raise ValueError('Cannot find target center')

    cv2.drawContours(img, contours=[target_contour], contourIdx=-1,
                     color=(255, 255, 255),
                     thickness=2)
    img_carpet = cv2.bitwise_and(img, img, mask=carpet_mask)
    cv2.addWeighted(img, 0.25, img_carpet, 0.75, 0, dst=img_carpet)

    cv2.imwrite(os.path.join(out_dir, 'carpet_target.jpg'), img_carpet)

    # Find ball
    img = read_from_capture(cap)
    ball_lowerb, ball_upperb = gui.get_hls_range(img,
                                                 winname='Choose HLS range for ball')

    logger.info('Ball HLS range: {} - {}'.format(ball_lowerb, ball_upperb))

    conn = robot.RobotConnector(winscp_path='C:/TRIKStudio/winscp/WinSCP.com')
    trik = robot.Robot(conn, angle_to_encoder=180 / np.pi)

    traj = None

    def get_coords():
        if traj is None:
            return False, (0, 0)
        else:
            idx, landing_point = tr.get_landing_point(traj)
            if landing_point is not None:
                return True, perspective_transform(landing_point,
                                                   perspective_matrix)
        return False, (0, 0)

    gun_params = {'x': 300, 'y': 50, 'z': 97, 'v': 300, 'g': 981,
                  'alpha_0': 15 * pi / 180, 'phi_0': 0.0, 'gun_length': 10}
    gun_layer = GunLayer(rotate_and_shoot=trik.rotate_and_shoot,
                         get_coords=get_coords, target=target_coords,
                         gun_params=gun_params)

    for i in xrange(10):
        out_prefix = os.path.join(out_dir, str(i))

        commands = tuple()
        stdout, stderr = trik.open_and_trikRun(*commands)
        logger.debug('trik: sent command = {}'.format(commands))
        logger.debug('trik: stdout = {}'.format(stdout))
        logger.debug('trik: stderr = {}'.format(stderr))

        logger.info('ready to shoot...')
        gun_layer.shoot_at_target()
        logger.info('ready to capture...')
        frames, mask_ball = video_util.extract_ball_from_capture(cap,
                                                                 max_frames_count=30,
                                                                 skip_count=0,
                                                                 carpet_mask=carpet_mask,
                                                                 get_mask=video_util.get_ball_mask,
                                                                 ball_lowerb=ball_lowerb,
                                                                 ball_upperb=ball_upperb,
                                                                 ball_size=3
                                                                 )

        logger.debug("Frames captured: {}".format(len(frames)))

        traj = tr.get_ball_trajectory(mask_ball)
        if not traj:
            logger.info('trajectory not found')
        else:
            logger.info('trajectory found')
            for x, y in traj:
                center = (int(round(x)), int(round(y)))
                for frame in frames:
                    cv2.circle(frame, center, 3, (0, 0, 255), thickness=-1)
            idx, landing_point = tr.get_landing_point(traj)
            if landing_point is not None:
                logger.info('landing point on plane: {}'.format(
                    perspective_transform(landing_point, perspective_matrix)))

                for frame in frames:
                    x, y = landing_point
                    center = (int(round(x)), int(round(y)))
                    cv2.circle(frame, center, radius=10, color=(255, 255, 255),
                               thickness=1)
        video_io.write_video(out_prefix + '_frames_ball.avi', frames)
        video_io.write_video(out_prefix + '_ball_mask.avi', mask_ball)

    cap.release()