Spaces:
Paused
Paused
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| 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: | |
| # batch : List[Dict[str: Tensor[N, C, H, W], str: Optional[str]]] | |
| # images_batch : Tensor[N, C, H, W] | |
| # image : Tensor[C, H, W] | |
| 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) # pyre-ignore[6] | |
| 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 | |