refcolorrestore 1.2.0__py3-none-any.whl

refcolorrestore/build_dataset.py

@@ -0,0 +1,61 @@
+from typing import Tuple
+import torch
+import os
+from gaussian_splatting import GaussianModel
+from gaussian_splatting.prepare import prepare_dataset, prepare_gaussians
+from extrinterp import ExtrinsicInterpolator
+from refcolorrestore.dataset import Extrinsic2DualCameraDataset, DualCamera2RestorationDataset, DualCameraDataset
+
+
+def prepare_rendering(
+        sh_degree: int, source: str, device: str, n: int, window_size: int,
+        trainable_camera: bool = False,
+        load_ply: str = None, load_ply_gt: str = None,
+        load_camera: str = None,
+        use_intrinsics: int | dict = 0
+) -> Tuple[DualCameraDataset, GaussianModel, GaussianModel]:
+    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_depth=False)
+    if isinstance(use_intrinsics, int):
+        i = use_intrinsics
+        use_intrinsics = dict(
+            image_height=dataset[i].image_height, image_width=dataset[i].image_width,
+            FoVx=dataset[i].FoVx, FoVy=dataset[i].FoVy)
+    elif not isinstance(use_intrinsics, dict):
+        raise ValueError("Invalid use_intrinsics format")
+    dataset = Extrinsic2DualCameraDataset(ExtrinsicInterpolator(dataset=dataset, n=n, window_size=window_size), **use_intrinsics)
+    gaussians = prepare_gaussians(sh_degree=sh_degree, source=source, device=device, trainable_camera=trainable_camera, load_ply=load_ply)
+    gaussians_gt = prepare_gaussians(sh_degree=sh_degree, source=source, device=device, trainable_camera=trainable_camera, load_ply=load_ply_gt)
+    return dataset, gaussians, gaussians_gt
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("--sh_degree", default=3, type=int)
+    parser.add_argument("-s", "--source", required=True, type=str)
+    parser.add_argument("-d", "--destination", required=True, type=str)
+    parser.add_argument("-i", "--iteration", required=True, type=int)
+    parser.add_argument("--load_camera", default=None, type=str)
+    parser.add_argument("--mode", choices=["base", "camera"], default="base")
+    parser.add_argument("--device", default="cuda", type=str)
+    parser.add_argument("--destination_gt", required=True, type=str)
+    parser.add_argument("--iteration_gt", required=True, type=int)
+    parser.add_argument("--interp_n", required=True, type=int)
+    parser.add_argument("--interp_window_size", type=int, default=3)
+    parser.add_argument("--use_intrinsics", type=str, default="0", help="Use intrinsics for rendering, can be an integer index or a dict with keys: image_height, image_width, FoVx, FoVy")
+    parser.add_argument("--downsample", default=4, type=int)
+    parser.add_argument("--data_dir", required=True, type=str)
+    args = parser.parse_args()
+    load_ply = os.path.join(args.destination, "point_cloud", "iteration_" + str(args.iteration), "point_cloud.ply")
+    load_ply_gt = os.path.join(args.destination_gt, "point_cloud", "iteration_" + str(args.iteration_gt), "point_cloud.ply")
+    with torch.no_grad():
+        cameras, gaussians, gaussians_gt = prepare_rendering(
+            sh_degree=args.sh_degree, source=args.source, device=args.device,
+            n=args.interp_n, window_size=args.interp_window_size,
+            trainable_camera=args.mode == "camera",
+            load_ply=load_ply, load_ply_gt=load_ply_gt, load_camera=args.load_camera,
+            use_intrinsics=eval(args.use_intrinsics)
+        )
+        dataset = DualCamera2RestorationDataset(cameras=cameras, color_distorted_gaussians=gaussians, ground_truth_gaussians=gaussians_gt)
+        os.makedirs(args.data_dir, exist_ok=True)
+        dataset.save_dataset(args.data_dir)

refcolorrestore/dataset/__init__.py

@@ -0,0 +1,3 @@
+from .abc import DualCameraDataset, RestorationDataset, RestorationTuple, SavedRestorationDataset
+from .interp import Extrinsic2DualCameraDataset
+from .dataset import DualCamera2RestorationDataset

refcolorrestore/dataset/abc.py

