def test_ensure_int():
np.testing.assert_array_equal(
normalization.ensure_int(tf.constant([0.0, 0.5, 1.0])),
np.array([0, 127, 255]))
np.testing.assert_array_equal(
normalization.ensure_int(tf.constant([0.0, 127.0, 255.0])),
np.array([0, 127, 255]),
)
np.testing.assert_array_equal(
normalization.ensure_int(tf.constant([0, 127, 255])),
np.array([0, 127, 255]))
def flow_shift_instances(
ref_instances: List[InstanceType],
ref_img: np.ndarray,
new_img: np.ndarray,
min_shifted_points: int = 0,
scale: float = 1.0,
window_size: int = 21,
max_levels: int = 3,
) -> List[ShiftedInstance]:
"""Generates instances in a new frame by applying optical flow displacements.
Args:
ref_instances: Reference instances in the previous frame.
ref_img: Previous frame image as a numpy array.
new_img: New frame image as a numpy array.
min_shifted_points: Minimum number of points that must be detected in the
new frame in order to generate a new shifted instance.
scale: Factor to scale the images by when computing optical flow. Decrease
this to increase performance at the cost of finer accuracy. Sometimes
decreasing the image scale can improve performance with fast movements.
window_size: Optical flow window size to consider at each pyramid scale
level.
max_levels: Number of pyramid scale levels to consider. This is different
from the scale parameter, which determines the initial image scaling.
Returns:
A list of ShiftedInstances with the optical flow displacements applied to
the reference instance points. Points that are not found will be represented
as NaNs in the points array for each shifted instance.
Notes:
This function relies on the Lucas-Kanade method for optical flow estimation.
"""
# Convert to uint8 for cv2.calcOpticalFlowPyrLK
ref_img = ensure_int(ref_img)
new_img = ensure_int(new_img)
# Convert tensors to ndarays
if hasattr(ref_img, "numpy"):
ref_img = ref_img.numpy()
if hasattr(new_img, "numpy"):
new_img = new_img.numpy()
# Ensure images are rank 2 in case there is a singleton channel dimension.
if ref_img.ndim > 3:
ref_img = np.squeeze(ref_img)
new_img = np.squeeze(new_img)
# Convert RGB to grayscale.
if ref_img.ndim > 2 and ref_img.shape[-1] == 3:
ref_img = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
new_img = cv2.cvtColor(new_img, cv2.COLOR_BGR2GRAY)
# Input image scaling.
if scale != 1:
ref_img = cv2.resize(ref_img, None, None, scale, scale)
new_img = cv2.resize(new_img, None, None, scale, scale)
# Gather reference points.
ref_pts = [inst.points_array for inst in ref_instances]
# Compute optical flow at all points.
shifted_pts, status, errs = cv2.calcOpticalFlowPyrLK(
ref_img,
new_img,
(np.concatenate(ref_pts, axis=0)).astype("float32") * scale,
None,
winSize=(window_size, window_size),
maxLevel=max_levels,
criteria=(
cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
30,
0.01,
),
)
shifted_pts /= scale
# Split results by instances.
sections = np.cumsum([len(x) for x in ref_pts])[:-1]
shifted_pts = np.split(shifted_pts, sections, axis=0)
status = np.split(status, sections, axis=0)
status_sum = [np.sum(x) for x in status]
errs = np.split(errs, sections, axis=0)
# Create shifted instances.
shifted_instances = []
for ref, pts, found, err in zip(ref_instances, shifted_pts, status,
errs):
if found.sum() > min_shifted_points:
# Exclude points that weren't found by optical flow.
found = found.squeeze().astype(bool)
pts[~found] = np.nan
# Create a shifted instance.
shifted_instances.append(
ShiftedInstance.from_instance(
ref,
new_points_array=pts,
shift_score=-np.mean(err[found])))
return shifted_instances