|
|
|
|
|
import random |
|
from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple |
|
import torch |
|
from torch import nn |
|
|
|
SampledData = Any |
|
ModelOutput = Any |
|
|
|
|
|
def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]: |
|
""" |
|
Group elements of an iterable by chunks of size `n`, e.g. |
|
grouper(range(9), 4) -> |
|
(0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None) |
|
""" |
|
it = iter(iterable) |
|
while True: |
|
values = [] |
|
for _ in range(n): |
|
try: |
|
value = next(it) |
|
except StopIteration: |
|
if values: |
|
values.extend([fillvalue] * (n - len(values))) |
|
yield tuple(values) |
|
return |
|
values.append(value) |
|
yield tuple(values) |
|
|
|
|
|
class ScoreBasedFilter: |
|
""" |
|
Filters entries in model output based on their scores |
|
Discards all entries with score less than the specified minimum |
|
""" |
|
|
|
def __init__(self, min_score: float = 0.8): |
|
self.min_score = min_score |
|
|
|
def __call__(self, model_output: ModelOutput) -> ModelOutput: |
|
for model_output_i in model_output: |
|
instances = model_output_i["instances"] |
|
if not instances.has("scores"): |
|
continue |
|
instances_filtered = instances[instances.scores >= self.min_score] |
|
model_output_i["instances"] = instances_filtered |
|
return model_output |
|
|
|
|
|
class InferenceBasedLoader: |
|
""" |
|
Data loader based on results inferred by a model. Consists of: |
|
- a data loader that provides batches of images |
|
- a model that is used to infer the results |
|
- a data sampler that converts inferred results to annotations |
|
""" |
|
|
|
def __init__( |
|
self, |
|
model: nn.Module, |
|
data_loader: Iterable[List[Dict[str, Any]]], |
|
data_sampler: Optional[Callable[[ModelOutput], List[SampledData]]] = None, |
|
data_filter: Optional[Callable[[ModelOutput], ModelOutput]] = None, |
|
shuffle: bool = True, |
|
batch_size: int = 4, |
|
inference_batch_size: int = 4, |
|
drop_last: bool = False, |
|
category_to_class_mapping: Optional[dict] = None, |
|
): |
|
""" |
|
Constructor |
|
|
|
Args: |
|
model (torch.nn.Module): model used to produce data |
|
data_loader (Iterable[List[Dict[str, Any]]]): iterable that provides |
|
dictionaries with "images" and "categories" fields to perform inference on |
|
data_sampler (Callable: ModelOutput -> SampledData): functor |
|
that produces annotation data from inference results; |
|
(optional, default: None) |
|
data_filter (Callable: ModelOutput -> ModelOutput): filter |
|
that selects model outputs for further processing |
|
(optional, default: None) |
|
shuffle (bool): if True, the input images get shuffled |
|
batch_size (int): batch size for the produced annotation data |
|
inference_batch_size (int): batch size for input images |
|
drop_last (bool): if True, drop the last batch if it is undersized |
|
category_to_class_mapping (dict): category to class mapping |
|
""" |
|
self.model = model |
|
self.model.eval() |
|
self.data_loader = data_loader |
|
self.data_sampler = data_sampler |
|
self.data_filter = data_filter |
|
self.shuffle = shuffle |
|
self.batch_size = batch_size |
|
self.inference_batch_size = inference_batch_size |
|
self.drop_last = drop_last |
|
if category_to_class_mapping is not None: |
|
self.category_to_class_mapping = category_to_class_mapping |
|
else: |
|
self.category_to_class_mapping = {} |
|
|
|
def __iter__(self) -> Iterator[List[SampledData]]: |
|
for batch in self.data_loader: |
|
|
|
|
|
|
|
images_and_categories = [ |
|
{"image": image, "category": category} |
|
for element in batch |
|
for image, category in zip(element["images"], element["categories"]) |
|
] |
|
if not images_and_categories: |
|
continue |
|
if self.shuffle: |
|
random.shuffle(images_and_categories) |
|
yield from self._produce_data(images_and_categories) |
|
|
|
def _produce_data( |
|
self, images_and_categories: List[Tuple[torch.Tensor, Optional[str]]] |
|
) -> Iterator[List[SampledData]]: |
|
""" |
|
Produce batches of data from images |
|
|
|
Args: |
|
images_and_categories (List[Tuple[torch.Tensor, Optional[str]]]): |
|
list of images and corresponding categories to process |
|
|
|
Returns: |
|
Iterator over batches of data sampled from model outputs |
|
""" |
|
data_batches: List[SampledData] = [] |
|
category_to_class_mapping = self.category_to_class_mapping |
|
batched_images_and_categories = _grouper(images_and_categories, self.inference_batch_size) |
|
for batch in batched_images_and_categories: |
|
batch = [ |
|
{ |
|
"image": image_and_category["image"].to(self.model.device), |
|
"category": image_and_category["category"], |
|
} |
|
for image_and_category in batch |
|
if image_and_category is not None |
|
] |
|
if not batch: |
|
continue |
|
with torch.no_grad(): |
|
model_output = self.model(batch) |
|
for model_output_i, batch_i in zip(model_output, batch): |
|
assert len(batch_i["image"].shape) == 3 |
|
model_output_i["image"] = batch_i["image"] |
|
instance_class = category_to_class_mapping.get(batch_i["category"], 0) |
|
model_output_i["instances"].dataset_classes = torch.tensor( |
|
[instance_class] * len(model_output_i["instances"]) |
|
) |
|
model_output_filtered = ( |
|
model_output if self.data_filter is None else self.data_filter(model_output) |
|
) |
|
data = ( |
|
model_output_filtered |
|
if self.data_sampler is None |
|
else self.data_sampler(model_output_filtered) |
|
) |
|
for data_i in data: |
|
if len(data_i["instances"]): |
|
data_batches.append(data_i) |
|
if len(data_batches) >= self.batch_size: |
|
yield data_batches[: self.batch_size] |
|
data_batches = data_batches[self.batch_size :] |
|
if not self.drop_last and data_batches: |
|
yield data_batches |
|
|
