chromium / catapult / ed57c0f77cce4c9dff0c66654f39e9e85daf3859 / . / telemetry / telemetry / internal / image_processing / screen_finder.py

#!/usr/bin/env python | |

# Copyright 2014 The Chromium Authors. All rights reserved. | |

# Use of this source code is governed by a BSD-style license that can be | |

# found in the LICENSE file. | |

# | |

# This script attempts to detect the region of a camera's field of view that | |

# contains the screen of the device we are testing. | |

# | |

# Usage: ./screen_finder.py path_to_video 0 0 --verbose | |

from __future__ import division | |

import copy | |

import logging | |

import os | |

import sys | |

from telemetry.internal.image_processing import cv_util | |

from telemetry.internal.image_processing import ( | |

frame_generator as frame_generator_module) | |

from telemetry.internal.image_processing import video_file_frame_generator | |

from telemetry.internal.util import external_modules | |

np = external_modules.ImportRequiredModule('numpy') | |

cv2 = external_modules.ImportRequiredModule('cv2') | |

class ScreenFinder(object): | |

"""Finds and extracts device screens from video. | |

Sample Usage: | |

sf = ScreenFinder(sys.argv[1]) | |

while sf.HasNext(): | |

ret, screen = sf.GetNext() | |

Attributes: | |

_lost_corners: Each index represents whether or not we lost track of that | |

corner on the previous frame. Ordered by [top-right, top-left, | |

bottom-left, bottom-right] | |

_frame: An unmodified copy of the frame we're currently processing. | |

_frame_debug: A copy of the frame we're currently processing, may be | |

modified at any time, used for debugging. | |

_frame_grey: A greyscale copy of the frame we're currently processing. | |

_frame_edges: A Canny Edge detected copy of the frame we're currently | |

processing. | |

_screen_size: The size of device screen in the video when first detected. | |

_avg_corners: Exponentially weighted average of the previous corner | |

locations. | |

_prev_corners: The location of the corners in the previous frame. | |

_lost_corner_frames: A counter of the number of successive frames in which | |

we've lost a corner location. | |

_border: See |border| above. | |

_min_line_length: The minimum length a line must be before we consider it | |

a possible screen edge. | |

_frame_generator: See |frame_generator| above. | |

_width, _height: The width and height of the frame. | |

_anglesp5, _anglesm5: The angles for each point we look at in the grid | |

when computing brightness, constant across frames.""" | |

class ScreenNotFoundError(Exception): | |

pass | |

# Square of the distance a corner can travel in pixels between frames | |

MAX_INTERFRAME_MOTION = 25 | |

# The minimum width line that may be considered a screen edge. | |

MIN_SCREEN_WIDTH = 40 | |

# Number of frames with lost corners before we ignore MAX_INTERFRAME_MOTION | |

RESET_AFTER_N_BAD_FRAMES = 2 | |

# The weight applied to the new screen location when exponentially averaging | |

# screen location. | |

# TODO(mthiesse): This should be framerate dependent, for lower framerates | |

# this value should approach 1. For higher framerates, this value should | |

# approach 0. The current 0.5 value works well in testing with 240 FPS. | |

CORNER_AVERAGE_WEIGHT = 0.5 | |

# TODO(mthiesse): Investigate how to select the constants used here. In very | |

# bright videos, twice as bright may be too high, and the minimum of 60 may | |

# be too low. | |

# The factor by which a quadrant at an intersection must be brighter than | |

# the other quadrants to be considered a screen corner. | |

MIN_RELATIVE_BRIGHTNESS_FACTOR = 1.5 | |

# The minimum average brightness required of an intersection quadrant to | |

# be considered a screen corner (on a scale of 0-255). | |

MIN_CORNER_ABSOLUTE_BRIGHTNESS = 60 | |

# Low and high hysteresis parameters to be passed to the Canny edge | |

# detection algorithm. | |

CANNY_HYSTERESIS_THRESH_LOW = 300 | |

CANNY_HYSTERESIS_THRESH_HIGH = 500 | |

SMALL_ANGLE = 5 / 180 * np.pi # 5 degrees in radians | |

DEBUG = False | |

def __init__(self, frame_generator, border=5): | |

"""Initializes the ScreenFinder object. | |

Args: | |

frame_generator: FrameGenerator, An initialized Video Frame Generator. | |

border: int, number of pixels of border to be kept when cropping the | |

detected screen. | |

Raises: | |

FrameReadError: The frame generator may output a read error during | |

initialization.""" | |

