File size: 2,317 Bytes
66003a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import numpy as np
def randomly_limit_trues(mask: np.ndarray, max_trues: int) -> np.ndarray:
"""
If mask has more than max_trues True values,
randomly keep only max_trues of them and set the rest to False.
"""
# 1D positions of all True entries
true_indices = np.flatnonzero(mask) # shape = (N_true,)
# if already within budget, return as-is
if true_indices.size <= max_trues:
return mask
# randomly pick which True positions to keep
sampled_indices = np.random.choice(true_indices, size=max_trues, replace=False) # shape = (max_trues,)
# build new flat mask: True only at sampled positions
limited_flat_mask = np.zeros(mask.size, dtype=bool)
limited_flat_mask[sampled_indices] = True
# restore original shape
return limited_flat_mask.reshape(mask.shape)
def create_pixel_coordinate_grid(num_frames, height, width):
"""
Creates a grid of pixel coordinates and frame indices for all frames.
Returns:
tuple: A tuple containing:
- points_xyf (numpy.ndarray): Array of shape (num_frames, height, width, 3)
with x, y coordinates and frame indices
- y_coords (numpy.ndarray): Array of y coordinates for all frames
- x_coords (numpy.ndarray): Array of x coordinates for all frames
- f_coords (numpy.ndarray): Array of frame indices for all frames
"""
# Create coordinate grids for a single frame
y_grid, x_grid = np.indices((height, width), dtype=np.float32)
x_grid = x_grid[np.newaxis, :, :]
y_grid = y_grid[np.newaxis, :, :]
# Broadcast to all frames
x_coords = np.broadcast_to(x_grid, (num_frames, height, width))
y_coords = np.broadcast_to(y_grid, (num_frames, height, width))
# Create frame indices and broadcast
f_idx = np.arange(num_frames, dtype=np.float32)[:, np.newaxis, np.newaxis]
f_coords = np.broadcast_to(f_idx, (num_frames, height, width))
# Stack coordinates and frame indices
points_xyf = np.stack((x_coords, y_coords, f_coords), axis=-1)
return points_xyf
|