@@ -0,0 +1,82 @@
+import os
+from abc import abstractmethod
+from typing import NamedTuple, Tuple
+from tqdm import tqdm
+
+import torch
+import torchvision
+
+from gaussian_splatting.camera import Camera
+
+
+class DualCameraDataset:
+
+    @abstractmethod
+    def to(self, device) -> 'DualCameraDataset':
+        return self
+
+    @abstractmethod
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def __getitem__(self, idx) -> Tuple[Camera, Camera]:
+        raise NotImplementedError
+
+
+class RestorationTuple(NamedTuple):
+    color_distorted: torch.FloatTensor
+    reference: torch.FloatTensor
+
+
+class RestorationDataset:
+
+    @abstractmethod
+    def to(self, device) -> 'RestorationDataset':
+        return self
+
+    @abstractmethod
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def __getitem__(self, idx) -> Tuple[RestorationTuple, torch.FloatTensor]:
+        raise NotImplementedError
+
+    def save_image_tuple(self, idx, image_dir):
+        restoration, gt = self[idx]
+        color_distorted_dir = os.path.join(image_dir, 'distorted')
+        os.makedirs(color_distorted_dir, exist_ok=True)
+        torchvision.utils.save_image(restoration.color_distorted, os.path.join(color_distorted_dir, '{0:05d}'.format(idx) + ".png"))
+        reference_dir = os.path.join(image_dir, 'reference')
+        os.makedirs(reference_dir, exist_ok=True)
+        torchvision.utils.save_image(restoration.reference, os.path.join(reference_dir, '{0:05d}'.format(idx) + ".png"))
+        gt_dir = os.path.join(image_dir, 'groundtruth')
+        os.makedirs(gt_dir, exist_ok=True)
+        torchvision.utils.save_image(gt, os.path.join(gt_dir, '{0:05d}'.format(idx) + ".png"))
+
+    def save_dataset(self, image_dir):
+        os.makedirs(image_dir, exist_ok=True)
+        for idx in tqdm(range(len(self)), desc="Saving images"):
+            self.save_image_tuple(idx, image_dir)
+
+
+class SavedRestorationDataset(RestorationDataset):
+    def __init__(self, data_dir: str):
+        self.color_distorted_dir = os.path.join(data_dir, 'distorted')
+        self.reference_dir = os.path.join(data_dir, 'reference')
+        self.ground_truth_dir = os.path.join(data_dir, 'groundtruth')
+
+    def __len__(self) -> int:
+        n = 0
+        while os.path.exists(os.path.join(self.color_distorted_dir, '{0:05d}'.format(n) + ".png")) and \
+                os.path.exists(os.path.join(self.reference_dir, '{0:05d}'.format(n) + ".png")) and \
+                os.path.exists(os.path.join(self.ground_truth_dir, '{0:05d}'.format(n) + ".png")):
+            n += 1
+        return n
+
+    def __getitem__(self, idx) -> Tuple[RestorationTuple, torch.FloatTensor]:
+        color_distorted = torchvision.io.read_image(os.path.join(self.color_distorted_dir, '{0:05d}'.format(idx) + ".png")) / 255.0
+        reference = torchvision.io.read_image(os.path.join(self.reference_dir, '{0:05d}'.format(idx) + ".png")) / 255.0
+        ground_truth = torchvision.io.read_image(os.path.join(self.ground_truth_dir, '{0:05d}'.format(idx) + ".png")) / 255.0
+        return RestorationTuple(color_distorted, reference), ground_truth

refcolorrestore/dataset/dataset.py

@@ -0,0 +1,30 @@
+from typing import Tuple
+
+import torch
+
+from gaussian_splatting import GaussianModel
+
+from .abc import DualCameraDataset, RestorationDataset, RestorationTuple
+
+
+class DualCamera2RestorationDataset(RestorationDataset):
+    def __init__(self, cameras: DualCameraDataset, color_distorted_gaussians: GaussianModel, ground_truth_gaussians: GaussianModel):
+        self.cameras = cameras
+        self.ground_truth_gaussians = ground_truth_gaussians
+        self.color_distorted_gaussians = color_distorted_gaussians
+
+    def to(self, device) -> 'DualCamera2RestorationDataset':
+        self.cameras = self.cameras.to(device)
+        self.ground_truth_gaussians = self.ground_truth_gaussians.to(device)
+        self.color_distorted_gaussians = self.color_distorted_gaussians.to(device)
+        return self
+
+    def __len__(self) -> int:
+        return len(self.cameras)
+
+    def __getitem__(self, idx) -> Tuple[RestorationTuple, torch.FloatTensor]:
+        lr_camera, hr_camera = self.cameras[idx]
+        color_distorted = self.color_distorted_gaussians(hr_camera)['render']
+        ground_truth = self.ground_truth_gaussians(hr_camera)['render']
+        reference = self.ground_truth_gaussians(lr_camera)['render']
+        return RestorationTuple(color_distorted, reference), ground_truth

refcolorrestore/dataset/interp.py