assert isinstance(frame_generator, frame_generator_module.FrameGenerator) | |

self._lost_corners = [False, False, False, False] | |

self._frame_debug = None | |

self._frame = None | |

self._frame_grey = None | |

self._frame_edges = None | |

self._screen_size = None | |

self._avg_corners = None | |

self._prev_corners = None | |

self._lost_corner_frames = 0 | |

self._border = border | |

self._min_line_length = self.MIN_SCREEN_WIDTH | |

self._frame_generator = frame_generator | |

self._anglesp5 = None | |

self._anglesm5 = None | |

if not self._InitNextFrame(): | |

logging.warn('Not enough frames in video feed!') | |

return | |

self._height, self._width = self._frame.shape[:2] | |

def _InitNextFrame(self): | |

"""Called after processing each frame, reads in the next frame to ensure | |

HasNext() is accurate.""" | |

self._frame_debug = None | |

self._frame = None | |

self._frame_grey = None | |

self._frame_edges = None | |

try: | |

frame = next(self._frame_generator.Generator) | |

except StopIteration: | |

return False | |

self._frame = frame | |

self._frame_debug = copy.copy(frame) | |

self._frame_grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) | |

self._frame_edges = cv2.Canny(self._frame_grey, | |

self.CANNY_HYSTERESIS_THRESH_LOW, | |

self.CANNY_HYSTERESIS_THRESH_HIGH) | |

return True | |

def HasNext(self): | |

"""True if there are more frames available to process. """ | |

return self._frame is not None | |

def GetNext(self): | |

"""Gets the next screen image. | |

Returns: | |

A numpy matrix containing the screen surrounded by the number of border | |

pixels specified in initialization, and the location of the detected | |

screen corners in the current frame, if a screen is found. The returned | |

screen is guaranteed to be the same size at each frame. | |

'None' and 'None' if no screen was found on the current frame. | |

Raises: | |

FrameReadError: An error occurred in the FrameGenerator. | |

RuntimeError: This method was called when no frames were available.""" | |

if self._frame is None: | |

raise RuntimeError('No more frames available.') | |

logging.info('Processing frame: %d', | |

self._frame_generator.CurrentFrameNumber) | |

# Finds straight lines in the image. | |

hlines = cv2.HoughLinesP(self._frame_edges, 1, np.pi / 180, 60, | |

minLineLength=self._min_line_length, | |

maxLineGap=100) | |

# Extends these straight lines to be long enough to ensure the screen edge | |

# lines intersect. | |

lines = cv_util.ExtendLines(np.float32(hlines[0]), 10000) \ | |

if hlines is not None else [] | |

# Find intersections in the lines; these are likely to be screen corners. | |

intersections = self._FindIntersections(lines) | |

if len(intersections[:, 0]) > 0: | |

points = np.vstack(intersections[:, 0].flat) | |

if (self._prev_corners is not None and len(points) >= 4 and | |

not self._HasMovedFast(points, self._prev_corners)): | |

corners = self._prev_corners | |

missing_corners = 0 | |

else: | |

# Extract the corners from all intersections. | |

corners, missing_corners = self._FindCorners( | |

intersections, self._frame_grey) | |

else: | |

corners = np.empty((4, 2), np.float32) | |

corners[:] = np.nan | |

missing_corners = 4 | |

screen = None | |

found_screen = True | |

final_corners = None | |

try: | |

# Handle the cases where we have missing corners. | |

screen_corners = self._NewScreenLocation( | |

corners, missing_corners, intersections) | |

final_corners = self._SmoothCorners(screen_corners) | |

# Create a perspective transform from our corners. | |

transform, w, h = self._GetTransform(final_corners, self._border) | |

# Apply the perspective transform to get our output. | |

screen = cv2.warpPerspective( | |

self._frame, transform, (int(w + 0.5), int(h + 0.5))) | |

self._prev_corners = final_corners | |

except self.ScreenNotFoundError as e: | |

found_screen = False | |

logging.info(e) | |

if self.DEBUG: | |

self._Debug(lines, corners, final_corners, screen) | |

self._InitNextFrame() | |

if found_screen: | |

return screen, self._prev_corners | |

return None, None | |

def _FindIntersections(self, lines): | |

