blob: edab5e565c48f86b5ac0fba2080bc40db0859c28 [file] [log] [blame]
#!/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()