in source/super-resolution-inf1/container/model/predictor.py [0:0]
def __call__(self, img, scale):
def do(idx_patch):
idx, patch = idx_patch
y_hat = model(patch)
return (idx, y_hat)
img_shape = img.shape
shave = 10
if scale == 2:
model_shape=(144, 144)
else:
model_shape=(72, 72)
input_shape = (img_shape[0]//(model_shape[0]-shave)*(model_shape[0]-shave) + model_shape[0],
img_shape[1]//(model_shape[1]-shave)*(model_shape[1]-shave) + model_shape[1])
pad_top = (input_shape[0] - img_shape[0])//2
pad_bottom = input_shape[0] - img_shape[0] - pad_top
pad_left = (input_shape[1] - img_shape[1])//2
pad_right = input_shape[1] - img_shape[1] - pad_left
img = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_REFLECT)
patches = []
for i in range(0, img_shape[0], model_shape[0]-shave):
for j in range(0, img_shape[1], model_shape[1]-shave):
patch = (slice(i, i+model_shape[0]), slice(j, j+model_shape[1]))
patches.append((
patch[0], patch[1],
np.expand_dims((np.transpose(img[patch[0], patch[1],...], (2, 0, 1)))/255, 0).astype('float32')
))
if scale == 4:
model = self.models[1]
else:
model = self.models[0]
temps = []
for i, patch in enumerate(patches):
temps.append(self.thread_pool.submit(do, (i, torch.from_numpy(patch[2]))))
y = np.zeros((input_shape[0]*scale, input_shape[1]*scale, 3)).astype('uint16')
weight = np.zeros(y.shape).astype('uint16')
for task in as_completed(temps):
i, rlt = task.result()
rlt = self.tensor2img(rlt)
y[slice(patches[i][0].start * scale,
patches[i][0].stop * scale),
slice(patches[i][1].start * scale,
patches[i][1].stop * scale),...] += rlt
weight[slice(patches[i][0].start * scale,
patches[i][0].stop * scale),
slice(patches[i][1].start * scale,
patches[i][1].stop * scale),...] += 1
y = y / weight
y = y[pad_top*scale:y.shape[0]-pad_bottom*scale,
pad_left*scale:y.shape[1]-pad_right*scale].astype('uint8')
return y