"""Finds intersections in a set of lines. | |

Filters pairs of lines that are less than 45 degrees apart. Filtering | |

these pairs helps dramatically reduce the number of points we have to | |

process, as these points could not represent screen corners anyways. | |

Returns: | |

The intersections, represented as a tuple of (point, line, line) of the | |

points and the lines that intersect there of all lines in the array that | |

are more than 45 degrees apart.""" | |

intersections = np.empty((0, 3), np.float32) | |

for i in range(0, len(lines)): | |

for j in range(i + 1, len(lines)): | |

# Filter lines that are less than 45 (or greater than 135) degrees | |

# apart. | |

if not cv_util.AreLinesOrthogonal(lines[i], lines[j], (np.pi / 4.0)): | |

continue | |

ret, point = cv_util.FindLineIntersection(lines[i], lines[j]) | |

point = np.float32(point) | |

if not ret: | |

continue | |

# If we know where the previous corners are, we can also filter | |

# intersections that are too far away from the previous corners to be | |

# where the screen has moved. | |

if self._prev_corners is not None and \ | |

self._lost_corner_frames <= self.RESET_AFTER_N_BAD_FRAMES and \ | |

not self._PointIsCloseToPreviousCorners(point): | |

continue | |

intersections = np.vstack((intersections, | |

np.array((point, lines[i], lines[j])))) | |

return intersections | |

def _PointIsCloseToPreviousCorners(self, point): | |

"""True if the point is close to the previous corners.""" | |

max_dist = self.MAX_INTERFRAME_MOTION | |

if cv_util.SqDistance(self._prev_corners[0], point) <= max_dist or \ | |

cv_util.SqDistance(self._prev_corners[1], point) <= max_dist or \ | |

cv_util.SqDistance(self._prev_corners[2], point) <= max_dist or \ | |

cv_util.SqDistance(self._prev_corners[3], point) <= max_dist: | |

return True | |

return False | |

def _HasMovedFast(self, corners, prev_corners): | |

min_dist = np.zeros(4, np.float32) | |

for i in range(4): | |

dist = np.min(cv_util.SqDistances(corners, prev_corners[i])) | |

min_dist[i] = dist | |

# 3 corners can move up to one pixel before we consider the screen to have | |

# moved. TODO(mthiesse): Should this be relaxed? Resolution dependent? | |

if np.sum(min_dist) < 3: | |

return False | |

return True | |

class CornerData(object): | |

def __init__(self, corner_index, corner_location, brightness_score, line1, | |

line2): | |

self.corner_index = corner_index | |

self.corner_location = corner_location | |

self.brightness_score = brightness_score | |

self.line1 = line1 | |

self.line2 = line2 | |

def __gt__(self, corner_data2): | |

return self.corner_index > corner_data2.corner_index | |

def __repr__(self): | |

return ('\nCorner index: ' + str(self.corner_index) + | |

',\nCorner location: ' + str(self.corner_location) + | |

',\nBrightness score: ' + str(self.brightness_score) + | |

',\nline1: ' + str(self.line1) + ',\nline2: ' + str(self.line2)) | |

def _FindCorners(self, intersections, grey_frame): | |

"""Finds the screen corners in the image. | |

Given the set of intersections in the image, finds the intersections most | |

likely to be corners. | |

Args: | |

intersections: The array of intersections in the image. | |

grey_frame: The greyscale frame we're processing. | |

Returns: | |

An array of length 4 containing the positions of the corners, or nan for | |

each index where a corner could not be found, and a count of the number | |

of missing corners. | |

The corners are ordered as follows: | |

1 | 0 | |

----- | |

2 | 3 | |

Ex. 3 corners are found from a square of width 2 centered at the origin, | |

the output would look like: | |

'[[1, 1], [np.nan, np.nan], [-1, -1], [1, -1]], 1'""" | |

filtered = [] | |

corners = np.empty((0, 2), np.float32) | |

for corner_pos, score, point, line1, line2 in \ | |

self._LooksLikeCorner(intersections, grey_frame): | |

if self.DEBUG: | |

center = (int(point[0] + 0.5), int(point[1] + 0.5)) | |

cv2.circle(self._frame_debug, center, 5, (0, 255, 0), 1) | |

point.resize(1, 2) | |

corners = np.append(corners, point, axis=0) | |

point.resize(2,) | |

corner_data = self.CornerData(corner_pos, point, score, line1, line2) | |

filtered.append(corner_data) | |

# De-duplicate corners because we may have found many false positives, or | |

# near-misses. | |

self._DeDupCorners(filtered, corners) | |

