|
|
|
|
|
|
|
|
|
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
|
|
|
