Spaces:
Running
Running
| import torch | |
| import numpy as np | |
| from isegm.model.metrics import _compute_iou | |
| from .brs_losses import BRSMaskLoss | |
| class BaseOptimizer: | |
| def __init__(self, optimizer_params, | |
| prob_thresh=0.49, | |
| reg_weight=1e-3, | |
| min_iou_diff=0.01, | |
| brs_loss=BRSMaskLoss(), | |
| with_flip=False, | |
| flip_average=False, | |
| **kwargs): | |
| self.brs_loss = brs_loss | |
| self.optimizer_params = optimizer_params | |
| self.prob_thresh = prob_thresh | |
| self.reg_weight = reg_weight | |
| self.min_iou_diff = min_iou_diff | |
| self.with_flip = with_flip | |
| self.flip_average = flip_average | |
| self.best_prediction = None | |
| self._get_prediction_logits = None | |
| self._opt_shape = None | |
| self._best_loss = None | |
| self._click_masks = None | |
| self._last_mask = None | |
| self.device = None | |
| def init_click(self, get_prediction_logits, pos_mask, neg_mask, device, shape=None): | |
| self.best_prediction = None | |
| self._get_prediction_logits = get_prediction_logits | |
| self._click_masks = (pos_mask, neg_mask) | |
| self._opt_shape = shape | |
| self._last_mask = None | |
| self.device = device | |
| def __call__(self, x): | |
| opt_params = torch.from_numpy(x).float().to(self.device) | |
| opt_params.requires_grad_(True) | |
| with torch.enable_grad(): | |
| opt_vars, reg_loss = self.unpack_opt_params(opt_params) | |
| result_before_sigmoid = self._get_prediction_logits(*opt_vars) | |
| result = torch.sigmoid(result_before_sigmoid) | |
| pos_mask, neg_mask = self._click_masks | |
| if self.with_flip and self.flip_average: | |
| result, result_flipped = torch.chunk(result, 2, dim=0) | |
| result = 0.5 * (result + torch.flip(result_flipped, dims=[3])) | |
| pos_mask, neg_mask = pos_mask[:result.shape[0]], neg_mask[:result.shape[0]] | |
| loss, f_max_pos, f_max_neg = self.brs_loss(result, pos_mask, neg_mask) | |
| loss = loss + reg_loss | |
| f_val = loss.detach().cpu().numpy() | |
| if self.best_prediction is None or f_val < self._best_loss: | |
| self.best_prediction = result_before_sigmoid.detach() | |
| self._best_loss = f_val | |
| if f_max_pos < (1 - self.prob_thresh) and f_max_neg < self.prob_thresh: | |
| return [f_val, np.zeros_like(x)] | |
| current_mask = result > self.prob_thresh | |
| if self._last_mask is not None and self.min_iou_diff > 0: | |
| diff_iou = _compute_iou(current_mask, self._last_mask) | |
| if len(diff_iou) > 0 and diff_iou.mean() > 1 - self.min_iou_diff: | |
| return [f_val, np.zeros_like(x)] | |
| self._last_mask = current_mask | |
| loss.backward() | |
| f_grad = opt_params.grad.cpu().numpy().ravel().astype(np.float) | |
| return [f_val, f_grad] | |
| def unpack_opt_params(self, opt_params): | |
| raise NotImplementedError | |
| class InputOptimizer(BaseOptimizer): | |
| def unpack_opt_params(self, opt_params): | |
| opt_params = opt_params.view(self._opt_shape) | |
| if self.with_flip: | |
| opt_params_flipped = torch.flip(opt_params, dims=[3]) | |
| opt_params = torch.cat([opt_params, opt_params_flipped], dim=0) | |
| reg_loss = self.reg_weight * torch.sum(opt_params**2) | |
| return (opt_params,), reg_loss | |
| class ScaleBiasOptimizer(BaseOptimizer): | |
| def __init__(self, *args, scale_act=None, reg_bias_weight=10.0, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.scale_act = scale_act | |
| self.reg_bias_weight = reg_bias_weight | |
| def unpack_opt_params(self, opt_params): | |
| scale, bias = torch.chunk(opt_params, 2, dim=0) | |
| reg_loss = self.reg_weight * (torch.sum(scale**2) + self.reg_bias_weight * torch.sum(bias**2)) | |
| if self.scale_act == 'tanh': | |
| scale = torch.tanh(scale) | |
| elif self.scale_act == 'sin': | |
| scale = torch.sin(scale) | |
| return (1 + scale, bias), reg_loss | |