# Strip everything but the corner location. | |

filtered_corners = np.array( | |

[corner_data.corner_location for corner_data in filtered]) | |

corner_indices = [corner_data.corner_index for corner_data in filtered] | |

# If we have found a corner to replace a lost corner, we want to check | |

# that the corner is not erroneous by ensuring it makes a rectangle with | |

# the 3 known good corners. | |

if len(filtered) == 4: | |

for i in range(4): | |

point_info = (filtered[i].corner_location, | |

filtered[i].line1, | |

filtered[i].line2) | |

if (self._lost_corners[i] and | |

not self._PointConnectsToCorners(filtered_corners, point_info)): | |

filtered_corners = np.delete(filtered_corners, i, 0) | |

corner_indices = np.delete(corner_indices, i, 0) | |

break | |

# Ensure corners are sorted properly, inserting nans for missing corners. | |

sorted_corners = np.empty((4, 2), np.float32) | |

sorted_corners[:] = np.nan | |

for i, filtered_corner in enumerate(filtered_corners): | |

sorted_corners[corner_indices[i]] = filtered_corner | |

# From this point on, our corners arrays are guaranteed to have 4 | |

# elements, though some may be nan. | |

# Filter corners that have moved too far from the previous corner if we | |

# are not resetting known corner information. | |

reset_corners = ( | |

(self._lost_corner_frames > self.RESET_AFTER_N_BAD_FRAMES) | |

and len(filtered_corners) == 4) | |

if self._prev_corners is not None and not reset_corners: | |

sqdists = cv_util.SqDistances(self._prev_corners, sorted_corners) | |

for i in range(4): | |

if np.isnan(sorted_corners[i][0]): | |

continue | |

if sqdists[i] > self.MAX_INTERFRAME_MOTION: | |

sorted_corners[i] = np.nan | |

real_corners = self._FindExactCorners(sorted_corners) | |

missing_corners = np.count_nonzero(np.isnan(real_corners)) / 2 | |

return real_corners, missing_corners | |

def _LooksLikeCorner(self, intersections, grey_frame): | |

"""Finds any intersections of lines that look like a screen corner. | |

Args: | |

intersections: The numpy array of points, and the lines that intersect | |

at the given point. | |

grey_frame: The greyscale frame we're processing. | |

Returns: | |

An array of: The corner location (0-3), the relative brightness score | |

(to be used to de-duplicate corners later), the point, and the lines | |

that make up the intersection, for all intersections that look like a | |

corner.""" | |

points = np.vstack(intersections[:, 0].flat) | |

lines1 = np.vstack(intersections[:, 1].flat) | |

lines2 = np.vstack(intersections[:, 2].flat) | |

# Map the image to four quadrants defined as the regions between each of | |

# the lines that make up the intersection. | |

line1a1 = np.pi - np.arctan2(lines1[:, 1] - points[:, 1], | |

lines1[:, 0] - points[:, 0]) | |

line1a2 = np.pi - np.arctan2(lines1[:, 3] - points[:, 1], | |

lines1[:, 2] - points[:, 0]) | |

line2a1 = np.pi - np.arctan2(lines2[:, 1] - points[:, 1], | |

lines2[:, 0] - points[:, 0]) | |

line2a2 = np.pi - np.arctan2(lines2[:, 3] - points[:, 1], | |

lines2[:, 2] - points[:, 0]) | |

line1a1 = line1a1.reshape(-1, 1) | |

line1a2 = line1a2.reshape(-1, 1) | |

line2a1 = line2a1.reshape(-1, 1) | |

line2a2 = line2a2.reshape(-1, 1) | |

line_angles = np.concatenate((line1a1, line1a2, line2a1, line2a2), axis=1) | |

np.ndarray.sort(line_angles) | |

# TODO(mthiesse): Investigate whether these should scale with image or | |

# screen size. My intuition is that these don't scale with image size, | |

# though they may be affected by image quality and how blurry the corners | |

# are. See stackoverflow.com/q/7765810/ for inspiration. | |

avg_range = 8.0 | |

num_points = 7 | |

points_m_avg = points - avg_range | |

points_p_avg = points + avg_range | |

# Exclude points near frame boundaries. | |

include = np.where((points_m_avg[:, 0] > 0) & (points_m_avg[:, 1] > 0) & | |

(points_p_avg[:, 0] < self._width) & | |

(points_p_avg[:, 1] < self._height)) | |

