| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| |
|
| | def iuvmap_clean(U_uv, V_uv, Index_UV, AnnIndex=None): |
| |
|
| | Index_UV_max = torch.argmax(Index_UV, dim=1).float() |
| | recon_Index_UV = [] |
| | for i in range(Index_UV.size(1)): |
| | if i == 0: |
| | recon_Index_UV_i = torch.min( |
| | F.threshold(Index_UV_max + 1, 0.5, 0), -F.threshold(-Index_UV_max - 1, -1.5, 0) |
| | ) |
| | else: |
| | recon_Index_UV_i = torch.min( |
| | F.threshold(Index_UV_max, i - 0.5, 0), -F.threshold(-Index_UV_max, -i - 0.5, 0) |
| | ) / float(i) |
| | recon_Index_UV.append(recon_Index_UV_i) |
| | recon_Index_UV = torch.stack(recon_Index_UV, dim=1) |
| |
|
| | if AnnIndex is None: |
| | recon_Ann_Index = None |
| | else: |
| | AnnIndex_max = torch.argmax(AnnIndex, dim=1).float() |
| | recon_Ann_Index = [] |
| | for i in range(AnnIndex.size(1)): |
| | if i == 0: |
| | recon_Ann_Index_i = torch.min( |
| | F.threshold(AnnIndex_max + 1, 0.5, 0), -F.threshold(-AnnIndex_max - 1, -1.5, 0) |
| | ) |
| | else: |
| | recon_Ann_Index_i = torch.min( |
| | F.threshold(AnnIndex_max, i - 0.5, 0), -F.threshold(-AnnIndex_max, -i - 0.5, 0) |
| | ) / float(i) |
| | recon_Ann_Index.append(recon_Ann_Index_i) |
| | recon_Ann_Index = torch.stack(recon_Ann_Index, dim=1) |
| |
|
| | recon_U = recon_Index_UV * U_uv |
| | recon_V = recon_Index_UV * V_uv |
| |
|
| | return recon_U, recon_V, recon_Index_UV, recon_Ann_Index |
| |
|
| |
|
| | def iuv_map2img(U_uv, V_uv, Index_UV, AnnIndex=None, uv_rois=None, ind_mapping=None, n_part=24): |
| | device_id = U_uv.get_device() |
| | batch_size = U_uv.size(0) |
| | K = U_uv.size(1) |
| | heatmap_size = U_uv.size(2) |
| |
|
| | Index_UV_max = torch.argmax(Index_UV, dim=1) |
| | if AnnIndex is None: |
| | Index_UV_max = Index_UV_max.to(torch.int64) |
| | else: |
| | AnnIndex_max = torch.argmax(AnnIndex, dim=1) |
| | Index_UV_max = Index_UV_max * (AnnIndex_max > 0).to(torch.int64) |
| |
|
| | outputs = [] |
| |
|
| | for batch_id in range(batch_size): |
| |
|
| | output = torch.zeros([3, U_uv.size(2), U_uv.size(3)], dtype=torch.float32).cuda(device_id) |
| | output[0] = Index_UV_max[batch_id].to(torch.float32) |
| | if ind_mapping is None: |
| | output[0] /= float(K - 1) |
| | else: |
| | for ind in range(len(ind_mapping)): |
| | output[0][output[0] == ind] = ind_mapping[ind] * (1. / n_part) |
| |
|
| | for part_id in range(0, K): |
| | CurrentU = U_uv[batch_id, part_id] |
| | CurrentV = V_uv[batch_id, part_id] |
| | output[1, |
| | Index_UV_max[batch_id] == part_id] = CurrentU[Index_UV_max[batch_id] == part_id] |
| | output[2, |
| | Index_UV_max[batch_id] == part_id] = CurrentV[Index_UV_max[batch_id] == part_id] |
| |
|
| | if uv_rois is None: |
| | outputs.append(output.unsqueeze(0)) |
| | else: |
| | roi_fg = uv_rois[batch_id][1:] |
| |
|
| | |
| | |
| | |
| | |
| |
|
| | w = roi_fg[2] - roi_fg[0] |
| | h = roi_fg[3] - roi_fg[1] |
| |
|
| | aspect_ratio = float(w) / h |
| |
|
| | if aspect_ratio < 1: |
| | new_size = [heatmap_size, max(int(heatmap_size * aspect_ratio), 1)] |
| | output = F.interpolate(output.unsqueeze(0), size=new_size, mode='nearest') |
| | paddingleft = int(0.5 * (heatmap_size - new_size[1])) |
| | output = F.pad( |
| | output, pad=(paddingleft, heatmap_size - new_size[1] - paddingleft, 0, 0) |
| | ) |
| | else: |
| | new_size = [max(int(heatmap_size / aspect_ratio), 1), heatmap_size] |
| | output = F.interpolate(output.unsqueeze(0), size=new_size, mode='nearest') |
| | paddingtop = int(0.5 * (heatmap_size - new_size[0])) |
| | output = F.pad( |
| | output, pad=(0, 0, paddingtop, heatmap_size - new_size[0] - paddingtop) |
| | ) |
| |
|
| | outputs.append(output) |
| |
|
| | return torch.cat(outputs, dim=0) |
| |
|
| |
|
| | def iuv_img2map(uvimages, uv_rois=None, new_size=None, n_part=24): |
| | device_id = uvimages.get_device() |
| | batch_size = uvimages.size(0) |
| | uvimg_size = uvimages.size(-1) |
| |
|
| | Index2mask = [ |
| | [0], [1, 2], [3], [4], [5], [6], [7, 9], [8, 10], [11, 13], [12, 14], [15, 17], [16, 18], |
| | [19, 21], [20, 22], [23, 24] |
| | ] |
| |
|
| | part_ind = torch.round(uvimages[:, 0, :, :] * n_part) |
| | part_u = uvimages[:, 1, :, :] |
| | part_v = uvimages[:, 2, :, :] |
| |
|
| | recon_U = [] |
| | recon_V = [] |
| | recon_Index_UV = [] |
| | recon_Ann_Index = [] |
| |
|
| | for i in range(n_part + 1): |
| | if i == 0: |
| | recon_Index_UV_i = torch.min( |
| | F.threshold(part_ind + 1, 0.5, 0), -F.threshold(-part_ind - 1, -1.5, 0) |
| | ) |
| | else: |
| | recon_Index_UV_i = torch.min( |
| | F.threshold(part_ind, i - 0.5, 0), -F.threshold(-part_ind, -i - 0.5, 0) |
| | ) / float(i) |
| | recon_U_i = recon_Index_UV_i * part_u |
| | recon_V_i = recon_Index_UV_i * part_v |
| |
|
| | recon_Index_UV.append(recon_Index_UV_i) |
| | recon_U.append(recon_U_i) |
| | recon_V.append(recon_V_i) |
| |
|
| | for i in range(len(Index2mask)): |
| | if len(Index2mask[i]) == 1: |
| | recon_Ann_Index_i = recon_Index_UV[Index2mask[i][0]] |
| | elif len(Index2mask[i]) == 2: |
| | p_ind0 = Index2mask[i][0] |
| | p_ind1 = Index2mask[i][1] |
| | |
| | |
| | recon_Ann_Index_i = recon_Index_UV[p_ind0] + recon_Index_UV[p_ind1] |
| |
|
| | recon_Ann_Index.append(recon_Ann_Index_i) |
| |
|
| | recon_U = torch.stack(recon_U, dim=1) |
| | recon_V = torch.stack(recon_V, dim=1) |
| | recon_Index_UV = torch.stack(recon_Index_UV, dim=1) |
| | recon_Ann_Index = torch.stack(recon_Ann_Index, dim=1) |
| |
|
| | if uv_rois is None: |
| | return recon_U, recon_V, recon_Index_UV, recon_Ann_Index |
| |
|
| | recon_U_roi = [] |
| | recon_V_roi = [] |
| | recon_Index_UV_roi = [] |
| | recon_Ann_Index_roi = [] |
| |
|
| | if new_size is None: |
| | M = uvimg_size |
| | else: |
| | M = new_size |
| |
|
| | for i in range(batch_size): |
| | roi_fg = uv_rois[i][1:] |
| |
|
| | |
| | |
| | |
| | |
| |
|
| | w = roi_fg[2] - roi_fg[0] |
| | h = roi_fg[3] - roi_fg[1] |
| |
|
| | aspect_ratio = float(w) / h |
| |
|
| | if aspect_ratio < 1: |
| | w_size = max(int(uvimg_size * aspect_ratio), 1) |
| | w_margin = int((uvimg_size - w_size) / 2) |
| |
|
| | recon_U_roi_i = recon_U[i, :, :, w_margin:w_margin + w_size] |
| | recon_V_roi_i = recon_V[i, :, :, w_margin:w_margin + w_size] |
| | recon_Index_UV_roi_i = recon_Index_UV[i, :, :, w_margin:w_margin + w_size] |
| | recon_Ann_Index_roi_i = recon_Ann_Index[i, :, :, w_margin:w_margin + w_size] |
| | else: |
| | h_size = max(int(uvimg_size / aspect_ratio), 1) |
| | h_margin = int((uvimg_size - h_size) / 2) |
| |
|
| | recon_U_roi_i = recon_U[i, :, h_margin:h_margin + h_size, :] |
| | recon_V_roi_i = recon_V[i, :, h_margin:h_margin + h_size, :] |
| | recon_Index_UV_roi_i = recon_Index_UV[i, :, h_margin:h_margin + h_size, :] |
| | recon_Ann_Index_roi_i = recon_Ann_Index[i, :, h_margin:h_margin + h_size, :] |
| |
|
| | recon_U_roi_i = F.interpolate(recon_U_roi_i.unsqueeze(0), size=(M, M), mode='nearest') |
| | recon_V_roi_i = F.interpolate(recon_V_roi_i.unsqueeze(0), size=(M, M), mode='nearest') |
| | recon_Index_UV_roi_i = F.interpolate( |
| | recon_Index_UV_roi_i.unsqueeze(0), size=(M, M), mode='nearest' |
| | ) |
| | recon_Ann_Index_roi_i = F.interpolate( |
| | recon_Ann_Index_roi_i.unsqueeze(0), size=(M, M), mode='nearest' |
| | ) |
| |
|
| | recon_U_roi.append(recon_U_roi_i) |
| | recon_V_roi.append(recon_V_roi_i) |
| | recon_Index_UV_roi.append(recon_Index_UV_roi_i) |
| | recon_Ann_Index_roi.append(recon_Ann_Index_roi_i) |
| |
|
| | recon_U_roi = torch.cat(recon_U_roi, dim=0) |
| | recon_V_roi = torch.cat(recon_V_roi, dim=0) |
| | recon_Index_UV_roi = torch.cat(recon_Index_UV_roi, dim=0) |
| | recon_Ann_Index_roi = torch.cat(recon_Ann_Index_roi, dim=0) |
| |
|
| | return recon_U_roi, recon_V_roi, recon_Index_UV_roi, recon_Ann_Index_roi |
| |
|
| |
|
| | def seg_img2map(segimages, uv_rois=None, new_size=None, n_part=24): |
| | device_id = segimages.get_device() |
| | batch_size = segimages.size(0) |
| | uvimg_size = segimages.size(-1) |
| |
|
| | part_ind = torch.round(segimages[:, 0, :, :] * n_part) |
| |
|
| | recon_Index_UV = [] |
| |
|
| | for i in range(n_part + 1): |
| | if i == 0: |
| | recon_Index_UV_i = torch.min( |
| | F.threshold(part_ind + 1, 0.5, 0), -F.threshold(-part_ind - 1, -1.5, 0) |
| | ) |
| | else: |
| | recon_Index_UV_i = torch.min( |
| | F.threshold(part_ind, i - 0.5, 0), -F.threshold(-part_ind, -i - 0.5, 0) |
| | ) / float(i) |
| |
|
| | recon_Index_UV.append(recon_Index_UV_i) |
| |
|
| | recon_Index_UV = torch.stack(recon_Index_UV, dim=1) |
| |
|
| | if uv_rois is None: |
| | return None, None, recon_Index_UV, None |
| |
|
| | recon_Index_UV_roi = [] |
| |
|
| | if new_size is None: |
| | M = uvimg_size |
| | else: |
| | M = new_size |
| |
|
| | for i in range(batch_size): |
| | roi_fg = uv_rois[i][1:] |
| |
|
| | |
| | |
| | |
| | |
| |
|
| | w = roi_fg[2] - roi_fg[0] |
| | h = roi_fg[3] - roi_fg[1] |
| |
|
| | aspect_ratio = float(w) / h |
| |
|
| | if aspect_ratio < 1: |
| | w_size = max(int(uvimg_size * aspect_ratio), 1) |
| | w_margin = int((uvimg_size - w_size) / 2) |
| |
|
| | recon_Index_UV_roi_i = recon_Index_UV[i, :, :, w_margin:w_margin + w_size] |
| | else: |
| | h_size = max(int(uvimg_size / aspect_ratio), 1) |
| | h_margin = int((uvimg_size - h_size) / 2) |
| |
|
| | recon_Index_UV_roi_i = recon_Index_UV[i, :, h_margin:h_margin + h_size, :] |
| |
|
| | recon_Index_UV_roi_i = F.interpolate( |
| | recon_Index_UV_roi_i.unsqueeze(0), size=(M, M), mode='nearest' |
| | ) |
| |
|
| | recon_Index_UV_roi.append(recon_Index_UV_roi_i) |
| |
|
| | recon_Index_UV_roi = torch.cat(recon_Index_UV_roi, dim=0) |
| |
|
| | return None, None, recon_Index_UV_roi, None |
| |
|
