Module wtracker.dataset.sample_extractor
View Source
import numpy as np
from typing import Collection
from wtracker.dataset.box_calculator import BoxCalculator
from wtracker.utils.bbox_utils import BoxUtils
from wtracker.utils.io_utils import FrameSaver
class SampleExtractor:
"""
A class that extracts samples from frames based on specified parameters.
Each sample is a cropped image around a bounding box which was detected in the frame.
The bounding boxes are calculated using the BoxCalculator class.
This class is used to create image datasets for training object detection models.
Args:
bbox_calculator (BoxCalculator): An instance of the BoxCalculator class.
"""
def __init__(self, bbox_calculator: BoxCalculator):
self._bbox_calculator = bbox_calculator
self._frame_reader = bbox_calculator._frame_reader
def move_bboxes_into_bounds(self, bboxes: np.ndarray, frame_size: tuple[int, int]) -> np.ndarray:
"""
Moves the bounding boxes into the bounds of the frame.
Args:
bboxes (np.ndarray): The bounding boxes to be moved.
frame_size (tuple[int, int]): The size of the frame in the format (w, h).
Returns:
np.ndarray: The updated bounding boxes.
Raises:
ValueError: If exists a bounding box which cannot be moved into the provided bounds without resizing it.
"""
max_w, max_h = frame_size
x, y, w, h = BoxUtils.unpack(bboxes)
x[x < 0] = 0
mask = (x + w) > max_w
x[mask] = max_w - w[mask]
y[y < 0] = 0
mask = (y + h) > max_h
y[mask] = max_h - h[mask]
if np.any(x < 0) or np.any(y < 0):
raise ValueError()
if np.any(x + w > frame_size[0]) or np.any(y + h > frame_size[1]):
raise ValueError()
return BoxUtils.pack(x, y, w, h)
def create_specified_samples(
self,
frame_indices: Collection[int],
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates specified samples based on the given frame indices.
Args:
frame_indices (Collection[int]): The indices of the frames to extract samples from.
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk.
"""
bboxes = self._bbox_calculator.calc_specified_boxes(
frame_indices=frame_indices,
num_workers=num_workers,
chunk_size=chunk_size,
)
x, y, w, h = BoxUtils.unpack(bboxes)
x -= np.random.randint(0, target_size[0] - w + 1)
y -= np.random.randint(0, target_size[1] - h + 1)
w = np.full_like(x, target_size[0])
h = np.full_like(x, target_size[1])
bboxes = BoxUtils.pack(x, y, w, h)
frame_size = tuple(reversed(self._frame_reader.frame_size)) # (h, w) -> (w, h)
bboxes = self.move_bboxes_into_bounds(bboxes, frame_size)
with FrameSaver(self._frame_reader, root_path=save_folder, desc="Saving samples", unit="fr") as saver:
for i, bbox in enumerate(bboxes):
saver.schedule_save(i, bbox, name_format.format(i))
def create_samples(
self,
count: int,
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates random samples based on a specified count.
Args:
count (int): The number of samples to create.
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk sent to each worker.
"""
length = len(self._frame_reader)
count = min(length, count)
frame_indices = np.random.choice(length, size=count, replace=False)
self.create_specified_samples(frame_indices, target_size, save_folder, name_format, num_workers, chunk_size)
def create_all_samples(
self,
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates samples for all frames.
Args:
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk.
"""
frame_indices = range(0, len(self._frame_reader))
self.create_specified_samples(frame_indices, target_size, save_folder, name_format, num_workers, chunk_size)
Classes
SampleExtractor
class SampleExtractor(
bbox_calculator: wtracker.dataset.box_calculator.BoxCalculator
)
A class that extracts samples from frames based on specified parameters.
Each sample is a cropped image around a bounding box which was detected in the frame. The bounding boxes are calculated using the BoxCalculator class. This class is used to create image datasets for training object detection models.
Attributes
| Name | Type | Description | Default |
|---|---|---|---|
| bbox_calculator | BoxCalculator | An instance of the BoxCalculator class. | None |
View Source
class SampleExtractor:
"""
A class that extracts samples from frames based on specified parameters.
Each sample is a cropped image around a bounding box which was detected in the frame.
The bounding boxes are calculated using the BoxCalculator class.
This class is used to create image datasets for training object detection models.
Args:
bbox_calculator (BoxCalculator): An instance of the BoxCalculator class.
"""
def __init__(self, bbox_calculator: BoxCalculator):
self._bbox_calculator = bbox_calculator
self._frame_reader = bbox_calculator._frame_reader
def move_bboxes_into_bounds(self, bboxes: np.ndarray, frame_size: tuple[int, int]) -> np.ndarray:
"""
Moves the bounding boxes into the bounds of the frame.
Args:
bboxes (np.ndarray): The bounding boxes to be moved.
frame_size (tuple[int, int]): The size of the frame in the format (w, h).
Returns:
np.ndarray: The updated bounding boxes.
Raises:
ValueError: If exists a bounding box which cannot be moved into the provided bounds without resizing it.
"""
max_w, max_h = frame_size
x, y, w, h = BoxUtils.unpack(bboxes)
x[x < 0] = 0
mask = (x + w) > max_w
x[mask] = max_w - w[mask]
y[y < 0] = 0
mask = (y + h) > max_h
y[mask] = max_h - h[mask]
if np.any(x < 0) or np.any(y < 0):
raise ValueError()
if np.any(x + w > frame_size[0]) or np.any(y + h > frame_size[1]):
raise ValueError()
return BoxUtils.pack(x, y, w, h)
def create_specified_samples(
self,
frame_indices: Collection[int],
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates specified samples based on the given frame indices.
Args:
frame_indices (Collection[int]): The indices of the frames to extract samples from.
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk.
"""
bboxes = self._bbox_calculator.calc_specified_boxes(
frame_indices=frame_indices,
num_workers=num_workers,
chunk_size=chunk_size,
)
x, y, w, h = BoxUtils.unpack(bboxes)
x -= np.random.randint(0, target_size[0] - w + 1)
y -= np.random.randint(0, target_size[1] - h + 1)
w = np.full_like(x, target_size[0])
h = np.full_like(x, target_size[1])
bboxes = BoxUtils.pack(x, y, w, h)
frame_size = tuple(reversed(self._frame_reader.frame_size)) # (h, w) -> (w, h)
bboxes = self.move_bboxes_into_bounds(bboxes, frame_size)
with FrameSaver(self._frame_reader, root_path=save_folder, desc="Saving samples", unit="fr") as saver:
for i, bbox in enumerate(bboxes):
saver.schedule_save(i, bbox, name_format.format(i))
def create_samples(
self,
count: int,
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates random samples based on a specified count.
Args:
count (int): The number of samples to create.
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk sent to each worker.
"""
length = len(self._frame_reader)
count = min(length, count)
frame_indices = np.random.choice(length, size=count, replace=False)
self.create_specified_samples(frame_indices, target_size, save_folder, name_format, num_workers, chunk_size)
def create_all_samples(
self,
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates samples for all frames.
Args:
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk.
"""
frame_indices = range(0, len(self._frame_reader))
self.create_specified_samples(frame_indices, target_size, save_folder, name_format, num_workers, chunk_size)
Methods
create_all_samples
def create_all_samples(
self,
target_size: tuple[int, int],
save_folder: str,
name_format: str = 'img_{:09d}.png',
num_workers: int = None,
chunk_size: int = 50
)
Creates samples for all frames.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
| target_size | tuple[int, int] | The target size of the samples in the format (w, h). | None |
| save_folder | str | The folder path to save the samples. | None |
| name_format | str | The format of the sample names. | None |
| num_workers | int | The number of workers to use for parallel processing. If None, the number of workers is determined automatically. |
None |
| chunk_size | int | The size of each processing chunk. | None |
View Source
def create_all_samples(
self,
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates samples for all frames.
Args:
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk.
"""
frame_indices = range(0, len(self._frame_reader))
self.create_specified_samples(frame_indices, target_size, save_folder, name_format, num_workers, chunk_size)
create_samples
def create_samples(
self,
count: int,
target_size: tuple[int, int],
save_folder: str,
name_format: str = 'img_{:09d}.png',
num_workers: int = None,
chunk_size: int = 50
)
Creates random samples based on a specified count.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
| count | int | The number of samples to create. | None |
| target_size | tuple[int, int] | The target size of the samples in the format (w, h). | None |
| save_folder | str | The folder path to save the samples. | None |
| name_format | str | The format of the sample names. | None |
| num_workers | int | The number of workers to use for parallel processing. If None, the number of workers is determined automatically. |
None |
| chunk_size | int | The size of each processing chunk sent to each worker. | None |
View Source
def create_samples(
self,
count: int,
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates random samples based on a specified count.
Args:
count (int): The number of samples to create.
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk sent to each worker.
"""
length = len(self._frame_reader)
count = min(length, count)
frame_indices = np.random.choice(length, size=count, replace=False)
self.create_specified_samples(frame_indices, target_size, save_folder, name_format, num_workers, chunk_size)
create_specified_samples
def create_specified_samples(
self,
frame_indices: Collection[int],
target_size: tuple[int, int],
save_folder: str,
name_format: str = 'img_{:09d}.png',
num_workers: int = None,
chunk_size: int = 50
)
Creates specified samples based on the given frame indices.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
| frame_indices | Collection[int] | The indices of the frames to extract samples from. | None |
| target_size | tuple[int, int] | The target size of the samples in the format (w, h). | None |
| save_folder | str | The folder path to save the samples. | None |
| name_format | str | The format of the sample names. | None |
| num_workers | int | The number of workers to use for parallel processing. If None, the number of workers is determined automatically. |
None |
| chunk_size | int | The size of each processing chunk. | None |
View Source
def create_specified_samples(
self,
frame_indices: Collection[int],
target_size: tuple[int, int],
save_folder: str,
name_format: str = "img_{:09d}.png",
num_workers: int = None,
chunk_size: int = 50,
):
"""
Creates specified samples based on the given frame indices.
Args:
frame_indices (Collection[int]): The indices of the frames to extract samples from.
target_size (tuple[int, int]): The target size of the samples in the format (w, h).
save_folder (str): The folder path to save the samples.
name_format (str, optional): The format of the sample names.
num_workers (int, optional): The number of workers to use for parallel processing.
If None, the number of workers is determined automatically.
chunk_size (int, optional): The size of each processing chunk.
"""
bboxes = self._bbox_calculator.calc_specified_boxes(
frame_indices=frame_indices,
num_workers=num_workers,
chunk_size=chunk_size,
)
x, y, w, h = BoxUtils.unpack(bboxes)
x -= np.random.randint(0, target_size[0] - w + 1)
y -= np.random.randint(0, target_size[1] - h + 1)
w = np.full_like(x, target_size[0])
h = np.full_like(x, target_size[1])
bboxes = BoxUtils.pack(x, y, w, h)
frame_size = tuple(reversed(self._frame_reader.frame_size)) # (h, w) -> (w, h)
bboxes = self.move_bboxes_into_bounds(bboxes, frame_size)
with FrameSaver(self._frame_reader, root_path=save_folder, desc="Saving samples", unit="fr") as saver:
for i, bbox in enumerate(bboxes):
saver.schedule_save(i, bbox, name_format.format(i))
move_bboxes_into_bounds
def move_bboxes_into_bounds(
self,
bboxes: numpy.ndarray,
frame_size: tuple[int, int]
) -> numpy.ndarray
Moves the bounding boxes into the bounds of the frame.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
| bboxes | np.ndarray | The bounding boxes to be moved. | None |
| frame_size | tuple[int, int] | The size of the frame in the format (w, h). | None |
Returns:
| Type | Description |
|---|---|
| np.ndarray | The updated bounding boxes. |
Raises:
| Type | Description |
|---|---|
| ValueError | If exists a bounding box which cannot be moved into the provided bounds without resizing it. |
View Source
def move_bboxes_into_bounds(self, bboxes: np.ndarray, frame_size: tuple[int, int]) -> np.ndarray:
"""
Moves the bounding boxes into the bounds of the frame.
Args:
bboxes (np.ndarray): The bounding boxes to be moved.
frame_size (tuple[int, int]): The size of the frame in the format (w, h).
Returns:
np.ndarray: The updated bounding boxes.
Raises:
ValueError: If exists a bounding box which cannot be moved into the provided bounds without resizing it.
"""
max_w, max_h = frame_size
x, y, w, h = BoxUtils.unpack(bboxes)
x[x < 0] = 0
mask = (x + w) > max_w
x[mask] = max_w - w[mask]
y[y < 0] = 0
mask = (y + h) > max_h
y[mask] = max_h - h[mask]
if np.any(x < 0) or np.any(y < 0):
raise ValueError()
if np.any(x + w > frame_size[0]) or np.any(y + h > frame_size[1]):
raise ValueError()
return BoxUtils.pack(x, y, w, h)