line_angles = line_angles[include] | |

points = points[include] | |

lines1 = lines1[include] | |

lines2 = lines2[include] | |

points_m_avg = points_m_avg[include] | |

points_p_avg = points_p_avg[include] | |

# Perform a 2-d linspace to generate the x, y ranges for each | |

# intersection. | |

arr1 = points_m_avg[:, 0].reshape(-1, 1) | |

arr2 = points_p_avg[:, 0].reshape(-1, 1) | |

lin = np.linspace(0, 1, num_points) | |

x_range = arr1 + (arr2 - arr1) * lin | |

arr1 = points_m_avg[:, 1].reshape(-1, 1) | |

arr2 = points_p_avg[:, 1].reshape(-1, 1) | |

y_range = arr1 + (arr2 - arr1) * lin | |

# The angles for each point we look at in the grid when computing | |

# brightness are constant across frames, so we can generate them once. | |

if self._anglesp5 is None: | |

ind = np.transpose([np.tile(x_range[0], num_points), | |

np.repeat(y_range[0], num_points)]) | |

vectors = ind - points[0] | |

angles = np.arctan2(vectors[:, 1], vectors[:, 0]) + np.pi | |

self._anglesp5 = angles + self.SMALL_ANGLE | |

self._anglesm5 = angles - self.SMALL_ANGLE | |

results = [] | |

for i, _ in enumerate(y_range): | |

# Generate our filters for which points belong to which quadrant. | |

one = np.where((self._anglesp5 <= line_angles[i, 1]) & | |

(self._anglesm5 >= line_angles[i, 0])) | |

two = np.where((self._anglesp5 <= line_angles[i, 2]) & | |

(self._anglesm5 >= line_angles[i, 1])) | |

thr = np.where((self._anglesp5 <= line_angles[i, 3]) & | |

(self._anglesm5 >= line_angles[i, 2])) | |

fou = np.where((self._anglesp5 <= line_angles[i, 0]) | | |

(self._anglesm5 >= line_angles[i, 3])) | |

# Take the cartesian product of our x and y ranges to get the full list | |

# of pixels to look at. | |

ind = np.transpose([np.tile(x_range[i], num_points), | |

np.repeat(y_range[i], num_points)]) | |

# Filter the full list by which indices belong to which quadrant, and | |

# convert to integers so we can index with them. | |

one_i = np.int32(np.rint(ind[one[0]])) | |

two_i = np.int32(np.rint(ind[two[0]])) | |

thr_i = np.int32(np.rint(ind[thr[0]])) | |

fou_i = np.int32(np.rint(ind[fou[0]])) | |

# Average the brightness of the pixels that belong to each quadrant. | |

q_1 = np.average(grey_frame[one_i[:, 1], one_i[:, 0]]) | |

q_2 = np.average(grey_frame[two_i[:, 1], two_i[:, 0]]) | |

q_3 = np.average(grey_frame[thr_i[:, 1], thr_i[:, 0]]) | |

q_4 = np.average(grey_frame[fou_i[:, 1], fou_i[:, 0]]) | |

avg_intensity = [(q_4, 0), (q_1, 1), (q_2, 2), (q_3, 3)] | |

# Sort by intensity. | |

avg_intensity.sort(reverse=True) | |

# Treat the point as a corner if one quadrant is at least twice as | |

# bright as the next brightest quadrant, with a minimum brightness | |

# requirement. | |

tau = (2.0 * np.pi) | |

min_factor = self.MIN_RELATIVE_BRIGHTNESS_FACTOR | |

min_brightness = self.MIN_RELATIVE_BRIGHTNESS_FACTOR | |

if avg_intensity[0][0] > avg_intensity[1][0] * min_factor and \ | |

avg_intensity[0][0] > min_brightness: | |

bright_corner = avg_intensity[0][1] | |

if bright_corner == 0: | |

angle = np.pi - (line_angles[i, 0] + line_angles[i, 3]) / 2.0 | |

if angle < 0: | |

angle = angle + tau | |

else: | |

angle = tau - (line_angles[i, bright_corner] + | |

line_angles[i, bright_corner - 1]) / 2.0 | |

score = avg_intensity[0][0] - avg_intensity[1][0] | |

# TODO(mthiesse): int(angle / (pi / 2.0)) will break if the screen is | |

# rotated through 45 degrees. Probably many other things will break as | |

# well, movement of corners from one quadrant to another hasn't been | |