@@ -0,0 +1,46 @@
+from typing import Tuple
+import torch
+
+from gaussian_splatting import Camera
+from extrinterp import ExtrinsicDataset
+
+from .abc import DualCameraDataset
+
+
+class Extrinsic2DualCameraDataset(DualCameraDataset):
+    def __init__(
+        self,
+            base_dataset: ExtrinsicDataset,
+            FoVx: float = 90.0*torch.pi/180, FoVy: float = 90.0*torch.pi/180,
+            image_height: int = 1000, image_width: int = 1000, downsample=4):
+        self.cameras = base_dataset
+        self.FoVx = FoVx
+        self.FoVy = FoVy
+        self.downsample = downsample
+        self.image_height = image_height // 4 * 4
+        self.image_width = image_width // 4 * 4
+
+    def __len__(self):
+        return len(self.cameras)
+
+    def __getitem__(self, idx) -> Tuple[Camera, Camera]:
+        extr = self.cameras[idx]
+        camera_hr = extr.to_camera(
+            image_height=self.image_height,
+            image_width=self.image_width,
+            FoVx=self.FoVx,
+            FoVy=self.FoVy,
+            device=self.cameras[idx].R.device
+        )
+        camera_lr = extr.to_camera(
+            image_height=self.image_height // self.downsample,
+            image_width=self.image_width // self.downsample,
+            FoVx=self.FoVx,
+            FoVy=self.FoVy,
+            device=self.cameras[idx].R.device
+        )
+        return camera_lr, camera_hr
+
+    def to(self, device):
+        self.cameras = self.cameras.to(device)
+        return self

refcolorrestore/prepare.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from refcolorrestore.dataset import SavedRestorationDataset
+from refcolorrestore.shelf import model_dict
+
+
+def load_dataset(
+        image_dir: str,
+        batch_size: int, num_workers: int, shuffle=False) -> DataLoader:
+    dataset = SavedRestorationDataset(data_dir=image_dir)
+    loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        shuffle=shuffle,
+        drop_last=False
+    )
+    return loader
+
+
+def build_model(model: str, load_path: str = None, **kwargs) -> nn.Module:
+    net = model_dict[model](**kwargs)
+    if load_path:
+        net.load_state_dict(torch.load(load_path))
+    return net

refcolorrestore/restore.py

