|
import math |
|
import torch |
|
import numpy as np |
|
from megatron import get_args |
|
|
|
def slidingcrops(img, mask): |
|
|
|
|
|
args = get_args() |
|
assert args.img_h == args.img_w |
|
crop_size = args.img_h |
|
stride = args.seg_stride |
|
ignore_index = args.ignore_index |
|
n, c, h, w = img.shape |
|
assert h >= crop_size |
|
assert w >= crop_size |
|
long_size = max(h, w) |
|
|
|
img_slices, mask_slices, slices_info = [], [], [] |
|
if long_size > crop_size: |
|
assert stride <= crop_size |
|
h_step_num = int(math.ceil((h - crop_size) / float(stride))) + 1 |
|
w_step_num = int(math.ceil((w - crop_size) / float(stride))) + 1 |
|
for yy in range(h_step_num): |
|
for xx in range(w_step_num): |
|
sy, sx = yy * stride, xx * stride |
|
ey, ex = sy + crop_size, sx + crop_size |
|
img_sub = img[:, :, sy: ey, sx: ex] |
|
mask_sub = mask[:, sy: ey, sx: ex] |
|
|
|
|
|
sub_h, sub_w = img_sub.shape[2:] |
|
pad_h = max(crop_size - sub_h, 0) |
|
pad_w = max(crop_size - sub_w, 0) |
|
img_sub = torch.nn.functional.pad(img_sub, pad=(0, pad_w, 0, pad_h), value=ignore_index) |
|
mask_sub = torch.nn.functional.pad(mask_sub, pad=(0, pad_w, 0, pad_h)) |
|
|
|
img_slices.append(img_sub) |
|
mask_slices.append(mask_sub) |
|
slices_info.append([sy, ey, sx, ex, sub_h, sub_w]) |
|
|
|
return torch.cat(img_slices), torch.cat(mask_slices), slices_info, (h, w) |
|
else: |
|
return img, mask, [[0, h, 0, w, h, w]], (h, w) |
|
|
|
|
|
def slidingjoins(preds, probs, labels, slices_info, img_size): |
|
args = get_args() |
|
num_slices = len(slices_info) |
|
|
|
if num_slices == 1: |
|
return preds, labels |
|
|
|
h, w = img_size |
|
split_size = args.micro_batch_size |
|
|
|
preds_split = torch.split(preds, split_size) |
|
probs_split = torch.split(probs, split_size) |
|
labels_split = torch.split(labels, split_size) |
|
|
|
assert(len(preds_split) == num_slices) |
|
|
|
total_max_probs = torch.zeros((split_size, h, w), dtype=torch.float, device='cuda') |
|
total_preds = torch.zeros((split_size, h, w), dtype=torch.int, device='cuda') |
|
total_labels = torch.zeros((split_size, h, w), dtype=torch.int, device='cuda') |
|
|
|
for i in range(num_slices): |
|
sy, ey, sx, ex, sub_h, sub_w = slices_info[i] |
|
assert sy + sub_h <= h |
|
assert sx + sub_w <= w |
|
curr_max_probs = total_max_probs[:, sy:sy + sub_h, sx:sx + sub_w] |
|
curr_preds = total_preds[:, sy:sy + sub_h, sx:sx + sub_w] |
|
|
|
local_max_probs = probs_split[i][:, :sub_h, : sub_w] |
|
local_preds = preds_split[i][:, :sub_h, :sub_w] |
|
|
|
result_max_probs = torch.maximum(curr_max_probs, local_max_probs) |
|
result_preds = torch.where(curr_max_probs >= local_max_probs, curr_preds, local_preds) |
|
|
|
total_max_probs[:, sy:sy + sub_h, sx:sx + sub_w] = result_max_probs |
|
total_preds[:, sy:sy + sub_h, sx:sx + sub_w] = result_preds |
|
total_labels[:, sy:sy + sub_h, sx:sx + sub_w] = labels_split[i][0, :sub_h, :sub_w] |
|
|
|
return total_preds, total_labels |
|
|
|
|