# tested. We should support this eventually, but this is unlikely to | |

# cause issues for any test setups. | |

results.append((int(angle / (np.pi / 2.0)), score, points[i], | |

lines1[i], lines2[i])) | |

return results | |

def _DeDupCorners(self, corner_data, corners): | |

"""De-duplicate corners based on corner_index. | |

For each set of points representing a corner: If one point is part of the | |

rectangle and the other is not, filter the other one. If both or none are | |

part of the rectangle, filter based on score (highest relative brightness | |

of a quadrant). The reason we allow for neither to be part of the | |

rectangle is because we may not have found all four corners of the | |

rectangle, and in degenerate cases like this it's better to find 3 likely | |

corners than none. | |

Modifies corner_data directly. | |

Args: | |

corner_data: CornerData for each potential corner in the frame. | |

corners: List of all potential corners in the frame.""" | |

# TODO(mthiesse): Ensure that the corners form a sensible rectangle. For | |

# example, it is currently possible (but unlikely) to detect a 'screen' | |

# where the bottom-left corner is above the top-left corner, while the | |

# bottom-right corner is below the top-right corner. | |

# Sort by corner_index to make de-duping easier. | |

corner_data.sort() | |

# De-dup corners. | |

c_old = None | |

for i in range(len(corner_data) - 1, 0, -1): | |

if corner_data[i].corner_index != corner_data[i - 1].corner_index: | |

c_old = None | |

continue | |

if c_old is None: | |

point_info = (corner_data[i].corner_location, | |

corner_data[i].line1, | |

corner_data[i].line2) | |

c_old = self._PointConnectsToCorners(corners, point_info, 2) | |

point_info_new = (corner_data[i - 1].corner_location, | |

corner_data[i - 1].line1, | |

corner_data[i - 1].line2) | |

c_new = self._PointConnectsToCorners(corners, point_info_new, 2) | |

if (not (c_old or c_new)) or (c_old and c_new): | |

if (corner_data[i].brightness_score < | |

corner_data[i - 1].brightness_score): | |

del corner_data[i] | |

c_old = c_new | |

else: | |

del corner_data[i - 1] | |

elif c_old: | |

del corner_data[i - 1] | |

else: | |

del corner_data[i] | |

c_old = c_new | |

def _PointConnectsToCorners(self, corners, point_info, tolerance=1): | |

"""Checks if the lines of an intersection intersect with corners. | |

This is useful to check if the point is part of a rectangle specified by | |

|corners|. | |

Args: | |

point_info: A tuple of (point, line, line) representing an intersection | |

of two lines. | |

corners: corners that (hopefully) make up a rectangle. | |

tolerance: The tolerance (approximately in pixels) of the distance | |

between the corners and the lines for detecting if the point is on | |

the line. | |

Returns: | |

True if each of the two lines that make up the intersection where the | |

point is located connect the point to other corners.""" | |

line1_connected = False | |

line2_connected = False | |

point, line1, line2 = point_info | |

for corner in corners: | |

if corner is None: | |

continue | |

# Filter out points that are too close to one another to be different | |

# corners. | |

sqdist = cv_util.SqDistance(corner, point) | |

if sqdist < self.MIN_SCREEN_WIDTH * self.MIN_SCREEN_WIDTH: | |

continue | |

line1_connected = line1_connected or \ | |

cv_util.IsPointApproxOnLine(corner, line1, tolerance) | |

line2_connected = line2_connected or \ | |

cv_util.IsPointApproxOnLine(corner, line2, tolerance) | |

if line1_connected and line2_connected: | |

return True | |

return False | |

def _FindExactCorners(self, sorted_corners): | |

"""Attempts to find more accurate corner locations. | |

Args: | |

sorted_corners: The four screen corners, sorted by corner_index. | |

Returns: | |

A list of 4 probably more accurate corners, still sorted.""" | |

real_corners = np.empty((4, 2), np.float32) | |

# Count missing corners, and search in a small area around our | |

# intersections representing corners to see if we can find a more exact | |

# corner, as the position of the intersections is noisy and not always | |

# perfectly accurate. | |

for i in range(4): | |

corner = sorted_corners[i] | |

if np.isnan(corner[0]): | |

real_corners[i] = np.nan | |

continue | |

# Almost unbelievably, in edge cases with floating point error, the | |

# width/height of the cropped corner image may be 2 or 4. This is fine | |

# though, as long as the width and height of the cropped corner are not | |

