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