Spaces:
Runtime error
Runtime error
| import logging | |
| import math | |
| from typing import Dict | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import tqdm | |
| from torch.utils.data import DataLoader | |
| from saicinpainting.evaluation.utils import move_to_device | |
| LOGGER = logging.getLogger(__name__) | |
| class InpaintingEvaluator(): | |
| def __init__(self, dataset, scores, area_grouping=True, bins=10, batch_size=32, device='cuda', | |
| integral_func=None, integral_title=None, clamp_image_range=None): | |
| """ | |
| :param dataset: torch.utils.data.Dataset which contains images and masks | |
| :param scores: dict {score_name: EvaluatorScore object} | |
| :param area_grouping: in addition to the overall scores, allows to compute score for the groups of samples | |
| which are defined by share of area occluded by mask | |
| :param bins: number of groups, partition is generated by np.linspace(0., 1., bins + 1) | |
| :param batch_size: batch_size for the dataloader | |
| :param device: device to use | |
| """ | |
| self.scores = scores | |
| self.dataset = dataset | |
| self.area_grouping = area_grouping | |
| self.bins = bins | |
| self.device = torch.device(device) | |
| self.dataloader = DataLoader(self.dataset, shuffle=False, batch_size=batch_size) | |
| self.integral_func = integral_func | |
| self.integral_title = integral_title | |
| self.clamp_image_range = clamp_image_range | |
| def _get_bin_edges(self): | |
| bin_edges = np.linspace(0, 1, self.bins + 1) | |
| num_digits = max(0, math.ceil(math.log10(self.bins)) - 1) | |
| interval_names = [] | |
| for idx_bin in range(self.bins): | |
| start_percent, end_percent = round(100 * bin_edges[idx_bin], num_digits), \ | |
| round(100 * bin_edges[idx_bin + 1], num_digits) | |
| start_percent = '{:.{n}f}'.format(start_percent, n=num_digits) | |
| end_percent = '{:.{n}f}'.format(end_percent, n=num_digits) | |
| interval_names.append("{0}-{1}%".format(start_percent, end_percent)) | |
| groups = [] | |
| for batch in self.dataloader: | |
| mask = batch['mask'] | |
| batch_size = mask.shape[0] | |
| area = mask.to(self.device).reshape(batch_size, -1).mean(dim=-1) | |
| bin_indices = np.searchsorted(bin_edges, area.detach().cpu().numpy(), side='right') - 1 | |
| # corner case: when area is equal to 1, bin_indices should return bins - 1, not bins for that element | |
| bin_indices[bin_indices == self.bins] = self.bins - 1 | |
| groups.append(bin_indices) | |
| groups = np.hstack(groups) | |
| return groups, interval_names | |
| def evaluate(self, model=None): | |
| """ | |
| :param model: callable with signature (image_batch, mask_batch); should return inpainted_batch | |
| :return: dict with (score_name, group_type) as keys, where group_type can be either 'overall' or | |
| name of the particular group arranged by area of mask (e.g. '10-20%') | |
| and score statistics for the group as values. | |
| """ | |
| results = dict() | |
| if self.area_grouping: | |
| groups, interval_names = self._get_bin_edges() | |
| else: | |
| groups = None | |
| for score_name, score in tqdm.auto.tqdm(self.scores.items(), desc='scores'): | |
| score.to(self.device) | |
| with torch.no_grad(): | |
| score.reset() | |
| for batch in tqdm.auto.tqdm(self.dataloader, desc=score_name, leave=False): | |
| batch = move_to_device(batch, self.device) | |
| image_batch, mask_batch = batch['image'], batch['mask'] | |
| if self.clamp_image_range is not None: | |
| image_batch = torch.clamp(image_batch, | |
| min=self.clamp_image_range[0], | |
| max=self.clamp_image_range[1]) | |
| if model is None: | |
| assert 'inpainted' in batch, \ | |
| 'Model is None, so we expected precomputed inpainting results at key "inpainted"' | |
| inpainted_batch = batch['inpainted'] | |
| else: | |
| inpainted_batch = model(image_batch, mask_batch) | |
| score(inpainted_batch, image_batch, mask_batch) | |
| total_results, group_results = score.get_value(groups=groups) | |
| results[(score_name, 'total')] = total_results | |
| if groups is not None: | |
| for group_index, group_values in group_results.items(): | |
| group_name = interval_names[group_index] | |
| results[(score_name, group_name)] = group_values | |
| if self.integral_func is not None: | |
| results[(self.integral_title, 'total')] = dict(mean=self.integral_func(results)) | |
| return results | |
| def ssim_fid100_f1(metrics, fid_scale=100): | |
| ssim = metrics[('ssim', 'total')]['mean'] | |
| fid = metrics[('fid', 'total')]['mean'] | |
| fid_rel = max(0, fid_scale - fid) / fid_scale | |
| f1 = 2 * ssim * fid_rel / (ssim + fid_rel + 1e-3) | |
| return f1 | |
| def lpips_fid100_f1(metrics, fid_scale=100): | |
| neg_lpips = 1 - metrics[('lpips', 'total')]['mean'] # invert, so bigger is better | |
| fid = metrics[('fid', 'total')]['mean'] | |
| fid_rel = max(0, fid_scale - fid) / fid_scale | |
| f1 = 2 * neg_lpips * fid_rel / (neg_lpips + fid_rel + 1e-3) | |
| return f1 | |
| class InpaintingEvaluatorOnline(nn.Module): | |
| def __init__(self, scores, bins=10, image_key='image', inpainted_key='inpainted', | |
| integral_func=None, integral_title=None, clamp_image_range=None): | |
| """ | |
| :param scores: dict {score_name: EvaluatorScore object} | |
| :param bins: number of groups, partition is generated by np.linspace(0., 1., bins + 1) | |
| :param device: device to use | |
| """ | |
| super().__init__() | |
| LOGGER.info(f'{type(self)} init called') | |
| self.scores = nn.ModuleDict(scores) | |
| self.image_key = image_key | |
| self.inpainted_key = inpainted_key | |
| self.bins_num = bins | |
| self.bin_edges = np.linspace(0, 1, self.bins_num + 1) | |
| num_digits = max(0, math.ceil(math.log10(self.bins_num)) - 1) | |
| self.interval_names = [] | |
| for idx_bin in range(self.bins_num): | |
| start_percent, end_percent = round(100 * self.bin_edges[idx_bin], num_digits), \ | |
| round(100 * self.bin_edges[idx_bin + 1], num_digits) | |
| start_percent = '{:.{n}f}'.format(start_percent, n=num_digits) | |
| end_percent = '{:.{n}f}'.format(end_percent, n=num_digits) | |
| self.interval_names.append("{0}-{1}%".format(start_percent, end_percent)) | |
| self.groups = [] | |
| self.integral_func = integral_func | |
| self.integral_title = integral_title | |
| self.clamp_image_range = clamp_image_range | |
| LOGGER.info(f'{type(self)} init done') | |
| def _get_bins(self, mask_batch): | |
| batch_size = mask_batch.shape[0] | |
| area = mask_batch.view(batch_size, -1).mean(dim=-1).detach().cpu().numpy() | |
| bin_indices = np.clip(np.searchsorted(self.bin_edges, area) - 1, 0, self.bins_num - 1) | |
| return bin_indices | |
| def forward(self, batch: Dict[str, torch.Tensor]): | |
| """ | |
| Calculate and accumulate metrics for batch. To finalize evaluation and obtain final metrics, call evaluation_end | |
| :param batch: batch dict with mandatory fields mask, image, inpainted (can be overriden by self.inpainted_key) | |
| """ | |
| result = {} | |
| with torch.no_grad(): | |
| image_batch, mask_batch, inpainted_batch = batch[self.image_key], batch['mask'], batch[self.inpainted_key] | |
| if self.clamp_image_range is not None: | |
| image_batch = torch.clamp(image_batch, | |
| min=self.clamp_image_range[0], | |
| max=self.clamp_image_range[1]) | |
| self.groups.extend(self._get_bins(mask_batch)) | |
| for score_name, score in self.scores.items(): | |
| result[score_name] = score(inpainted_batch, image_batch, mask_batch) | |
| return result | |
| def process_batch(self, batch: Dict[str, torch.Tensor]): | |
| return self(batch) | |
| def evaluation_end(self, states=None): | |
| """:return: dict with (score_name, group_type) as keys, where group_type can be either 'overall' or | |
| name of the particular group arranged by area of mask (e.g. '10-20%') | |
| and score statistics for the group as values. | |
| """ | |
| LOGGER.info(f'{type(self)}: evaluation_end called') | |
| self.groups = np.array(self.groups) | |
| results = {} | |
| for score_name, score in self.scores.items(): | |
| LOGGER.info(f'Getting value of {score_name}') | |
| cur_states = [s[score_name] for s in states] if states is not None else None | |
| total_results, group_results = score.get_value(groups=self.groups, states=cur_states) | |
| LOGGER.info(f'Getting value of {score_name} done') | |
| results[(score_name, 'total')] = total_results | |
| for group_index, group_values in group_results.items(): | |
| group_name = self.interval_names[group_index] | |
| results[(score_name, group_name)] = group_values | |
| if self.integral_func is not None: | |
| results[(self.integral_title, 'total')] = dict(mean=self.integral_func(results)) | |
| LOGGER.info(f'{type(self)}: reset scores') | |
| self.groups = [] | |
| for sc in self.scores.values(): | |
| sc.reset() | |
| LOGGER.info(f'{type(self)}: reset scores done') | |
| LOGGER.info(f'{type(self)}: evaluation_end done') | |
| return results | |