# hard-coded anywhere. | |

corner_image = self._frame_edges[corner[1] - 1:corner[1] + 2, | |

corner[0] - 1:corner[0] + 2] | |

ret, p = self._FindExactCorner(i <= 1, i == 1 or i == 2, corner_image) | |

if ret: | |

if self.DEBUG: | |

self._frame_edges[corner[1] - 1 + p[1]][corner[0] - 1 + p[0]] = 128 | |

real_corners[i] = corner - 1 + p | |

else: | |

real_corners[i] = corner | |

return real_corners | |

def _FindExactCorner(self, top, left, img): | |

"""Tries to finds the exact corner location for a given corner. | |

Searches for the top or bottom, left or right most lit | |

pixel in an edge-detected image, which should represent, with pixel | |

precision, as accurate a corner location as possible. (Though perhaps | |

up-sampling using cubic spline interpolation could get sub-pixel | |

precision) | |

TODO(mthiesse): This algorithm could be improved by including a larger | |

region to search in, but would have to be made smarter about which lit | |

pixels are on the detected screen edge and which are a not as it's | |

currently extremely easy to fool by things like notification icons in | |

screen corners. | |

Args: | |

top: boolean, whether or not we're looking for a top corner. | |

left: boolean, whether or not we're looking for a left corner. | |

img: A small cropping of the edge detected image in which to search. | |

Returns: | |

True and the location if a better corner location is found, | |

False otherwise.""" | |

h, w = img.shape[:2] | |

cy = 0 | |

starting_x = w - 1 if left else 0 | |

cx = starting_x | |

if top: | |

y_range = range(h - 1, -1, -1) | |

else: | |

y_range = range(0, h, 1) | |

if left: | |

x_range = range(w - 1, -1, -1) | |

else: | |

x_range = range(0, w, 1) | |

for y in y_range: | |

for x in x_range: | |

if img[y][x] == 255: | |

cy = y | |

if (left and x <= cx) or (not left and x >= cx): | |

cx = x | |

if cx == starting_x and cy == 0 and img[0][starting_x] != 255: | |

return False, (0, 0) | |

return True, (cx, cy) | |

def _NewScreenLocation(self, new_corners, missing_corners, intersections): | |

"""Computes the new screen location with best effort. | |

Creates the final list of corners that represents the best effort attempt | |

to find the new screen location. Handles degenerate cases where 3 or fewer | |

new corners are present, using previous corner and intersection data. | |

Args: | |

new_corners: The corners found by our search for corners. | |

missing_corners: The count of how many corners we're missing. | |

intersections: The intersections of straight lines found in the current | |

frame. | |

Returns: | |

An array of 4 new_corners hopefully representing the screen, or throws | |

an error if this is not possible. | |

Raises: | |

ValueError: Finding the screen location was not possible.""" | |

screen_corners = copy.copy(new_corners) | |

if missing_corners == 0: | |

self._lost_corner_frames = 0 | |

self._lost_corners = [False, False, False, False] | |

return screen_corners | |

if self._prev_corners is None: | |

raise self.ScreenNotFoundError( | |

'Could not locate screen on frame %d' % | |

self._frame_generator.CurrentFrameNumber) | |

self._lost_corner_frames += 1 | |

if missing_corners > 1: | |

logging.info('Unable to properly detect screen corners, making ' | |

'potentially false assumptions on frame %d', | |

self._frame_generator.CurrentFrameNumber) | |

# Replace missing new_corners with either nearest intersection to previous | |

# corner, or previous corner if no intersections are found. | |

for i in range(0, 4): | |

if not np.isnan(new_corners[i][0]): | |

self._lost_corners[i] = False | |

continue | |

self._lost_corners[i] = True | |

min_dist = self.MAX_INTERFRAME_MOTION | |

min_corner = None | |

for isection in intersections: | |

dist = cv_util.SqDistance(isection[0], self._prev_corners[i]) | |

if dist >= min_dist: | |

continue | |

if missing_corners == 1: | |

# We know in this case that we have 3 corners present, meaning | |

# all 4 screen lines, and therefore intersections near screen | |

# corners present, so our new corner must connect to these | |

# other corners. | |

if not self._PointConnectsToCorners(new_corners, isection, 3): | |

continue | |

min_corner = isection[0] | |

min_dist = dist | |

screen_corners[i] = min_corner if min_corner is not None else \ | |

self._prev_corners[i] | |

return screen_corners | |

