Spaces:
Build error
Build error
| import torch | |
| # transpose | |
| FLIP_LEFT_RIGHT = 0 | |
| FLIP_TOP_BOTTOM = 1 | |
| class KES(object): | |
| def __init__(self, kes, size, mode=None): | |
| # FIXME remove check once we have better integration with device | |
| # in my version this would consistently return a CPU tensor | |
| device = kes.device if isinstance(kes, torch.Tensor) else torch.device('cpu') | |
| kes = torch.as_tensor(kes, dtype=torch.float32, device=device) | |
| if len(kes.size()) == 2: | |
| kes = kes.unsqueeze(2) | |
| if not kes.size()[0] ==0: | |
| assert(kes.size()[-2] == 12), str(kes.size()) # 12kes | |
| num_kes = kes.shape[0] | |
| kes_x = kes[:, :6, 0] # 4+2=6 | |
| kes_y = kes[:, 6:, 0] | |
| # TODO remove once support or zero in dim is in | |
| if not kes.size()[0] ==0: | |
| assert(kes_x.size() == kes_y.size()), str(kes_x.size())+' '+str(kes_y.size()) | |
| if num_kes > 0: | |
| kes = kes.view(num_kes, -1, 1) | |
| kes_x = kes_x.view(num_kes, -1, 1) | |
| kes_y = kes_y.view(num_kes, -1, 1) | |
| # TODO should I split them? | |
| self.kes = kes | |
| self.kes_x = kes_x | |
| self.kes_y = kes_y | |
| self.size = size | |
| self.mode = mode | |
| def crop(self, box): | |
| w, h = box[2] - box[0], box[3] - box[1] | |
| k = self.kes.clone() | |
| k[:, :6, 0] -= box[0] | |
| k[:, 6:, 0] -= box[1] | |
| return type(self)(k, (w, h), self.mode) | |
| def resize(self, size, *args, **kwargs): | |
| ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(size, self.size)) | |
| ratio_w, ratio_h = ratios | |
| resized_data_x = self.kes_x.clone() | |
| resized_data_x[..., :] *= ratio_w | |
| resized_data_y = self.kes_y.clone() | |
| resized_data_y[..., :] *= ratio_h | |
| resized_data = torch.cat((resized_data_x, resized_data_y), dim=-2) | |
| return type(self)(resized_data, size, self.mode) | |
| def transpose(self, method): | |
| if method not in (FLIP_LEFT_RIGHT,): | |
| raise NotImplementedError( | |
| "Only FLIP_LEFT_RIGHT implemented") | |
| flip_inds = type(self).FLIP_INDS | |
| flipped_data_x = self.kes_x[:, flip_inds] | |
| width = self.size[0] | |
| TO_REMOVE = 1 | |
| # Flip x coordinates | |
| flipped_data_x[..., :] = width - flipped_data_x[..., :] - TO_REMOVE | |
| flipped_data_y = self.kes_y.clone() | |
| flipped_data = torch.cat((flipped_data_x, flipped_data_y), dim=-2) | |
| return type(self)(flipped_data, self.size, self.mode) | |
| def to(self, *args, **kwargs): | |
| return type(self)(self.kes.to(*args, **kwargs), self.size, self.mode) | |
| def __getitem__(self, item): | |
| return type(self)(self.kes[item], self.size, self.mode) | |
| def __repr__(self): | |
| s = self.__class__.__name__ + '(' | |
| s += 'num_instances_x={}, '.format(len(self.kes_x)) | |
| s += 'num_instances_y={}, '.format(len(self.kes_y)) | |
| s += 'image_width={}, '.format(self.size[0]) | |
| s += 'image_height={})'.format(self.size[1]) | |
| return s | |
| def _create_flip_indices(names, flip_map): | |
| full_flip_map = flip_map.copy() | |
| full_flip_map.update({v: k for k, v in flip_map.items()}) | |
| flipped_names = [i if i not in full_flip_map else full_flip_map[i] for i in names] | |
| flip_indices = [names.index(i) for i in flipped_names] | |
| return torch.tensor(flip_indices) | |
| class textKES(KES): | |
| NAMES = [ # x and y | |
| 'meanx', | |
| 'xmin', | |
| 'x2', | |
| 'x3', | |
| 'xmax', | |
| 'cx' | |
| # 'meany', | |
| # 'ymin', | |
| # 'y2', | |
| # 'y3', | |
| # 'ymax', | |
| # 'cy' | |
| ] | |
| FLIP_MAP = { | |
| 'xmin': 'xmax', | |
| 'x2': 'x3', | |
| } | |
| # TODO this doesn't look great | |
| textKES.FLIP_INDS = _create_flip_indices(textKES.NAMES, textKES.FLIP_MAP) | |
| # TODO make this nicer, this is a direct translation from C2 (but removing the inner loop) | |
| def kes_to_heat_map(kes_x, kes_y, mty, rois, heatmap_size): | |
| if rois.numel() == 0: | |
| return rois.new().long(), rois.new().long() | |
| offset_x = rois[:, 0] | |
| offset_y = rois[:, 1] | |
| scale_x = heatmap_size / (rois[:, 2] - rois[:, 0]) | |
| scale_y = heatmap_size / (rois[:, 3] - rois[:, 1]) | |
| offset_x = offset_x[:, None] | |
| offset_y = offset_y[:, None] | |
| scale_x = scale_x[:, None] | |
| scale_y = scale_y[:, None] | |
| x = kes_x[..., 0] | |
| y = kes_y[..., 0] | |
| x_boundary_inds = x == rois[:, 2][:, None] | |
| y_boundary_inds = y == rois[:, 3][:, None] | |
| x = (x - offset_x) * scale_x | |
| x = x.floor().long() | |
| y = (y - offset_y) * scale_y | |
| y = y.floor().long() | |
| x[x_boundary_inds] = heatmap_size - 1 | |
| y[y_boundary_inds] = heatmap_size - 1 | |
| valid_loc_x = (x >= 0) & (x < heatmap_size) | |
| valid_x = (valid_loc_x).long() | |
| valid_loc_y = (y >= 0) & (y < heatmap_size) | |
| valid_y = (valid_loc_y).long() | |
| valid_mty = ((x >= 0) & (x < heatmap_size)) & ((y >= 0) & (y < heatmap_size)) | |
| valid_mty = valid_mty.sum(dim=1)>0 | |
| valid_mty = (valid_mty).long() | |
| heatmap_x = x | |
| heatmap_y = y | |
| mty = mty | |
| return heatmap_x, heatmap_y, valid_x, valid_y, mty, valid_mty | |