Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from invertnet import InvNet | |
| import cv2 | |
| from utils import cut_image, recover_image | |
| class Denoise(nn.Module): | |
| def __init__(self): | |
| super(Denoise, self).__init__() | |
| self.net = InvNet() | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def forward(self, x): | |
| if self.training: | |
| return self.forward_train(x) | |
| else: | |
| return self.forward_test(x) | |
| def forward_train(self, x): | |
| output_forward = self.net(x, reverse=False) | |
| output_lr = output_forward[:, :3, :, :] | |
| gaussian_noise = torch.randn(output_forward[:, 3:, :, :].shape, device=self.device) | |
| x_backward = torch.cat([output_lr, gaussian_noise], dim=1) | |
| y_backward = self.net(x_backward, reverse=True) | |
| output_denoise = y_backward[:, :3, :, :] | |
| return output_lr, output_denoise | |
| def forward_test(self, x): | |
| x = x.cpu().numpy() | |
| xs = [x] | |
| xs.append(x[:, :, :, ::-1].copy()) | |
| xs.append(x[:, :, ::-1, :].copy()) | |
| xs.append(x.transpose((0, 1, 3, 2)).copy()) | |
| output_forward = [self.net(torch.tensor(aug, device=self.device), reverse=False) for aug in xs] | |
| x_backwards = [] | |
| for data in output_forward: | |
| output_lr = data[:, :3, :, :] | |
| gaussian_noise = torch.randn(data[:, 3:, :, :].shape, device=self.device) | |
| x_backward = torch.cat([output_lr, gaussian_noise], dim=1) | |
| x_backwards.append(x_backward) | |
| output_denoise = [self.net(data, reverse=True) for data in x_backwards] | |
| output_denoise[0] = np.array(output_denoise[0].cpu()) | |
| output_denoise[1] = np.array(output_denoise[1].cpu())[:, :, :, ::-1].copy() | |
| output_denoise[2] = np.array(output_denoise[2].cpu())[:, :, ::-1, :].copy() | |
| output_denoise[3] = np.array(output_denoise[3].cpu()).transpose((0, 1, 3, 2)).copy() | |
| output_cat = np.concatenate(output_denoise, axis=0) | |
| output = output_cat.mean(axis=0, keepdims=True) | |
| return output | |
| def inference_invdn(img_dir): | |
| model = Denoise() | |
| model.eval() | |
| model.net.load_state_dict(torch.load("invdn_checkpoint.pth", map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))) | |
| image = cv2.imread(img_dir) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255 | |
| cuts, indexes = cut_image(image, (256, 256)) | |
| h, w = image.shape[:2] | |
| outputs = [] | |
| with torch.no_grad(): | |
| for cut in cuts: | |
| h1, w1 = cut.shape[:2] | |
| cut = cv2.resize(cut, (256, 256), interpolation=cv2.INTER_LINEAR) | |
| cut = cut.transpose(2, 0, 1) | |
| output = model(torch.tensor(cut).unsqueeze(0).float()) | |
| output = output[0].transpose(1, 2, 0) * 255 | |
| output = output.astype(dtype=np.uint8) | |
| output = cv2.resize(output, (w1, h1), interpolation=cv2.INTER_LINEAR) | |
| outputs.append(output) | |
| full = recover_image(outputs, indexes, (h, w), (256, 256)) | |
| return full.astype(np.uint8) | |