def _SmoothCorners(self, corners): | |

"""Smoothes the motion of corners, reduces noise. | |

Smoothes the motion of corners by computing an exponentially weighted | |

moving average of corner positions over time. | |

Args: | |

corners: The corners of the detected screen. | |

Returns: | |

The final corner positions.""" | |

if self._avg_corners is None: | |

self._avg_corners = np.asfarray(corners, np.float32) | |

for i in range(0, 4): | |

# Keep an exponential moving average of the corner location to reduce | |

# noise. | |

new_contrib = np.multiply(self.CORNER_AVERAGE_WEIGHT, corners[i]) | |

old_contrib = np.multiply(1 - self.CORNER_AVERAGE_WEIGHT, | |

self._avg_corners[i]) | |

self._avg_corners[i] = np.add(new_contrib, old_contrib) | |

return self._avg_corners | |

def _GetTransform(self, corners, border): | |

"""Gets the perspective transform of the screen. | |

Args: | |

corners: The corners of the detected screen. | |

border: The number of pixels of border to crop along with the screen. | |

Returns: | |

A perspective transform and the width and height of the target | |

transform. | |

Raises: | |

ScreenNotFoundError: Something went wrong in detecting the screen.""" | |

if self._screen_size is None: | |

w = np.sqrt(cv_util.SqDistance(corners[1], corners[0])) | |

h = np.sqrt(cv_util.SqDistance(corners[1], corners[2])) | |

if w < 1 or h < 1: | |

raise self.ScreenNotFoundError( | |

'Screen detected improperly (bad corners)') | |

if min(w, h) < self.MIN_SCREEN_WIDTH: | |

raise self.ScreenNotFoundError('Detected screen was too small.') | |

self._screen_size = (w, h) | |

# Extend min line length, if we can, to reduce the number of extraneous | |

# lines the line finder finds. | |

self._min_line_length = max(self._min_line_length, min(w, h) / 1.75) | |

w = self._screen_size[0] | |

h = self._screen_size[1] | |

target = np.zeros((4, 2), np.float32) | |

width = w + border | |

height = h + border | |

target[0] = np.asfarray((width, border)) | |

target[1] = np.asfarray((border, border)) | |

target[2] = np.asfarray((border, height)) | |

target[3] = np.asfarray((width, height)) | |

transform_w = width + border | |

transform_h = height + border | |

transform = cv2.getPerspectiveTransform(corners, target) | |

return transform, transform_w, transform_h | |

def _Debug(self, lines, corners, final_corners, screen): | |

for line in lines: | |

intline = ((int(line[0]), int(line[1])), | |

(int(line[2]), int(line[3]))) | |

cv2.line(self._frame_debug, intline[0], intline[1], (0, 0, 255), 1) | |

i = 0 | |

for corner in corners: | |

if not np.isnan(corner[0]): | |

cv2.putText( | |

self._frame_debug, str(i), (int(corner[0]), int(corner[1])), | |

cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (255, 255, 0), 1, cv2.CV_AA) | |

i += 1 | |

if final_corners is not None: | |

for corner in final_corners: | |

cv2.circle(self._frame_debug, | |

(int(corner[0]), int(corner[1])), 5, (255, 0, 255), 1) | |

cv2.imshow('original', self._frame) | |

cv2.imshow('debug', self._frame_debug) | |

if screen is not None: | |

cv2.imshow('screen', screen) | |

cv2.waitKey() | |

# For being run as a script. | |

# TODO(mthiesse): To be replaced with a better standalone script. | |

# Ex: ./screen_finder.py path_to_video 0 5 --verbose | |

def main(): | |

start_frame = int(sys.argv[2]) if len(sys.argv) >= 3 else 0 | |

vf = video_file_frame_generator.VideoFileFrameGenerator(sys.argv[1], | |

start_frame) | |

if len(sys.argv) >= 4: | |

sf = ScreenFinder(vf, int(sys.argv[3])) | |

else: | |

sf = ScreenFinder(vf) | |

# TODO(mthiesse): Use argument parser to improve command line parsing. | |

if len(sys.argv) > 4 and sys.argv[4] == '--verbose': | |

logging.basicConfig(format='%(message)s', level=logging.INFO) | |

else: | |

logging.basicConfig(format='%(message)s', level=logging.WARN) | |

while sf.HasNext(): | |

sf.GetNext() | |

if __name__ == '__main__': | |

sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..')) | |

main() |