def process_frame(frame, prev_mask, carpet_mask): if carpet_mask is None: carpet_mask = img_util.green_carpet_mask(frame, carpet_lowerb, carpet_upperb) mask = get_mask(frame, lowerb=ball_lowerb, upperb=ball_upperb, carpet_mask=carpet_mask) if prev_mask is None: move_mask = np.zeros_like(mask) else: move_mask = cv2.bitwise_not(prev_mask, mask=mask) prev_mask = mask is_ball, move_mask = img_util.detect_ball(move_mask) return is_ball, move_mask, prev_mask, frame
def test_orange_mask(input_path, out_prefix=None): if out_prefix is None: out_prefix = os.path.join('output', os.path.basename(input_path)) frames_orig = video_io.get_frames(input_path) carpet_mask = img_util.green_carpet_mask(frames_orig[30], lowerb_hls=(55, 0, 0), upperb_hls=(70, 255, 255)) cv2.imwrite(out_prefix + '_carpet.jpg', carpet_mask) frames_orange = [ video_util.preprocess(frame, carpet_mask=carpet_mask) for frame in frames_orig] write_video(out_prefix + '_orange_mask.avi', frames_orange) return
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()
def extract_ball_from_capture(cap, max_frames_count=-1, skip_count=0, carpet_mask=None, get_mask=None, carpet_lowerb=None, carpet_upperb=None, ball_lowerb=None, ball_upperb=None, ball_size=3): """ Read frames from capture until we detect motion of the ball Return tuple of (original frames, ball mask frames) :param carpet_lowerb: :rtype : (list(np.ndarray), list(np.ndarray)) """ frames = [] mask_frames = [] # mog = cv2.BackgroundSubtractorMOG(history=5, nmixtures=4,backgroundRatio=0.7) mog = cv2.BackgroundSubtractorMOG() motion_started = False for _ in xrange(skip_count): if cap.isOpened(): cap.read() prev_mask = None while cap.isOpened(): if 0 < max_frames_count <= len(frames): logger.debug('max frames count reached') break ret, frame = cap.read() if ret: if get_mask is not None: if carpet_mask is None: carpet_mask = img_util.green_carpet_mask(frame, carpet_lowerb, carpet_upperb) mask = get_mask(frame, lowerb=ball_lowerb, upperb=ball_upperb, carpet_mask=carpet_mask) if prev_mask is None: move_mask = np.zeros_like(mask) else: move_mask = cv2.bitwise_not(prev_mask, mask=mask) prev_mask = mask else: logger.debug('get_mask = None') move_mask = mog.apply(frame) is_ball, move_mask = img_util.detect_ball(move_mask, ball_size) if is_ball: if not motion_started: logger.debug('ball appeared') motion_started = True frames.append(frame) mask_frames.append(move_mask) else: if motion_started: logger.debug( 'ball disappeared. Frames count: {}'.format( len(frames))) break else: continue else: logger.debug('Cannot read more frames') break return frames, mask_frames