@@ -0,0 +1,56 @@
+import shutil
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import os
+from tqdm import tqdm
+import torchvision
+from gaussian_splatting.utils import psnr, ssim
+from gaussian_splatting.utils.lpipsPyTorch import lpips
+
+from refcolorrestore.shelf import model_dict
+from refcolorrestore.prepare import load_dataset, build_model
+
+
+def rendering(net: nn.Module, dataset: DataLoader, render_path: str, log_path: str, device: str, save_image: bool):
+    net = net.eval().to(device)
+    shutil.rmtree(render_path, ignore_errors=True)
+    os.makedirs(render_path, exist_ok=True)
+    pbar = tqdm(dataset, desc="Rendering progress")
+    with open(log_path, "w") as f:
+        f.write(f"frame,psnr,ssim,lpips\n")
+    for idx, (tup, ground_truth) in enumerate(pbar):
+        restore = net(tup.color_distorted.to(device), tup.reference.to(device))
+        ground_truth = ground_truth.to(device)
+        scores = {
+            "PSNR": psnr(restore, ground_truth).mean().item(),
+            "SSIM": ssim(restore, ground_truth).mean().item(),
+            "LPIPS": lpips(restore, ground_truth).mean().item(),
+        }
+        pbar.set_postfix(scores)
+        with open(log_path, "a") as f:
+            f.write(f"{idx},{scores['PSNR']},{scores['SSIM']},{scores['LPIPS']}\n")
+        if save_image:
+            torchvision.utils.save_image(restore, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    parser.add_argument("-s", "--source", required=True, type=str)
+    parser.add_argument("-d", "--destination", required=True, type=str)
+    parser.add_argument("--model", choices=list(model_dict.keys()), default="dualresnet")
+    parser.add_argument("--model_destination", required=True, type=str)
+    parser.add_argument("--device", default="cuda", type=str)
+    parser.add_argument("-o", "--option", default=[], action='append', type=str)
+    parser.add_argument("--batch_size", default=1, type=int)
+    parser.add_argument("--num_workers", default=4, type=int)
+    parser.add_argument("--epoch", default=30, type=int)
+    parser.add_argument("--save_image", action='store_true')
+    args = parser.parse_args()
+    configs = {o.split("=", 1)[0]: eval(o.split("=", 1)[1]) for o in args.option}
+    loader = load_dataset(image_dir=args.source, batch_size=args.batch_size, num_workers=args.num_workers, shuffle=True)
+    model_path = os.path.join(args.model_destination, f"{args.model}.pth")
+    model = build_model(model=args.model, load_path=model_path, **configs)
+    with torch.no_grad():
+        rendering(model, loader, os.path.join(args.destination, args.model), os.path.join(args.destination, args.model + '.csv'), device=args.device, save_image=args.save_image)

refcolorrestore/shelf/__init__.py

@@ -0,0 +1,15 @@
+from .arch import rt4ksr_rep, rt4krestore_rep
+from .archnerf import rt4kdual_rep
+from .srvgg_arch import srvgg, restorevgg, dualvgg
+from .srresnet_arch import srresnet, restoreresnet, dualresnet
+model_dict = {
+    "rt4ksr": rt4ksr_rep,
+    "rt4krestore": rt4krestore_rep,
+    "rt4kdual": rt4kdual_rep,
+    "srvgg": srvgg,
+    "restoreresnet": restoreresnet,
+    "dualresnet": dualresnet,
+    "srresnet": srresnet,
+    "restorevgg": restorevgg,
+    "dualvgg": dualvgg
+}

refcolorrestore/shelf/arch.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+from .modules import *
+
+
+####################################
+# DEFINE MODEL INSTANCES
+####################################
+
+
+class RT4KSR_Rep(nn.Module):
+    def __init__(self,
+                 num_channels_in,
+                 num_channels_out,
+                 num_feats,
+                 num_blocks,
+                 upscale,
+                 act,
+                 eca_gamma,
+                 is_train,
+                 forget,
+                 layernorm,
+                 residual) -> None:
+        super().__init__()
+        self.forget = forget
+        self.gamma = nn.Parameter(torch.zeros(1))
+        self.gaussian = torchvision.transforms.GaussianBlur(kernel_size=5, sigma=1)
+        self.upscale = upscale
+
+        self.down = nn.PixelUnshuffle(2)
+        self.up = nn.PixelShuffle(2)
+        self.head = nn.Sequential(nn.Conv2d(num_channels_in * (2**2), num_feats, 3, padding=1))
+
+        hfb = []
+        if is_train:
+            hfb.append(ResBlock(num_feats, ratio=2))
+        else:
+            hfb.append((RepResBlock(num_feats)))
+        hfb.append(act)
+        self.hfb = nn.Sequential(*hfb)
+
+        body = []
+        for i in range(num_blocks):
+            if is_train:
+                body.append(SimplifiedNAFBlock(in_c=num_feats, act=act, exp=2, eca_gamma=eca_gamma, layernorm=layernorm, residual=residual))
+            else:
+                body.append(SimplifiedRepNAFBlock(in_c=num_feats, act=act, exp=2, eca_gamma=eca_gamma, layernorm=layernorm, residual=residual))
+
+        self.body = nn.Sequential(*body)
+
+        tail = [LayerNorm2d(num_feats)]
+        if is_train:
+            tail.append(ResBlock(num_feats, ratio=2))
+        else:
+            tail.append(RepResBlock(num_feats))
+        self.tail = nn.Sequential(*tail)
+
+        self.upsample = nn.Sequential(
+            nn.Conv2d(num_feats, num_channels_out * ((2 * upscale) ** 2), 3, padding=1),
+            nn.PixelShuffle(upscale*2)
+        )
+
+    def forward(self, _, out_lr):  # only lr is used
+        x = out_lr
+        # stage 1
+        hf = x - self.gaussian(x)
+
+        # unshuffle to save computation
+        x_unsh = self.down(x)
+        hf_unsh = self.down(hf)
+
+        shallow_feats_hf = self.head(hf_unsh)
+        shallow_feats_lr = self.head(x_unsh)
+
+        # stage 2
+        deep_feats = self.body(shallow_feats_lr)
+        hf_feats = self.hfb(shallow_feats_hf)
+
+        # stage 3
+        if self.forget:
+            deep_feats = self.tail(self.gamma * deep_feats + hf_feats)
+        else:
+            deep_feats = self.tail(deep_feats)
+
+        out = self.upsample(deep_feats)
+        out += F.interpolate(x, scale_factor=self.upscale, mode='bilinear', align_corners=False)
+        return F.interpolate(out, size=_.shape[2:], mode='bilinear', align_corners=False).clamp(0, 1.)
+
+
+class RT4KRestore_Rep(RT4KSR_Rep):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, upscale=1, **kwargs)
+
+    def forward(self, x, _):  # only lr is used
+        shape = x.shape[-2:]
+        x = F.interpolate(x, size=(x.shape[-2]//2*2, x.shape[-1]//2*2), mode='bilinear', align_corners=False)
+        # stage 1
+        hf = x - self.gaussian(x)
+
+        # unshuffle to save computation
+        x_unsh = self.down(x)
+        hf_unsh = self.down(hf)
+
+        shallow_feats_hf = self.head(hf_unsh)
+        shallow_feats_lr = self.head(x_unsh)
+
+        # stage 2
+        deep_feats = self.body(shallow_feats_lr)
+        hf_feats = self.hfb(shallow_feats_hf)
+
+        # stage 3
+        if self.forget:
+            deep_feats = self.tail(self.gamma * deep_feats + hf_feats)
+        else:
+            deep_feats = self.tail(deep_feats)
+
+        out = self.upsample(deep_feats)
+        out += x
+        return F.interpolate(out, size=shape, mode='bilinear', align_corners=False).clamp(0, 1.)
+
+####################################
+# RETURN INITIALIZED MODEL INSTANCES
+####################################
+
+
+def rt4ksr_rep(
+        act_type: str = "gelu",
+        feature_channels: int = 24,
+        num_blocks: int = 4,
+        scale: int = 4,
+        is_train: bool = True,
+):
+    act = activation(act_type)
+    model = RT4KSR_Rep(num_channels_in=3,
+                       num_channels_out=3,
+                       num_feats=feature_channels,
+                       num_blocks=num_blocks,
+                       upscale=scale,
+                       act=act,
+                       eca_gamma=0,
+                       forget=False,
+                       is_train=is_train,
+                       layernorm=True,
+                       residual=True)
+    return model
+
+
+def rt4krestore_rep(
+        act_type: str = "gelu",
+        feature_channels: int = 24,
+        num_blocks: int = 4,
+        is_train: bool = True,
+):
+    act = activation(act_type)
+    model = RT4KRestore_Rep(num_channels_in=3,
+                            num_channels_out=3,
+                            num_feats=feature_channels,
+                            num_blocks=num_blocks,
+                            act=act,
+                            eca_gamma=0,
+                            forget=False,
+                            is_train=is_train,
+                            layernorm=True,
+                            residual=True)
+    return model

refcolorrestore/shelf/archnerf.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn.functional as F
+
+from .arch import *
+
+
+class RT4KDual_Rep(RT4KSR_Rep):
+    def __init__(self, num_channels_color, num_channels_gray, *args, **kwargs):
+        super().__init__(*args, **kwargs, upscale=1, num_channels_in=num_channels_color + num_channels_gray)
+        self.num_channels_color = num_channels_color
+        self.num_channels_gray = num_channels_gray
+
+    def forward(self, x, lr):
+        shape = x.shape[-2:]
+        x = F.interpolate(x, size=(x.shape[-2]//2*2, x.shape[-1]//2*2), mode='bilinear', align_corners=False)
+        base = F.interpolate(lr, size=(x.shape[-2]//2*2, x.shape[-1]//2*2), mode='bicubic', align_corners=False)
+        x = torch.cat([base, x], dim=1)
+
+        # Following is just copy from RT4KSR_Rep
+        # stage 1
+        hf = x - self.gaussian(x)
+
+        # unshuffle to save computation
+        x_unsh = self.down(x)
+        hf_unsh = self.down(hf)
+
+        shallow_feats_hf = self.head(hf_unsh)
+        shallow_feats_lr = self.head(x_unsh)
+
+        # stage 2
+        deep_feats = self.body(shallow_feats_lr)
+        hf_feats = self.hfb(shallow_feats_hf)
+
+        # stage 3
+        if self.forget:
+            deep_feats = self.tail(self.gamma * deep_feats + hf_feats)
+        else:
+            deep_feats = self.tail(deep_feats)
+
+        out = self.upsample(deep_feats)
+        out += base
+        return F.interpolate(out, size=shape, mode='bilinear', align_corners=False).clamp(0, 1.)
+
+####################################
+# RETURN INITIALIZED MODEL INSTANCES
+####################################
+
+
+def rt4kdual_rep(
+        act_type: str = "gelu",
+        feature_channels: int = 24,
+        num_blocks: int = 4,
+        is_train: bool = True,
+):
+    act = activation(act_type)
+    model = RT4KDual_Rep(num_channels_color=3,
+                         num_channels_gray=3,
+                         num_channels_out=3,
+                         num_feats=feature_channels,
+                         num_blocks=num_blocks,
+                         act=act,
+                         eca_gamma=0,
+                         forget=False,
+                         is_train=is_train,
+                         layernorm=True,
+                         residual=True)
+    return model