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