Add application file

app.py

@@ -0,0 +1,167 @@
+import gradio as gr
+from PIL import Image
+import numpy as np
+import cv2
+from zipfile import ZipFile
+# Функция обработки изображения
+import gradio as gr
+from PIL import Image
+import numpy as np
+import os
+
+import shutil
+import yaml
+from pathlib import Path
+from fabric_diffusion import FabricDiffusionPipeline
+import urllib.request
+import tarfile
+import capture
+from pytorch_lightning import Trainer
+import random
+import torch
+
+ZIP_FOLDER = "./ZIPS"
+os.makedirs(ZIP_FOLDER, exist_ok=True)
+
+def zip_folder(folder_path, output_zip):
+    with ZipFile(output_zip, 'w') as zipf:
+        for root, dirs, files in os.walk(folder_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                zipf.write(file_path, os.path.relpath(file_path, folder_path))
+
+def set_deterministic(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+def load_config(config_path):
+    with open(config_path, 'r') as file:
+        return yaml.safe_load(file)
+
+def run_flatten_texture(pipeline, input_image_path, output_path, n_samples=3):
+    os.makedirs(output_path, exist_ok=True)
+    texture_name = os.path.splitext(os.path.basename(input_image_path))[0]
+    texture_patch = pipeline.load_patch_data(input_image_path)
+    gen_imgs = pipeline.flatten_texture(texture_patch, n_samples=n_samples)
+    for i, gen_img in enumerate(gen_imgs):
+        gen_img.save(os.path.join(output_path, f'{texture_name}_gen_{i}.png'))
+
+def organize_images_into_structure(source_folder, new_folder):
+    os.makedirs(new_folder, exist_ok=True)
+    for file_name in os.listdir(source_folder):
+        source_file = os.path.join(source_folder, file_name)
+        if os.path.isfile(source_file) and file_name.lower().endswith(('.png', '.jpg', '.jpeg')):
+            folder_name = os.path.splitext(file_name)[0]
+            subfolder_path = os.path.join(new_folder, folder_name, "outputs")
+            os.makedirs(subfolder_path, exist_ok=True)
+            destination_file = os.path.join(subfolder_path, file_name)
+            shutil.copy(source_file, destination_file)
+
+# Create a directory for saving if it doesn't exist
+if not os.path.exists("saved_images"):
+    os.makedirs("saved_images")
+
+def decode_rgba(image: Image.Image) -> Image.Image:
+    image_array = np.array(image)
+    alpha_channel = image_array[:, :, 3]
+    coords = np.argwhere(alpha_channel == 255)
+    y_min, x_min = coords.min(axis=0)
+    y_max, x_max = coords.max(axis=0) + 1  # Include the max boundary
+    cropped_image_array = image_array[y_min:y_max, x_min:x_max]
+    cropped_rgb_array = cropped_image_array[:, :, :3]
+    
+    cropped_rgb_image = Image.fromarray(cropped_rgb_array, "RGB")
+    return cropped_rgb_image
+
+
+def process_image(orig_image, image_data):
+    """Processes and saves the original and cropped images."""
+    # Get the original image.  Convert to PIL Image if needed
+    
+    # Get the cropped image (composite). Convert to PIL Image if needed
+    cropped_image = image_data['composite']
+    # print(type(cropped_image), cropped_image.size, np.array(crop_image).shape)
+    # Generate unique filenames using timestamps to avoid overwriting
+    import datetime
+    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+    original_filename = os.path.join("saved_images", f"original_{timestamp}.png")
+    cropped_filename = os.path.join("saved_images", f"cropped_{timestamp}.png")
+
+    orig_image.save(original_filename)
+
+    decoded_cropped_image = decode_rgba(cropped_image)
+    decoded_cropped_image.save(cropped_filename)
+    return orig_image, original_filename, decoded_cropped_image, cropped_filename, timestamp
+
+def web_main(orig_image, image_data):
+    orig_image, original_filename, cropped_image, cropped_filename, timestamp = process_image(orig_image, image_data)
+    set_deterministic(seed=42)
+    config = load_config("config_demo.yaml")
+    device = config["hyperparameters"]["device"]
+    texture_checkpoint = config["hyperparameters"]["fb_checkpoint"]
+    print_checkpoint = config["hyperparameters"].get("print_checkpoint", None)
+    input_image = config["hyperparameters"]["input_image"]
+    save_fd_dir = config["hyperparameters"]["save_fd_dir"]
+    save_mp_dir = config["hyperparameters"]["save_mp_dir"]
+    # save_mp_dir = timestamp
+    n_samples = config["hyperparameters"]["n_samples"]
+
+    pipeline = FabricDiffusionPipeline(device, texture_checkpoint, print_checkpoint=print_checkpoint)
+    os.makedirs(save_fd_dir, exist_ok=True)
+    run_flatten_texture(pipeline, cropped_filename, output_path=save_fd_dir, n_samples=n_samples)
+
+    organize_images_into_structure(save_fd_dir, save_mp_dir)
+
+    data = capture.get_data(predict_dir=Path(save_mp_dir), predict_ds='sd')
+    module = capture.get_inference_module(pt=config["hyperparameters"]["checkpoint_name"])
+
+    decomp = Trainer(default_root_dir=Path(save_mp_dir), accelerator='gpu', devices=1, precision=16)
+    decomp.predict(module, data)
+
+    folder = f"cropped_{timestamp}_gen_0"
+    # for folder in os.listdir(save_mp_dir):
+    folder_path = os.path.join(save_mp_dir, folder)
+    
+    
+    if os.path.isdir(folder_path):
+        target_path = os.path.join(folder_path, "weights", "mask", "an_object_with_azertyuiop_texture",
+                                   "checkpoint-800", "outputs")
+        if os.path.exists(target_path):
+            for file_name in os.listdir(target_path):
+                file_path = os.path.join(target_path, file_name)
+                if os.path.isfile(file_path):
+                    shutil.move(file_path, folder_path)
+    
+    print(f"FOLDER: {folder_path}")
+    print(os.path.exists(original_filename))
+    shutil.copyfile(
+        original_filename,
+        os.path.join(folder_path, "outputs", original_filename.split("/")[-1])
+    )
+
+    shutil.copyfile(
+        cropped_filename,
+        os.path.join(folder_path, "outputs", cropped_filename.split("/")[-1])
+    )
+    zip_folder(
+        folder_path,
+        os.path.join(ZIP_FOLDER, f"{folder_path.split('/')[-1]}.zip")
+    )
+    return orig_image, cropped_image, os.path.join(ZIP_FOLDER, f"{folder_path.split('/')[-1]}.zip")
+    
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        orig_image = gr.Image(label="Orig image", type="pil")
+        image_editor = gr.ImageEditor(type="pil", crop_size="1:1", label="Edit Image: Crop the desired element")
+    with gr.Row():
+        output_image = gr.Image(label="Result")
+        crop_image = gr.Image(label="Crop_image")
+        zip_file = gr.File(label="Download Zip File")
+
+    process_button = gr.Button("Crop Element")
+    process_button.click(web_main, inputs=[orig_image, image_editor], outputs=[output_image, crop_image, zip_file])
+
+demo.launch(share=True, server_port=22856)

capture/__init__.py

@@ -0,0 +1,6 @@
+
+from .utils.model import get_inference_module
+from .utils.exp import get_data
+
+
+__all__ = ['get_inference_module', 'get_data']

capture/callbacks/__init__.py

@@ -0,0 +1,4 @@
+from .metrics import MetricLogging
+from .visualize import VisualizeCallback
+
+__all__ = ['MetricLogging', 'VisualizeCallback']

capture/callbacks/metrics.py

@@ -0,0 +1,29 @@
+import os
+from pathlib import Path
+from collections import OrderedDict
+
+from pytorch_lightning.callbacks import Callback
+
+from ..utils.log import append_csv, get_info
+
+
+class MetricLogging(Callback):
+    def __init__(self, weights: str, test_list: str, outdir: Path):
+        super().__init__()
+        print(outdir)
+        assert outdir.is_dir()
+
+        self.weights = weights
+        self.test_list = test_list
+        self.outpath = outdir/'eval.csv'
+
+    def on_test_end(self, trainer, pl_module):
+        weight_name, epoch = get_info(str(self.weights))
+        *_, test_set = self.test_list.parts
+
+        parsed = {k: f'{v}' for k,v in trainer.logged_metrics.items()}
+
+        odict = OrderedDict(name=weight_name, epoch=epoch, test_set=test_set)
+        odict.update(parsed)
+        append_csv(self.outpath, odict)
+        print(f'logged metrics in: {self.outpath}')

capture/callbacks/visualize.py

@@ -0,0 +1,101 @@
+from pathlib import Path
+
+import torch
+import pytorch_lightning as pl
+from torchvision.utils import make_grid, save_image
+from torchvision.transforms import Resize
+
+from capture.render import encode_as_unit_interval, gamma_encode
+
+
+class VisualizeCallback(pl.Callback):
+    def __init__(self, exist_ok: bool, out_dir: Path, log_every_n_epoch: int, n_batches_shown: int):
+        super().__init__()
+
+        self.out_dir = out_dir/'images_1'
+        if not exist_ok and (self.out_dir.is_dir() and len(list(self.out_dir.iterdir())) > 0):
+            print(f'directory {out_dir} already exists, press \'y\' to proceed')
+            x = input()
+            if x != 'y':
+                exit(1)
+
+        self.out_dir.mkdir(parents=True, exist_ok=True)
+
+        self.log_every_n_epoch = log_every_n_epoch
+        self.n_batches_shown = n_batches_shown
+        self.resize = Resize(size=[128,128], antialias=True)
+
+    def setup(self, trainer, module, stage):
+        self.logger = trainer.logger
+
+    def on_train_batch_end(self, *args):
+        self._on_batch_end(*args, split='train')
+
+    def on_validation_batch_end(self, *args):
+        self._on_batch_end(*args, split='valid')
+
+    def _on_batch_end(self, trainer, module, outputs, inputs, batch, *args, split):
+        x_src, x_tgt = inputs
+
+        # optim_idx:0=discr & optim_idx:1=generator
+        y_src, y_tgt = outputs[1]['y'] if isinstance(outputs, list) else outputs['y']
+
+        epoch = trainer.current_epoch
+        if epoch % self.log_every_n_epoch == 0 and batch <= self.n_batches_shown:
+            if x_src and y_src:
+                self._visualize_src(x_src, y_src, split=split, epoch=epoch, batch=batch, ds='src')
+            if x_tgt and y_tgt:
+                self._visualize_tgt(x_tgt, y_tgt, split=split, epoch=epoch, batch=batch, ds='tgt')
+
+    def _visualize_src(self, x, y, split, epoch, batch, ds):
+        #zipped = zip(x.albedo, x.roughness, x.normals, x.displacement, x.input, x.image,
+        #             y.albedo, y.roughness, y.normals, y.displacement, y.reco, y.image)
+
+        zipped = zip(x.albedo, x.roughness, x.normals, x.input, x.image,
+                     y.albedo, y.roughness, y.normals, y.reco, y.image)
+
+        grid = [self._visualize_single_src(*z) for z in zipped]
+
+        name = self.out_dir/f'{split}{epoch:05d}_{ds}_{batch}.jpg'
+        save_image(grid, name, nrow=1, padding=5)
+
+    @torch.no_grad()
+    def _visualize_single_src(self, a, r, n, input, mv, a_p, r_p, n_p, reco, mv_p):
+        n = encode_as_unit_interval(n)
+        n_p = encode_as_unit_interval(n_p)
+
+        mv_gt = [gamma_encode(o) for o in mv]
+        mv_pred = [gamma_encode(o) for o in mv_p]
+        reco = gamma_encode(reco)
+
+        maps = [input, a, r, n] + mv_gt + [reco, a_p, r_p, n_p] + mv_pred
+        maps = [self.resize(x.cpu()) for x in maps]
+        return make_grid(maps, nrow=len(maps) // 2, padding=0)
+
+    @torch.no_grad()
+    def _visualize_single_src_previous(self, a, r, n, d, input, mv, a_p, r_p, n_p, d_p, reco, mv_p):
+        n = encode_as_unit_interval(n)
+        n_p = encode_as_unit_interval(n_p)
+
+        mv_gt = [gamma_encode(o) for o in mv]
+        mv_pred = [gamma_encode(o) for o in mv_p]
+        reco = gamma_encode(reco)
+
+        maps = [input, a, r, n, d] + mv_gt + [reco, a_p, r_p, n_p, d_p] + mv_pred
+        maps = [self.resize(x.cpu()) for x in maps]
+        return make_grid(maps, nrow=len(maps)//2, padding=0)
+
+    def _visualize_tgt(self, x, y, split, epoch, batch, ds):
+        zipped = zip(x.input, y.albedo, y.roughness, y.normals, y.displacement)
+
+        grid = [self._visualize_single_tgt(*z) for z in zipped]
+
+        name = self.out_dir/f'{split}{epoch:05d}_{ds}_{batch}.jpg'
+        save_image(grid, name, nrow=1, padding=5)
+
+    @torch.no_grad()
+    def _visualize_single_tgt(self, input, a_p, r_p, n_p, d_p):
+        n_p = encode_as_unit_interval(n_p)
+        maps = [input, a_p, r_p, n_p, d_p]
+        maps = [self.resize(x.cpu()) for x in maps]
+        return make_grid(maps, nrow=len(maps), padding=0)

capture/data/.ipynb_checkpoints/source-checkpoint.py

@@ -0,0 +1,158 @@
+import typing
+from pathlib import Path
+from PIL import Image
+import cv2
+
+import torch
+from easydict import EasyDict
+import torchvision.transforms.functional as tf
+from torch.utils.data import Dataset
+
+from ..utils.log import get_matlist
+from . import augment as Aug
+
+
+class FabricsDataset(Dataset):
+    def __init__(
+            self,
+            split,
+            transform,
+            renderer,
+            matlist,
+            dir: typing.Optional[Path] = None,
+            set_seed_render: bool = False,
+            **kwargs
+    ):
+        assert dir.is_dir()
+        assert matlist.is_file()
+        assert split in ['train', 'valid', 'all']
+        self.set_seed_render = set_seed_render
+
+        files = get_matlist(matlist, dir)
+
+        # train/val/ split
+        self.split = split
+        k = int(len(files) * .95)
+        if split == 'train':
+            self.files = files[:k]
+        elif split == 'valid':
+            self.files = files[k:]
+        elif split == 'all':
+            self.files = files
+
+        print(f'FabricsDataset list={matlist}:{self.split}=[{len(self.files)}/{len(files)}]')
+
+        dtypes = ['normals', 'albedo', 'input', 'input']
+        self.tf = Aug.Pipeline(*transform, dtypes=dtypes)
+        self.renderer = renderer
+
+    def __getitem__(self, index, quick=False):
+        folder = self.folders[index]
+
+        N_path = folder / 'normal.png'
+        A_path = folder / 'basecolor.png'
+        R_path = folder / 'roughness.png'
+        D_path = folder / 'height.png'
+
+        N = tf.to_tensor(Image.open(N_path).convert('RGB'))
+        A = tf.to_tensor(Image.open(A_path).convert('RGB'))
+        R = tf.to_tensor(Image.open(R_path).convert('RGB'))
+        D_pil = cv2.imread(str(D_path), cv2.IMREAD_GRAYSCALE)
+        D = torch.from_numpy(D_pil)[None].repeat(3, 1, 1) / 255
+
+        # augmentation
+        N, A, R, D = self.tf([N, A, R, D])
+
+        if self.set_seed_render:
+            torch.manual_seed(hash(folder.name))
+        I, params = self.renderer([N, A, R, D], n_samples=1)
+        params = torch.stack(params)
+
+        # return homogenous object whatever the source: acg or sd
+        return EasyDict(
+            input=I[0],
+            input_params=params[:, 0],
+            normals=N,
+            albedo=A,
+            roughness=R,
+            displacement=D,
+            name=folder.name,
+        )
+
+    def __len__(self):
+        return len(self.files)
+
+
+class AmbientCG(Dataset):
+    def __init__(
+            self,
+            split,
+            transform,
+            renderer,
+            matlist,
+            dir: typing.Optional[Path] = None,
+            set_seed_render: bool = False,
+            **kwargs
+    ):
+        assert dir.is_dir()
+        assert matlist.is_file()
+        assert split in ['train', 'valid', 'all']
+        self.set_seed_render = set_seed_render
+
+        files = get_matlist(matlist, dir)
+
+        # train/val/ split
+        self.split = split
+        k = int(len(files) * .95)
+        if split == 'train':
+            self.files = files[:k]
+        elif split == 'valid':
+            self.files = files[k:]
+        elif split == 'all':
+            self.files = files
+
+        print(f'AmbientCG list={matlist}:{self.split}=[{len(self.files)}/{len(files)}]')
+
+        dtypes = ['normals', 'albedo', 'input', 'input']
+        self.tf = Aug.Pipeline(*transform, dtypes=dtypes)
+        self.renderer = renderer
+
+    def __getitem__(self, index, quick=False):
+        path = self.files[index]
+        name = path.stem.split('_')[0]
+        root = path.parent
+
+        N_path = root / f'{name}_2K-PNG_NormalGL.png'
+        N = tf.to_tensor(Image.open(N_path).convert('RGB'))
+
+        A_path = root / f'{name}_2K-PNG_Color.png'
+        A = tf.to_tensor(Image.open(A_path).convert('RGB'))
+
+        R_path = root / f'{name}_2K-PNG_Roughness.png'
+        R = tf.to_tensor(Image.open(R_path).convert('RGB'))
+
+        D_path = root / f'{name}_2K-PNG_Displacement.png'
+        D_pil = cv2.imread(str(D_path), cv2.IMREAD_GRAYSCALE)
+        D = torch.from_numpy(D_pil)[None].repeat(3, 1, 1) / 255
+
+        # augmentation
+        N, A, R, D = self.tf([N, A, R, D])
+
+        if self.set_seed_render:
+            torch.manual_seed(hash(name))
+        I, params = self.renderer([N, A, R, D], n_samples=1)
+        params = torch.stack(params)
+
+        # return homogenous object whatever the source: acg or sd
+        return EasyDict(
+            input=I[0],
+            input_params=params[:, 0],
+            normals=N,
+            albedo=A,
+            roughness=R,
+            displacement=D,
+            name=name,
+        )
+
+    def __len__(self):
+        return len(self.files)

capture/data/augment.py

@@ -0,0 +1,155 @@
+import random
+
+import torch
+import torchvision.transforms as T
+import torchvision.transforms.functional as tf
+import torch.nn.functional as F
+
+
+class RandomResizedCrop(T.RandomResizedCrop):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, x, dtypes):
+        """WARNING: torchvision v0.11. Wrapper to T.RandomResizedCrop.__call__"""
+        i, j, h, w = self.get_params(x[0], self.scale, self.ratio)
+        return [tf.resized_crop(img, i, j, h, w, self.size, self.interpolation) for img in x]
+
+class ColorJitter(T.ColorJitter):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x, dtypes):
+        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
+            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)
+
+        o = []
+        for img, dtype in zip(x, dtypes):
+            if dtype == 'albedo':
+                for fn_id in fn_idx:
+                    if fn_id == 0 and brightness_factor is not None:
+                        img = tf.adjust_brightness(img, brightness_factor)
+                    elif fn_id == 1 and contrast_factor is not None:
+                        img = tf.adjust_contrast(img, contrast_factor)
+                    elif fn_id == 2 and saturation_factor is not None:
+                        img = tf.adjust_saturation(img, saturation_factor)
+                    elif fn_id == 3 and hue_factor is not None:
+                        img = tf.adjust_hue(img, hue_factor)
+            o.append(img)
+        return o
+
+class RandomHorizontalFlip(T.RandomHorizontalFlip):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def flip_x(self, img, dtype):
+        if dtype == 'normals':
+            img[0] *= -1
+        return tf.hflip(img)
+
+    def forward(self, x, dtypes):
+        if torch.rand(1) < self.p:
+            return [self.flip_x(img, dtype) for img, dtype in zip(x, dtypes)]
+        return x
+
+class RandomVerticalFlip(T.RandomVerticalFlip):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def flip_y(self, img, dtype):
+        if dtype == 'normals':
+            img[1] *= -1
+        return tf.vflip(img)
+
+    def forward(self, x, dtypes):
+        if torch.rand(1) < self.p:
+            return [self.flip_y(img, dtype) for img, dtype in zip(x, dtypes)]
+        return x
+
+def deg0(x, y, z):
+    return torch.stack([ x,  y,  z])
+def deg90(x, y, z):
+    return torch.stack([-y,  x,  z])
+def deg180(x, y, z):
+    return torch.stack([-x, -y,  z])
+def deg270(x, y, z):
+    return torch.stack([ y, -x,  z])
+
+class RandomIncrementRotate:
+    def __init__(self, p):
+        self.p = p
+        self.angles = [0, 90, 180, 270]
+
+        # adjusts surface normals vector depending on rotation angle
+        self.f = { 0: deg0, 90: deg90, 180: deg180, 270: deg270 }
+
+    def rotate(self, img, theta, dtype):
+        if dtype == 'normals':
+            img = self.f[theta](*img)
+        return tf.rotate(img, theta)
+
+    def __call__(self, x, dtypes):
+        if torch.rand(1) < self.p:
+            theta = random.choice(self.angles)
+            return [self.rotate(img, theta, dtype) for img, dtype in zip(x, dtypes)]
+        return x
+
+class NormalizeGeometry:
+    def normalize(self, img, dtype):
+        if dtype == 'normals':
+            # perform [0, 1] -> [-1, 1] mapping
+            img = 2*img - 1
+            # normalize vector to unit sphere
+            img = F.normalize(img, dim=0)
+        return img
+
+    def __call__(self, x, dtypes):
+        return [self.normalize(img, dtype) for img, dtype in zip(x, dtypes)]
+
+class RandomCrop(T.RandomCrop):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x, dtypes):
+        img_size = tf.get_image_size(x[0])
+        assert all(tf.get_image_size(y) == img_size  for y in x)
+        i, j, h, w = self.get_params(x[0], self.size)
+        return [tf.crop(img, i, j, h, w) for img in x]
+
+class CenterCrop:
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, x, dtypes):
+        return [tf.center_crop(img, self.size) for img in x]
+
+class Resize(T.Resize):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x, dtypes):
+        return [super(Resize, self).forward(img) for img in x]
+
+class Identity():
+    def __call__(self, x, dtypes):
+        return x
+
+class ToTensor:
+    def __call__(self, x, dtypes):
+        return [tf.to_tensor(img) for img in x]
+
+class Pipeline:
+    DATA_TYPES = ['input', 'normals', 'albedo']
+
+    def __init__(self, *transforms, dtypes=None):
+        assert all(d in Pipeline.DATA_TYPES for d in dtypes)
+        self.dtypes = dtypes
+        self.transforms = transforms
+
+    def __call__(self, x):
+        #print(self.dtypes)
+        assert len(self.dtypes) == len(x)
+        assert all(y.shape[1:] == x[0].shape[1:] for y in x)
+        for f in self.transforms:
+            x = f(x, self.dtypes)
+        return x

capture/data/download.py

@@ -0,0 +1,39 @@
+import requests
+import argparse
+from pathlib import Path
+from io import BytesIO
+from zipfile import ZipFile
+
+from tqdm import tqdm
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('outdir', type=Path)
+args = parser.parse_args()
+
+url = 'https://ambientcg.com/get?file={}_2K-PNG.zip'
+
+print(f'saving files into {args.outdir}')
+args.outdir.mkdir(parents=True, exist_ok=True)
+
+cwd = Path(__file__).parent/'matlist/ambientcg'
+with open(cwd.resolve(), 'r') as file:
+    materials = [line.strip() for line in file]
+
+mtypes = ['Color', 'Roughness', 'NormalGL']
+for uid in tqdm(materials):
+    maps = [f'{uid}_2K-PNG_{x}.png' for x in mtypes]
+
+    if all((args.outdir/m).exists() for m in maps):
+        continue
+
+    link = url.format(uid)
+    try:
+        r = requests.get(link)
+        archive = ZipFile(BytesIO(r.content))
+    except:
+        print(f'{uid},{link}\n')
+        continue
+
+    for m in maps:
+        archive.extract(m, args.outdir)

capture/data/matlist/ambientcg

@@ -0,0 +1,1182 @@
+PavingStones063
+Ivory001B
+PavingStones127
+Fabric057
+Tiles095
+Ground022
+Tiles048
+Snow002
+Tiles018
+Tiles025
+Bricks035
+RoofingTiles001
+Gravel022
+Rock049
+Fabric019
+Concrete038
+Rock019
+Tiles014
+PaintedPlaster003
+Bricks031
+Leather034C
+Paper004
+Asphalt024C
+Tiles003
+Tiles053
+Tiles067
+Asphalt023S
+WoodFloor013
+Ground025
+Ground044
+PavingStones105
+Rock003
+Bricks060
+Rock017
+Bricks080C
+Tiles034
+Concrete027
+WoodFloor049
+PavingStones006
+Gravel026
+Rock037
+Rock045
+PaintedPlaster018
+Fabric037
+PaintedWood005
+Tiles098
+Asphalt014
+Ivory002A
+WoodFloor046
+Plastic009
+Facade014
+Rock025
+PaintedBricks004
+Rock018
+Paint005
+WoodSiding005
+Bricks065
+Carpet002
+Cork001
+PavingStones051
+Bricks006
+Marble002
+Tiles043
+Tiles102
+Tiles012
+Tiles074
+Rock042L
+Fabric064
+Leather013
+Asphalt028A
+Leather018
+PavingStones084
+Rock039
+Fabric028
+Tiles066
+Wood071
+Wood015
+Ground058
+Bricks051
+Wood043
+Plastic013A
+PaintedPlaster007
+Leather001
+Fabric032
+PavingStones050
+Concrete023
+PavingStones094
+Bricks017
+Bricks087
+Rock014
+Fabric010
+Plastic008
+Ground042
+PavingStones090
+Fabric067
+Concrete017
+Fabric025
+Bricks063
+Carpet004
+Bricks068
+Porcelain001
+PavingStones013
+Asphalt005
+WoodFloor055
+Bark012
+Granite002B
+Fabric079
+Ground012
+PavingStones098
+Gravel013
+Wicker011B
+WoodFloor016
+Lava003
+PavingStones046
+Asphalt008
+PavingStones061
+Leather035A
+Tiles026
+Gravel009
+Rocks004
+Fabric044
+Concrete025
+Fabric027
+Moss004
+PavingStones053
+Gravel015
+PavingStones132
+Porcelain002
+Leather011
+Plastic005
+Carpet010
+Plastic012B
+Ground020
+Marble024
+WoodFloor040
+Fabric039
+Tiles119
+Rock022
+PavingStones120
+Tiles013
+Asphalt010
+Leather005
+Bricks027
+Wood003
+PavingStones083
+Fabric071
+PavingStones095
+WoodFloor048
+Fabric059
+Concrete019
+Foam003
+Gravel024
+Tiles099
+Tiles065
+PaintedWood006C
+WoodFloor024
+Marble014
+Lava001
+Tiles006
+Ground014
+Tiles129A
+Tiles108
+Ground057
+WoodSiding006
+WoodFloor025
+Fabric041
+Wood066
+PavingStones014
+Grass002
+WoodFloor058
+Ground054
+Gravel032
+Marble023
+Tiles052
+Planks021
+PavingStones091
+Wood060
+PavingStones042
+PavingStones136
+TactilePaving003
+Asphalt016
+WoodFloor008
+Terrazzo007
+Tiles116
+Tiles017
+Asphalt013
+Wood033
+Wood053
+Fabric045
+Tiles068
+Plaster001
+Wallpaper001C
+Wood074
+Concrete020
+Rock042S
+Tiles094
+Rock047
+Lava005
+Gravel036S
+Concrete026
+Fabric062
+WoodFloor036
+PaintedPlaster002
+Bricks020
+Concrete036
+Snow001
+Tatami002
+Wood069
+Tiles051
+Gravel003
+Carpet011
+Concrete012
+Lava004
+Tiles113
+Bricks043
+Asphalt023L
+Asphalt025B
+Chipboard004
+Wood023
+PavingStones128
+Gravel035
+Tiles037
+PavingStones067
+Wood054
+WoodFloor054
+Plastic003
+Snow008C
+WoodFloor014
+WoodFloor015
+Grass001
+Marble016
+Leather021
+Tiles114
+Ground068
+Wood029
+Bark013
+Fabric033
+Plastic017A
+Leather034B
+Grass004
+Rocks020
+Gravel029
+Candy003
+Ground062L
+PavingStones102
+Concrete014
+Tiles122
+Tiles042
+Wood010
+Rock023
+WoodFloor056
+Gravel020
+Snow010C
+PavingStones115C
+Fabric054
+Asphalt029B
+Asphalt009
+Gravel012
+WoodFloor017
+Leather002
+Tiles044
+Rock046L
+PavingStones097
+PavingStones029
+Concrete039
+PavingStones074
+Tatami001
+Fabric056
+Tiles076
+WoodFloor020
+Concrete030
+PavingStones134
+PaintedWood007B
+Gravel017
+Marble021
+Wood051
+PavingStones117
+Tiles045
+PavingStones103
+Bricks022
+Asphalt021
+Planks029L
+Wicker008B
+Wood048
+Tiles125
+Rocks019
+Rock002
+Rock029
+PavingStones009
+Wood075
+Granite006B
+Wood078
+Fabric030
+Ground023
+Rocks012
+Granite003B
+Ground024
+Asphalt018
+Tiles029
+Tatami003
+PavingStones040
+Wood037
+Plastic004
+Leather029
+Leather036D
+Facade012
+Plastic010
+PavingStones052
+Grass003
+Rock011
+Leather004
+Cardboard004
+Bricks082B
+PavingStones077
+Fabric063
+Marble018
+Tiles022
+PavingStones082
+Bricks034
+Fabric007
+Fabric013
+Rock038
+Tiles120
+Tiles077
+Ivory001A
+PavingStones059
+Ground053
+Foam001
+Gravel005
+Terrazzo003
+PaintedBricks001
+Wood049
+Wicker007B
+Paint004
+PaintedWood008A
+Tiles072
+Rock008
+Gravel008
+PaintedPlaster014
+Wicker009B
+Tiles023
+Carpet006
+Planks003
+PavingStones033
+Bricks070
+Tiles083
+PavingStones039
+Wood031
+Tatami005
+Wicker007A
+PavingStones130
+WoodFloor041
+Leather020
+WoodFloor035
+Bricks061
+PaintedWood006B
+Wood041
+Road005
+Bricks037
+Leather034A
+Wood018
+Cardboard003
+Tiles085
+Tiles103
+Leather017
+PavingStones088
+Paint001
+Fabric047
+Tiles118
+Terrazzo012
+PavingStones078
+PaintedWood007C
+Wood042
+Granite001A
+Planks010
+Concrete040
+Wood084B
+Fabric058
+Marble004
+Bricks081A
+PavingStones020
+Plastic006
+PavingStones119
+Wood081
+Tiles050
+Concrete011
+ThatchedRoof002B
+Wood045
+Ground043
+WoodFloor021
+Wallpaper002A
+Bricks012
+PavingStones056
+Marble025
+Wood085B
+Cork002
+Terrazzo006
+Rocks024S
+PaintedPlaster001
+GlazedTerracotta002
+Tiles100
+Wood079
+WoodFloor052
+Gravel036L
+Chipboard001
+Granite004A
+Gravel011
+Gravel010
+PavingStones115B
+Fabric023
+PavingStones079
+Plastic014B
+Tiles089
+Asphalt028B
+Plastic012A
+Plastic014A
+Tiles004
+WoodFloor026
+PaintedWood001
+Wicker003
+WoodFloor004
+Terrazzo019M
+Clay003
+Asphalt029A
+PavingStones055
+Bricks011
+PaintedPlaster016
+Paint003
+Wood022
+Bricks040
+Bricks010
+WoodFloor047
+Marble003
+Ground036
+Leather033A
+PavingStones062
+WoodFloor033
+Bricks007
+Wood087
+Plastic011
+Fabric051
+PavingStones104
+WoodSiding013
+Wood052
+Tiles009
+Ground049A
+PaintedPlaster010
+PavingStones108
+WoodFloor060
+Ground011
+Granite005B
+PaintedWood008C
+Wood025
+Bricks076B
+Terrazzo013
+Rock026
+WoodFloor059
+PavingStones126B
+PavingStones015
+Rocks005
+Marble017
+Shells001
+Asphalt019
+Bricks032
+Rock010
+Bricks009
+Snow010B
+GlazedTerracotta001
+Rock052
+PavingStones038
+Wicker010B
+Ground039
+Tiles020
+Granite001B
+Asphalt026B
+Clay001
+PavingStones089
+Rock006
+Leather024
+Concrete022
+Concrete028
+PavingStones126A
+Ice003
+Ground069
+Wood007
+Carpet014
+Rock044
+Concrete021
+Tiles101
+Fabric075
+PavingStones112
+Tiles071
+Ice001
+PavingStones035
+Wood014
+Ground028
+Asphalt004
+Rocks008
+Tiles078
+Wood064
+Asphalt020L
+Bricks076A
+Rocks016
+WoodFloor034
+Terrazzo009
+Leather030
+WoodFloor053
+Wood068
+Snow007C
+Asphalt028C
+PaintedPlaster005
+Rocks001
+Marble026
+Fabric055
+Planks007
+Tiles106
+PavingStones125B
+Ground040
+Snow004
+Fabric034
+Ice004
+PavingStones065
+Rock016
+WoodSiding010
+Facade013
+Plaster004
+Fabric043
+Granite007B
+Wood009
+Wallpaper001A
+Bricks073C
+Marble009
+PavingStones124
+Tiles032
+Ground017
+Tiles040
+Marble005
+Leather033C
+Wicker002
+PavingStones008
+Plastic016B
+Bricks073A
+Asphalt025C
+Lava002
+Planks002
+Snow009A
+Wood035
+PavingStones060
+PavingStones133
+PavingStones129
+Tiles121
+Road007
+Snow008A
+Gravel007
+Tiles069
+Ground050
+Bamboo002C
+Wood008
+Asphalt022
+Fabric072
+Bark006
+Fabric021
+WoodChips001
+RoofingTiles004
+Wood026
+Bricks008
+Facade017
+Ground037
+RoofingTiles003
+Snow008B
+Carpet007
+Paint002
+PavingStones075
+PaintedWood007A
+ThatchedRoof001A
+Tiles110
+Wallpaper001B
+Terrazzo008
+WoodFloor051
+Tiles024
+Fabric011
+PavingStones135
+Gravel038
+Fabric022
+Ground018
+Wallpaper002B
+Gravel031
+Bricks005
+PavingStones058
+Rock050
+Bricks019
+Marble020
+WoodFloor022
+Leather019
+Ground051
+Tiles011
+WoodFloor011
+Leather023
+Fabric006
+Concrete003
+Bricks025
+Tiles107
+Wood020
+PavingStones093
+Wicker001
+PaintedPlaster009
+Tiles091
+PaintedPlaster006
+Marble001
+Road002
+Leather022
+Bricks064
+Wood013
+Leather010
+Rock005
+Tiles064
+Tiles054
+Bricks076C
+WoodFloor007
+Bricks067
+Wood057
+Bricks045
+Gravel021
+Plaster002
+Bark004
+Snow007A
+Asphalt030
+Tatami006
+Bark007
+Bricks041
+Wood017
+Bricks073B
+Ground033
+Bricks004
+Plastic015B
+Concrete045
+Wicker005
+Bricks074
+Tiles041
+Terrazzo004
+Clay002
+WoodFloor023
+Rock043S
+Rocks023
+Asphalt015
+Tiles035
+Wicker010A
+Bricks071
+Rocks013
+Planks029S
+Granite003A
+Tiles008
+Fabric068
+Tiles128A
+PavingStones086
+Fabric009
+Rocks017
+WoodFloor005
+Carpet005
+Wood070
+Rock040
+Clay004
+Wicker006
+Marble007
+Rocks011
+Wicker008A
+Ground066
+TactilePaving004
+Tiles019
+Fabric024
+PavingStones018
+Marble010
+Tiles046
+Asphalt012
+Fabric073
+Tiles086
+PavingStones011
+Rock024
+PaintedWood008B
+Rocks022
+Wood044
+PaintedBricks003
+Wood050
+PavingStones125A
+PavingStones025
+Ground056
+Wicker004
+Wood086
+Ground060
+Gravel028
+Wicker012A
+Rocks010
+Fabric005
+Rock007
+Tiles028
+Wood040
+Terrazzo011
+Leather014
+Wood004
+Sponge003
+Concrete032
+Carpet015
+Plastic007
+Tiles031
+Rock009
+Ground015
+Wood021
+PavingStones057
+Leather035D
+Tiles109
+Terrazzo019L
+Leather032
+Planks005
+Tiles111
+Rocks021
+Concrete033
+PavingStones101
+Road004
+PavingStones019
+PaintedWood009A
+PavingStones072
+Plaster005
+Gravel025
+PavingStones080
+Moss003
+TactilePaving002
+Wicker011A
+Asphalt024A
+Wood046
+PavingStones037
+Leather012
+Gravel016
+Plastic016A
+Asphalt017
+Leather025
+PavingStones036
+Bricks033
+Ground016
+Terrazzo005
+WoodFloor028
+Asphalt027C
+WoodFloor045
+Bricks014
+Fabric035
+Bricks028
+RoofingTiles010
+Asphalt027A
+PavingStones071
+WoodFloor050
+Ground049C
+RoofingTiles007
+Fabric074
+Bamboo002A
+Tiles112
+Bricks088
+Ground064
+Wicker009A
+WoodFloor027
+RoofingTiles008
+Gravel033
+Tiles084
+Concrete015
+Leather009
+Fabric014
+Wood089
+Rocks025
+Planks032
+Wood001
+Fabric069
+Carpet013
+PavingStones026
+PavingStones043
+Gravel034
+Terrazzo017
+WoodFloor012
+Chipboard007
+Concrete031
+Fabric016
+Ground007
+Tiles079
+Ivory002B
+PavingStones047
+Bricks013
+Marble006
+PavingStones016
+TactilePaving006
+Fabric050
+Rock033
+Ground034
+Rocks003
+Fabric026
+PavingStones024
+Rocks024L
+Bricks052
+Concrete005
+Tiles036
+Wicker012B
+Candy002
+Tiles090
+Ground059
+PavingStones045
+Rock035
+Wood067
+Plastic001
+Fabric077
+Tiles081
+Wood030
+Paper003
+Tiles073
+Marble013
+Fabric036
+Leather033B
+Tiles047
+Rocks014
+Tiles115
+Ground038
+PavingStones123
+Wood084A
+Leather028
+WoodSiding007
+PavingStones041
+Leather031
+Planks022
+PavingStones081
+Tiles097
+Chipboard006
+Snow010A
+Asphalt025A
+Tiles070
+Leather026
+Ground055S
+Wood065
+Tiles105
+Ground027
+Tiles082
+WoodFloor039
+Granite007A
+Tiles021
+Bricks029
+Leather034D
+Rock015
+PaintedPlaster004
+Plaster006
+Wood038
+PaintedPlaster012
+Wood085A
+Bricks057
+Marble015
+Ground013
+Gravel037
+Leather006
+Plastic015A
+WoodSiding012
+Wood028
+Gravel004
+Wood056
+WoodFloor009
+Wallpaper002C
+Leather003
+Tiles126
+Gravel018
+Wood063
+Bricks069
+Concrete037
+RoofingTiles002
+Chipboard002
+PavingStones031
+Facade016
+Tiles088
+PavingStones044
+Rock004
+Terrazzo016
+PavingStones066
+Snow009C
+Bricks049
+Tiles027
+PavingStones010
+Terrazzo010
+Wood082A
+Ground031
+Fabric031
+Ground047
+Fabric078
+Ground035
+Bricks039
+PaintedBricks002
+Granite002A
+Bark010
+Rock034
+PavingStones107
+Marble022
+PavingStones106
+Bricks080B
+Sponge001
+Cardboard002
+Fabric029
+Bricks023
+Paint006
+Bricks015
+Tiles124
+Leather035C
+Gravel019
+Bricks044
+PavingStones087
+Bricks038
+Terrazzo015
+Rock043L
+Bricks036
+PavingStones048
+Tiles016
+Tiles117
+Ground019
+Tiles030
+Rock048
+WoodFloor018
+Wood032
+Ice002
+Wood062
+Fabric015
+PavingStones099
+Terrazzo001
+ThatchedRoof002A
+Candy001
+Ground021
+Tiles075
+Rocks007
+PavingStones070
+PaintedPlaster015
+Fabric017
+Wood039
+Ground010
+WoodFloor019
+Fabric020
+Fabric070
+PavingStones076
+Rock013
+Leather008
+Tatami004
+Fabric076
+Fabric053
+PavingStones049
+PaintedPlaster008
+Marble011
+Bricks081C
+Asphalt024B
+PaintedPlaster013
+Ground029
+Tiles039
+Carpet003
+Plaster003
+Fabric048
+Fabric049
+Wood083A
+Wood005
+Road003
+PavingStones017
+TactilePaving001
+PavingStones116
+Wood002
+Leather016
+Leather036C
+Leather035B
+Snow003
+Fabric065
+Bricks066
+Rocks015
+Wood061
+Tiles093
+Wood019
+Concrete024
+Rock012
+PavingStones118
+Bricks050
+Road001
+Asphalt006
+Rock028
+PavingStones023
+Wood016
+PavingStones021
+Granite004B
+PavingStones005
+Bamboo001A
+GlazedTerracotta003
+Marble008
+Concrete006
+Tiles063
+Rock053
+Facade015
+Ground063
+PavingStones085
+Rocks018
+PaintedBricks005
+Terrazzo018
+Cardboard001
+Tiles080
+Tiles010
+Bricks021
+Tiles104
+Rock021
+PavingStones007
+Concrete004
+Leather007
+PaintedPlaster011
+Gravel027
+Moss001
+PaintedWood003
+Bamboo002B
+Concrete029
+WoodFloor003
+Gravel023
+Asphalt011
+Fabric061
+Wood072
+Rocks006
+WoodFloor029
+WoodFloor037
+Marble019
+Asphalt027B
+Wood058
+Bricks056
+Wood076
+PavingStones028
+Planks008
+Road006
+Tiles033
+Wood024
+Carpet001
+PavingStones027
+PavingStones115A
+Tiles096
+Asphalt026C
+Bricks018
+Plastic018A
+Marble012
+PavingStones100
+Asphalt007
+Tiles005
+Wood088
+Planks009
+PaintedWood004
+PavingStones073
+Chipboard003
+Bricks042
+Foam002
+WoodFloor057
+Ground062S
+Paper006
+Ground030
+RoofingTiles005
+PavingStones064
+Fabric046
+Gravel006
+Tiles128B
+Bark003
+Leather015
+Fabric038
+Bricks062
+Fabric060
+Fabric040
+WoodFloor006
+Rock032
+Leather027
+Wood006
+Snow006
+Wood027
+Ground046
+Rock046S
+Chipboard008
+Fabric052
+Bricks030
+Porcelain003
+Gravel014
+Chipboard005
+Tiles127
+Bricks053
+Fabric008
+Plastic018B
+Leather036A
+Granite006A
+Tiles129B
+Sponge002
+Ground052
+WoodFloor031
+Wood055
+Wood059
+Bamboo001B
+Bricks082C
+Planks004
+Concrete016
+Ground070
+Bricks026
+Plastic002
+PavingStones012
+Leather036B
+Snow007B
+Bricks048
+Leather037
+Concrete009
+WoodFloor032
+Ground055L
+Tiles087
+Rock031
+Fabric018
+PavingStones032
+WoodFloor030
+PavingStones114
+PavingStones022
+Granite005A
+Terrazzo019S
+Ground026
+Fabric012
+PaintedWood002
+PavingStones092
+Fabric066
+Bricks016
+Rock036
+Plastic013B
+Wood036
+Fabric042
+Rock020
+Ground045
+Ground049B
+Asphalt020S
+Concrete007
+TactilePaving005
+PavingStones113
+Terrazzo002
+Bamboo001C
+Wood073
+ThatchedRoof001B
+Concrete008
+PavingStones068
+Paper005
+Ground067
+Bricks081B
+Carpet008
+Moss002
+RoofingTiles006
+Bricks054
+PavingStones096
+Concrete018
+Wood082B
+Snow009B
+PavingStones109
+Bricks058
+PavingStones034
+Bricks047
+Ground009
+Rock030
+PaintedWood009B
+Carpet012
+RoofingTiles009
+Carpet009
+Tiles015
+Bricks024
+Gravel030
+PaintedWood009C
+Bark009
+PavingStones030
+Snow005
+Tiles049
+PaintedWood006A
+Ground041
+Rocks002
+Asphalt026A
+Ground048
+Bricks082A
+Terrazzo014
+Bricks046
+Bricks059
+Rocks009
+Tiles092
+Bricks080A
+Wood083B
+WoodFloor038
+Bricks055
+PavingStones131
+Tiles123
+Plastic017B
+Concrete010
+WoodFloor010
+Wood077

capture/data/matlist/texsd

@@ -0,0 +1,353 @@
+acrylic_000
+acrylic_002
+acrylic_003
+acrylic_004
+acrylic_005
+acrylic_006
+acrylic_007
+acrylic_008
+acrylic_009
+acrylic_010
+acrylic_011
+acrylic_012
+acrylic_013
+acrylic_014
+acrylic_016
+acrylic_017
+acrylic_018
+acrylic_019
+acrylic_020
+acrylic_021
+acrylic_022
+acrylic_023
+acrylic_024
+acrylic_025
+acrylic_026
+acrylic_027
+acrylic_028
+acrylic_029
+acrylic_030
+acrylic_031
+acrylic_032
+acrylic_033
+acrylic_034
+acrylic_035
+acrylic_036
+acrylic_037
+acrylic_038
+acrylic_039
+asphalt_000
+asphalt_001
+asphalt_002
+asphalt_003
+asphalt_004
+asphalt_005
+asphalt_006
+asphalt_007
+asphalt_008
+asphalt_009
+asphalt_010
+asphalt_011
+asphalt_012
+asphalt_013
+asphalt_014
+asphalt_015
+asphalt_016
+asphalt_017
+asphalt_018
+asphalt_019
+asphalt_020
+asphalt_021
+asphalt_022
+asphalt_023
+asphalt_025
+asphalt_026
+asphalt_027
+asphalt_028
+asphalt_029
+asphalt_030
+asphalt_031
+asphalt_032
+asphalt_033
+asphalt_034
+asphalt_035
+asphalt_036
+asphalt_037
+asphalt_038
+asphalt_039
+bamboo_000
+bamboo_001
+bamboo_002
+bamboo_003
+bamboo_004
+bamboo_005
+bamboo_006
+bamboo_007
+bamboo_008
+bamboo_009
+bamboo_010
+bamboo_011
+bamboo_012
+bamboo_013
+bamboo_014
+bamboo_015
+bamboo_016
+bamboo_017
+bamboo_018
+bamboo_019
+bamboo_020
+bamboo_021
+bamboo_022
+bamboo_023
+bamboo_024
+bamboo_025
+bamboo_026
+bamboo_027
+bamboo_028
+bamboo_029
+bamboo_030
+bamboo_031
+bamboo_032
+bamboo_033
+bamboo_034
+bamboo_035
+bamboo_036
+bamboo_037
+bamboo_038
+bamboo_039
+bambooflooring_000
+bambooflooring_001
+bambooflooring_002
+bambooflooring_003
+bambooflooring_004
+bambooflooring_005
+bambooflooring_006
+bambooflooring_007
+bambooflooring_008
+bambooflooring_009
+bambooflooring_010
+bambooflooring_011
+bambooflooring_012
+bambooflooring_013
+bambooflooring_014
+bambooflooring_015
+bambooflooring_016
+bambooflooring_017
+bambooflooring_018
+bambooflooring_019
+bambooflooring_020
+bambooflooring_021
+bambooflooring_022
+bambooflooring_023
+bambooflooring_024
+bambooflooring_025
+bambooflooring_026
+bambooflooring_027
+bambooflooring_028
+bambooflooring_029
+bambooflooring_030
+bambooflooring_031
+bambooflooring_032
+bambooflooring_033
+bambooflooring_034
+bambooflooring_035
+bambooflooring_036
+bambooflooring_037
+bambooflooring_038
+bambooflooring_039
+bark_000
+bark_001
+bark_002
+bark_003
+bark_004
+bark_005
+bark_006
+bark_007
+bark_008
+bark_009
+bark_010
+bark_011
+bark_012
+bark_013
+bark_014
+bark_015
+bark_016
+bark_017
+bark_018
+bark_019
+bark_020
+bark_021
+bark_022
+bark_023
+bark_024
+bark_025
+bark_026
+bark_027
+bark_028
+bark_029
+bark_030
+bark_031
+bark_032
+bark_033
+bark_034
+bark_035
+bark_036
+bark_037
+bark_038
+bark_039
+bricks_000
+bricks_001
+bricks_002
+bricks_003
+bricks_004
+bricks_005
+bricks_006
+bricks_007
+bricks_008
+bricks_009
+bricks_010
+bricks_011
+bricks_012
+bricks_013
+bricks_014
+bricks_015
+bricks_016
+bricks_017
+bricks_018
+bricks_019
+bricks_020
+bricks_021
+bricks_022
+bricks_023
+bricks_024
+bricks_025
+bricks_026
+bricks_028
+bricks_029
+bricks_030
+bricks_031
+bricks_032
+bricks_033
+bricks_034
+bricks_035
+bricks_036
+bricks_037
+bricks_038
+bricks_039
+brocade_000
+brocade_001
+brocade_002
+brocade_003
+brocade_004
+brocade_005
+brocade_006
+brocade_007
+brocade_008
+brocade_009
+brocade_010
+brocade_011
+brocade_012
+brocade_013
+brocade_014
+brocade_015
+brocade_016
+brocade_017
+brocade_018
+brocade_019
+brocade_020
+brocade_021
+brocade_022
+brocade_023
+brocade_024
+brocade_025
+brocade_026
+brocade_027
+brocade_028
+brocade_029
+brocade_030
+brocade_031
+brocade_032
+brocade_033
+brocade_034
+brocade_035
+brocade_036
+brocade_037
+brocade_038
+brocade_039
+burlap_000
+burlap_001
+burlap_002
+burlap_003
+burlap_004
+burlap_005
+burlap_006
+burlap_007
+burlap_008
+burlap_009
+burlap_010
+burlap_011
+burlap_012
+burlap_013
+burlap_014
+burlap_015
+burlap_016
+burlap_017
+burlap_018
+burlap_019
+burlap_020
+burlap_021
+burlap_022
+burlap_023
+burlap_024
+burlap_025
+burlap_026
+burlap_027
+burlap_028
+burlap_029
+burlap_030
+burlap_031
+burlap_033
+burlap_034
+burlap_035
+burlap_036
+burlap_037
+burlap_038
+burlap_039
+burnoutvelvet_000
+burnoutvelvet_001
+burnoutvelvet_002
+burnoutvelvet_003
+burnoutvelvet_004
+burnoutvelvet_005
+burnoutvelvet_006
+burnoutvelvet_007
+burnoutvelvet_008
+burnoutvelvet_009
+burnoutvelvet_010
+burnoutvelvet_011
+burnoutvelvet_012
+burnoutvelvet_013
+burnoutvelvet_014
+burnoutvelvet_015
+burnoutvelvet_016
+burnoutvelvet_017
+burnoutvelvet_018
+burnoutvelvet_019
+burnoutvelvet_020
+burnoutvelvet_021
+burnoutvelvet_022
+burnoutvelvet_023
+burnoutvelvet_024
+burnoutvelvet_025
+burnoutvelvet_027
+burnoutvelvet_028
+burnoutvelvet_029
+burnoutvelvet_031
+burnoutvelvet_032
+burnoutvelvet_033
+burnoutvelvet_034
+burnoutvelvet_036
+burnoutvelvet_037
+burnoutvelvet_038
+burnoutvelvet_039
+candy_000

capture/data/module.py

@@ -0,0 +1,218 @@
+import typing
+from pathlib import Path
+from pprint import pprint
+from torch.utils.data import DataLoader
+from pytorch_lightning import LightningDataModule
+
+from . import augment as Aug
+from ..render import Renderer
+from .utils import MultiLoader, EmptyDataset, collate_fn
+from .source import AmbientCG
+from .source import FabricsDataset
+from .target import StableDiffusion
+
+
+def is_set(x):
+    return x is not None
+
+class DataModule(LightningDataModule):
+    def __init__(
+        self,
+        batch_size: int = 1,
+        num_workers: int = 1,
+        transform: bool = None,
+        source_ds: str = '',
+        target_ds: str = '',
+        test_ds: str = '',
+        predict_ds: str = '',
+        source_list: typing.Optional[Path] = None,
+        target_list: typing.Optional[Path] = None,
+        target_val_list: typing.Optional[Path] = None,
+        test_list: typing.Optional[Path] = None,
+        predict_list: typing.Optional[Path] = None,
+        source_dir: typing.Optional[Path] = None,
+        target_dir: typing.Optional[Path] = None,
+        predict_dir: typing.Optional[Path] = None,
+        test_dir: typing.Optional[Path] = None,
+        pseudo_labels: bool = False,
+        input_size: int = 512,
+    ):
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.transform = transform
+        self.pseudo_labels = pseudo_labels
+
+        self.source_ds = source_ds
+        self.target_ds = target_ds
+        self.test_ds = test_ds
+        self.predict_ds = predict_ds
+
+        self.source_list = source_list
+        self.target_list = target_list
+        self.target_val_list = target_val_list
+        self.predict_list = predict_list
+        self.test_list = test_list
+
+        self.source_dir = source_dir
+        self.target_dir = target_dir
+        self.predict_dir = predict_dir
+        self.test_dir = test_dir
+
+        self.input_size = input_size
+        # self.use_ref = use_ref
+
+        assert self.source_ds or self.target_ds or self.test_ds or self.predict_ds
+        if self.source_ds:
+            assert is_set(source_list)
+        if self.target_ds:
+            assert is_set(target_list)
+            if self.target_ds != 'sd':
+                assert is_set(target_val_list)
+
+    def setup(self, stage: str):
+        renderer = Renderer(return_params=True)
+        eval_tf = [
+            Aug.NormalizeGeometry(),
+            # Aug.CenterCrop((2048,2048)),
+            Aug.Resize([self.input_size, self.input_size], antialias=True)]
+
+
+        if stage == 'fit':
+            if self.transform:
+                train_tf = [
+                    Aug.RandomResizedCrop((512,512), scale=(1/16, 1/4), ratio=(1.,1.)),
+                    # Aug.RandomCrop(self.input_size),
+                    Aug.NormalizeGeometry(),
+                    Aug.RandomHorizontalFlip(),
+                    Aug.RandomVerticalFlip(),
+                    Aug.RandomIncrementRotate(p=1.),
+                    Aug.ColorJitter(brightness=.2, hue=.05, contrast=0.1)
+                ]
+            else:
+                train_tf = [
+                    Aug.CenterCrop((self.input_size, self.input_size)),
+                    Aug.NormalizeGeometry()]
+            train_kwargs = dict(pseudo_labels=self.pseudo_labels,
+                                renderer=renderer,
+                                transform=train_tf)
+            print('stage fit:')
+            pprint(train_kwargs)
+
+            ## SOURCE train dataset
+            if self.source_ds == 'acg':
+                self.src_train = FabricsDataset(split='train',
+                                           dir=self.source_dir,
+                                           matlist=self.source_list,
+                                           **train_kwargs)
+            ## TARGET train dataset
+            if self.target_ds == 'sd':
+                self.tgt_train = StableDiffusion(split='train',
+                                                #  use_ref=self.use_ref,
+                                                 dir=self.target_dir,
+                                                 matlist=self.target_list,
+                                                 **train_kwargs)
+
+            if not self.source_ds:
+                self.src_train = EmptyDataset(len(self.tgt_train))
+            if not self.target_ds:
+                self.tgt_train = EmptyDataset(len(self.src_train))
+
+        if stage == 'fit' or stage == 'validate':
+            validate_kwargs = dict(transform=eval_tf,
+                                   renderer=renderer,
+                                   set_seed_render=True)
+
+            ## SOURCE validation dataset
+            if self.source_ds == 'acg':
+                self.src_valid = FabricsDataset(split='valid',
+                                           dir=self.source_dir,
+                                           matlist=self.source_list,
+                                           **validate_kwargs)
+            ## TARGET validation dataset
+            if self.target_ds == 'sd':
+                self.tgt_valid = StableDiffusion(split='valid',
+                                                 pseudo_labels=False,
+                                                 dir=self.target_dir,
+                                                #  use_ref=self.use_ref,
+                                                 matlist=self.target_list,
+                                                 **validate_kwargs)
+
+            if not self.source_ds:
+                self.src_valid = EmptyDataset(len(self.tgt_valid))
+            if not self.target_ds:
+                self.tgt_valid = EmptyDataset(len(self.src_valid))
+
+        elif stage == 'test':
+            assert self.test_ds
+
+            test_kwargs = dict(pseudo_labels=False,
+                               matlist=self.test_list,
+                               transform=eval_tf,
+                               renderer=renderer,
+                               dir=self.test_dir,
+                               set_seed_render=True)
+
+            if self.test_ds == 'acg':
+                self.eval = [FabricsDataset(split='all', **test_kwargs)]
+            elif self.test_ds == 'sd':
+                self.eval = [StableDiffusion(split='all', **test_kwargs)]
+
+        elif stage == 'predict':
+            predict_kwargs = dict(split='all',
+                                  pseudo_labels=False,
+                                  dir=self.predict_dir,
+                                  matlist=None,
+                                  transform=eval_tf,
+                                  renderer=renderer)
+
+            if self.predict_ds == 'sd':
+                self.ds = StableDiffusion(**predict_kwargs)
+
+    def train_dataloader(self):
+        src_dl = DataLoader(dataset=self.src_train,
+                            batch_size=self.batch_size,
+                            drop_last=True,
+                            shuffle=True,
+                            num_workers=self.num_workers,
+                            collate_fn=collate_fn)
+        tgt_dl = DataLoader(dataset=self.tgt_train,
+                            batch_size=self.batch_size,
+                            drop_last=True,
+                            shuffle=True,
+                            num_workers=self.num_workers,
+                            collate_fn=collate_fn)
+
+        mix = MultiLoader(src_dl, tgt_dl)
+        return mix
+
+    def val_dataloader(self):
+        src_dl = DataLoader(dataset=self.src_valid,
+                            batch_size=self.batch_size,
+                            drop_last=False,
+                            shuffle=False,
+                            num_workers=self.num_workers,
+                            collate_fn=collate_fn)
+        tgt_dl = DataLoader(dataset=self.tgt_valid,
+                            batch_size=self.batch_size,
+                            drop_last=False,
+                            shuffle=False,
+                            num_workers=self.num_workers,
+                            collate_fn=collate_fn)
+
+        mix = MultiLoader(src_dl, tgt_dl)
+        return mix
+
+    def test_dataloader(self):
+        return [DataLoader(dataset=ds,
+                           batch_size=self.batch_size,
+                           drop_last=False,
+                           shuffle=False,
+                           num_workers=self.num_workers) for ds in self.eval]
+
+    def predict_dataloader(self):
+        return DataLoader(dataset=self.ds,
+                          batch_size=1,
+                          drop_last=False,
+                          shuffle=False,
+                          num_workers=1)

capture/data/source.py

@@ -0,0 +1,184 @@
+import typing
+from pathlib import Path
+from PIL import Image
+import cv2
+
+import torch
+from easydict import EasyDict
+import torchvision.transforms.functional as tf
+from torch.utils.data import Dataset
+
+from ..utils.log import get_matlist
+from . import augment as Aug
+import numpy as np
+def get_files_from_directory(dir: Path):
+    return [f for f in dir.iterdir() if f.is_dir()]
+
+class FabricsDataset(Dataset):
+    def __init__(
+            self,
+            split,
+            transform,
+            renderer,
+            matlist,
+            dir: typing.Optional[Path] = None,
+            set_seed_render: bool = False,
+            **kwargs
+    ):
+        assert dir.is_dir()
+        assert matlist.is_file()
+        assert split in ['train', 'valid', 'all']
+        self.set_seed_render = set_seed_render
+
+        folders = get_files_from_directory(dir)
+        #folders = [folder for folder in folders if (folder / 'height.png').is_file()]
+        folders = [folder for folder in folders if (folder / 'normal.png').is_file()
+                   and (folder / 'roughness.png').is_file()
+                   and (folder / 'basecolor.png').is_file()]
+
+        valid_folders = []
+        for folder in folders:
+            R_path = folder / 'roughness.png'
+            R = Image.open(R_path).convert('RGB')
+            R_array = np.array(R)
+
+
+            if np.all(R_array == 255):
+                continue
+
+            valid_folders.append(folder)
+
+        folders = valid_folders
+
+        print("Размер датасета: ", len(folders))
+
+
+        # train/val/ split
+        self.split = split
+        k = int(len(folders) * .95)
+        if split == 'train':
+            self.folders = folders[:k]
+        elif split == 'valid':
+            self.folders = folders[k:]
+        elif split == 'all':
+            self.folders = folders
+
+        print(f'FabricsDataset list={matlist}:{self.split}=[{len(self.folders)}/{len(folders)}]')
+
+        #dtypes = ['normals', 'albedo', 'input', 'input']
+        dtypes = ['normals', 'albedo', 'input']
+        self.tf = Aug.Pipeline(*transform, dtypes=dtypes)
+        self.renderer = renderer
+
+    def __getitem__(self, index, quick=False):
+        folder = self.folders[index]
+
+        N_path = folder / 'normal.png'
+        A_path = folder / 'basecolor.png'
+        R_path = folder / 'roughness.png'
+        #D_path = folder / 'height.png'
+
+        N = tf.to_tensor(Image.open(N_path).convert('RGB'))
+        A = tf.to_tensor(Image.open(A_path).convert('RGB'))
+        R = tf.to_tensor(Image.open(R_path).convert('RGB'))
+        #D_pil = cv2.imread(str(D_path), cv2.IMREAD_GRAYSCALE)
+        #D = torch.from_numpy(D_pil)[None].repeat(3, 1, 1) / 255
+
+        # augmentation
+        #N, A, R, D = self.tf([N, A, R, D])
+        N, A, R = self.tf([N, A, R])
+
+        if self.set_seed_render:
+            torch.manual_seed(hash(folder.name))
+       # I, params = self.renderer([N, A, R, D], n_samples=1)
+        I, params = self.renderer([N, A, R], n_samples=1)
+        params = torch.stack(params)
+
+        # return homogenous object whatever the source: acg or sd
+        return EasyDict(
+            input=I[0],
+            input_params=params[:, 0],
+            normals=N,
+            albedo=A,
+            roughness=R,
+            #displacement=D,
+            name=folder.name,
+        )
+
+    def __len__(self):
+        return len(self.folders)
+
+
+class AmbientCG(Dataset):
+    def __init__(
+            self,
+            split,
+            transform,
+            renderer,
+            matlist,
+            dir: typing.Optional[Path] = None,
+            set_seed_render: bool = False,
+            **kwargs
+    ):
+        assert dir.is_dir()
+        assert matlist.is_file()
+        assert split in ['train', 'valid', 'all']
+        self.set_seed_render = set_seed_render
+
+        files = get_matlist(matlist, dir)
+
+        # train/val/ split
+        self.split = split
+        k = int(len(files) * .95)
+        if split == 'train':
+            self.files = files[:k]
+        elif split == 'valid':
+            self.files = files[k:]
+        elif split == 'all':
+            self.files = files
+
+        print(f'AmbientCG list={matlist}:{self.split}=[{len(self.files)}/{len(files)}]')
+
+        dtypes = ['normals', 'albedo', 'input', 'input']
+        self.tf = Aug.Pipeline(*transform, dtypes=dtypes)
+        self.renderer = renderer
+
+    def __getitem__(self, index, quick=False):
+        path = self.files[index]
+        name = path.stem.split('_')[0]
+        root = path.parent
+
+        N_path = root / f'{name}_2K-PNG_NormalGL.png'
+        N = tf.to_tensor(Image.open(N_path).convert('RGB'))
+
+        A_path = root / f'{name}_2K-PNG_Color.png'
+        A = tf.to_tensor(Image.open(A_path).convert('RGB'))
+
+        R_path = root / f'{name}_2K-PNG_Roughness.png'
+        R = tf.to_tensor(Image.open(R_path).convert('RGB'))
+
+        D_path = root / f'{name}_2K-PNG_Displacement.png'
+        D_pil = cv2.imread(str(D_path), cv2.IMREAD_GRAYSCALE)
+        D = torch.from_numpy(D_pil)[None].repeat(3, 1, 1) / 255
+
+        # augmentation
+        N, A, R, D = self.tf([N, A, R, D])
+
+        if self.set_seed_render:
+            torch.manual_seed(hash(name))
+        I, params = self.renderer([N, A, R, D], n_samples=1)
+        params = torch.stack(params)
+
+        # return homogenous object whatever the source: acg or sd
+        return EasyDict(
+            input=I[0],
+            input_params=params[:, 0],
+            normals=N,
+            albedo=A,
+            roughness=R,
+            displacement=D,
+            name=name,
+        )
+
+    def __len__(self):
+        return len(self.files)

capture/data/target.py

@@ -0,0 +1,74 @@
+import typing
+from pathlib import Path
+from PIL import Image
+import random
+
+import torch
+from easydict import EasyDict
+import torchvision.transforms.functional as tf
+from torch.utils.data import Dataset
+
+from ..utils.log import get_matlist
+from . import augment as Aug
+
+
+class StableDiffusion(Dataset):
+    def __init__(
+        self,
+        split,
+        pseudo_labels,
+        transform,
+        renderer,
+        matlist,
+        use_ref=False,
+        dir: typing.Optional[Path] = None,
+        **kwargs
+    ):
+        assert dir.is_dir()
+        assert split in ['train', 'valid', 'all']
+
+        self.split = split
+        self.renderer = renderer
+        self.pseudo_labels = pseudo_labels
+        self.use_ref = use_ref
+        # self.pl_dir = pl_dir
+
+        if matlist == None:
+            files = sorted(dir.rglob('**/outputs/*[0-9].png'))
+            files += sorted(dir.rglob('**/out_renorm/*[0-9].png'))
+            print(f'total={len(files)}')
+            files = [x for x in files if not (x.parent/f'{x.stem}_roughness.png').is_file()]
+            print(f'after={len(files)}')
+        else:
+            files = get_matlist(matlist, dir)
+
+        ### Train/Validation Split
+        k = int(len(files)*.98)
+        if split == 'train':
+            self.files = files[:k]
+        elif split == 'valid':
+            self.files = files[k:]
+        elif split == 'all':
+            self.files = files
+
+        random.shuffle(self.files)
+
+        print(f'StableDiffusion list={matlist}:{self.split}=[{len(self.files)}/{len(files)}]')
+
+        dtypes = ['input']
+        self.tf = Aug.Pipeline(*transform, dtypes=dtypes)
+
+    def __getitem__(self, index):
+        path = self.files[index]
+        name = path.stem
+
+        o = EasyDict(dir=str(path.parent), name=name)
+
+        I = tf.to_tensor(Image.open(path).convert('RGB'))
+
+        o.path = str(path)
+        o.input, *_ = self.tf([I])
+        return o
+
+    def __len__(self):
+        return len(self.files)

capture/data/utils.py

@@ -0,0 +1,47 @@
+import math
+
+from easydict import EasyDict
+from torch.utils.data import Dataset, default_collate
+
+
+class EmptyDataset(Dataset):
+    def __init__(self, length):
+        self.length = length
+    def __getitem__(self, _):
+        return None
+    def __len__(self):
+        return self.length
+
+class MultiLoader:
+    """Iterator wrapper to iterate over multiple dataloaders at the same time."""
+    def __init__(self, a, b):
+        # a = self._repeat(a, b)
+        self.loaders = [a,b]
+
+    def __iter__(self):
+        return zip(*self.loaders)
+
+    def __len__(self):
+        return min(map(len, self.loaders))
+
+    def _repeat(self, a, b):
+        if len(a) < len(b):
+            k = math.ceil(len(b)/len(a))
+            return RepeatLoader(a, k)
+        return a
+
+class RepeatLoader:
+    def __init__(self, loader, k):
+        self.loader = loader
+        self.k = k
+
+    def __iter__(self):
+        for _ in range(self.k):
+            for x in self.loader:
+                yield x
+
+    def __len__(self):
+        return self.k*len(self.loader)
+
+def collate_fn(data):
+    return data if None in data else EasyDict(default_collate(data))

capture/predict.yml

@@ -0,0 +1,26 @@
+archi: densemtl
+mode: predict
+logger:
+  project: ae_acg
+data:
+  batch_size: 1
+  num_workers: 10
+  input_size: 512
+  predict_ds: sd
+  predict_list: data/matlist/pbrsd_v2
+trainer:
+  accelerator: gpu
+  devices: 1
+  precision: 16
+routine:
+  lr: 2e-5
+loss:
+  use_source: True
+  use_target: False
+  reg_weight: 1
+  render_weight: 1
+  n_random_configs: 3
+  n_symmetric_configs: 6
+viz:
+  n_batches_shown: 5
+  log_every_n_epoch: 5

capture/render/__init__.py

@@ -0,0 +1,4 @@
+from .main import Renderer
+from .scene import Scene, generate_random_scenes, generate_specular_scenes, gamma_decode, gamma_encode, encode_as_unit_interval, decode_from_unit_interval
+
+__all__ = ['Renderer', 'Scene', 'generate_random_scenes', 'generate_specular_scenes', 'gamma_decode', 'gamma_encode', 'encode_as_unit_interval', 'decode_from_unit_interval']

capture/render/main.py

@@ -0,0 +1,154 @@
+import torch
+import numpy as np
+
+from .scene import Light, Scene, Camera, dot_product, normalize, generate_normalized_random_direction, gamma_encode
+
+
+class Renderer:
+    def __init__(self, return_params=False):
+        self.use_augmentation = False
+        self.return_params = return_params
+
+    def xi(self, x):
+        return (x > 0.0) * torch.ones_like(x)
+
+    def compute_microfacet_distribution(self, roughness, NH):
+        alpha = roughness**2
+        alpha_squared = alpha**2
+        NH_squared = NH**2
+        denominator_part = torch.clamp(NH_squared * (alpha_squared + (1 - NH_squared) / NH_squared), min=0.001)
+        return (alpha_squared * self.xi(NH)) / (np.pi * denominator_part**2)
+
+    def compute_fresnel(self, F0, VH):
+        # https://cdn2.unrealengine.com/Resources/files/2013SiggraphPresentationsNotes-26915738.pdf
+        return F0 + (1.0 - F0) * (1.0 - VH)**5
+
+    def compute_g1(self, roughness, XH, XN):
+        alpha = roughness**2
+        alpha_squared = alpha**2
+        XN_squared = XN**2
+        return 2 * self.xi(XH / XN) / (1 + torch.sqrt(1 + alpha_squared * (1.0 - XN_squared) / XN_squared))
+
+    def compute_geometry(self, roughness, VH, LH, VN, LN):
+        return self.compute_g1(roughness, VH, VN) * self.compute_g1(roughness, LH, LN)
+
+    def compute_specular_term(self, wi, wo, albedo, normals, roughness, metalness):
+        F0 = 0.04 * (1. - metalness) + metalness * albedo
+
+        # Compute the half direction
+        H = normalize((wi + wo) / 2.0)
+
+        # Precompute some dot product
+        NH = torch.clamp(dot_product(normals, H), min=0.001)
+        VH = torch.clamp(dot_product(wo, H), min=0.001)
+        LH = torch.clamp(dot_product(wi, H), min=0.001)
+        VN = torch.clamp(dot_product(wo, normals), min=0.001)
+        LN = torch.clamp(dot_product(wi, normals), min=0.001)
+
+        F = self.compute_fresnel(F0, VH)
+        G = self.compute_geometry(roughness, VH, LH, VN, LN)
+        D = self.compute_microfacet_distribution(roughness, NH)
+
+        return F * G * D / (4.0 * VN * LN)
+
+    def compute_diffuse_term(self, albedo, metalness):
+        return  albedo * (1. - metalness) / np.pi
+
+    def evaluate_brdf(self, wi, wo, normals, albedo, roughness, metalness):
+        diffuse_term = self.compute_diffuse_term(albedo, metalness)
+        specular_term = self.compute_specular_term(wi, wo, albedo, normals, roughness, metalness)
+        return diffuse_term, specular_term
+
+    def render(self, scene, svbrdf):
+        #normals, albedo, roughness, displacement = svbrdf
+        normals, albedo, roughness = svbrdf
+        device = albedo.device
+
+        # Generate surface coordinates for the material patch
+        # The center point of the patch is located at (0, 0, 0) which is the center of the global coordinate system.
+        # The patch itself spans from (-1, -1, 0) to (1, 1, 0).
+        xcoords_row = torch.linspace(-1, 1, albedo.shape[-1], device=device)
+        xcoords = xcoords_row.unsqueeze(0).expand(albedo.shape[-2], albedo.shape[-1]).unsqueeze(0)
+        ycoords = -1 * torch.transpose(xcoords, dim0=1, dim1=2)
+        coords = torch.cat((xcoords, ycoords, torch.zeros_like(xcoords)), dim=0)
+
+        # We treat the center of the material patch as focal point of the camera
+        camera_pos = scene.camera.pos.unsqueeze(-1).unsqueeze(-1).to(device)
+        relative_camera_pos = camera_pos - coords
+        wo = normalize(relative_camera_pos)
+
+        # Avoid zero roughness (i. e., potential division by zero)
+        roughness = torch.clamp(roughness, min=0.001)
+
+        light_pos = scene.light.pos.unsqueeze(-1).unsqueeze(-1).to(device)
+        relative_light_pos = light_pos - coords
+        wi = normalize(relative_light_pos)
+
+        fdiffuse, fspecular  = self.evaluate_brdf(wi, wo, normals, albedo, roughness, metalness=0)
+        f = fdiffuse + fspecular
+
+        color = scene.light.color if torch.is_tensor(scene.light.color) else torch.tensor(scene.light.color)
+        light_color = color.unsqueeze(-1).unsqueeze(-1).unsqueeze(0).to(device)
+        falloff     = 1.0 / torch.sqrt(dot_product(relative_light_pos, relative_light_pos))**2 # Radial light intensity falloff
+        LN = torch.clamp(dot_product(wi, normals), min=0.0) # Only consider the upper hemisphere
+        radiance    = torch.mul(torch.mul(f, light_color * falloff), LN)
+
+        return radiance
+
+    def _get_input_params(self, n_samples, light, pose):
+        min_eps = 0.001
+        max_eps = 0.02
+        light_distance = 2.197
+        view_distance = 2.75
+
+        # Generate scenes (camera and light configurations)
+        # In the first configuration, the light and view direction are guaranteed to be perpendicular to the material sample.
+        # For the remaining cases, both are randomly sampled from a hemisphere.
+        view_dist = torch.ones(n_samples-1) * view_distance
+        if pose is None:
+            view_poses = torch.cat([torch.Tensor(2).uniform_(-0.25, 0.25), torch.ones(1) * view_distance], dim=-1).unsqueeze(0)
+            if n_samples > 1:
+                hemi_views = generate_normalized_random_direction(n_samples - 1, min_eps=min_eps, max_eps=max_eps) * view_distance
+                view_poses = torch.cat([view_poses, hemi_views])
+        else:
+            assert torch.is_tensor(pose)
+            view_poses = pose.cpu()
+
+        if light is None:
+            light_poses = torch.cat([torch.Tensor(2).uniform_(-0.75, 0.75), torch.ones(1) * light_distance], dim=-1).unsqueeze(0)
+            if n_samples > 1:
+                hemi_lights = generate_normalized_random_direction(n_samples - 1, min_eps=min_eps, max_eps=max_eps) * light_distance
+                light_poses = torch.cat([light_poses, hemi_lights])
+        else:
+            assert torch.is_tensor(light)
+            light_poses = light.cpu()
+
+        light_colors = torch.Tensor([10.0]).unsqueeze(-1).expand(n_samples, 3)
+
+        return view_poses, light_poses, light_colors
+
+    def __call__(self, svbrdf, n_samples=1, lights=None, poses=None):
+        view_poses, light_poses, light_colors = self._get_input_params(n_samples, lights, poses)
+
+        renderings = []
+        for wo, wi, c in zip(view_poses, light_poses, light_colors):
+            scene = Scene(Camera(wo), Light(wi, c))
+            rendering = self.render(scene, svbrdf)
+
+            # Simulate noise
+            std_deviation_noise = torch.exp(torch.Tensor(1).normal_(mean = np.log(0.005), std=0.3)).numpy()[0]
+            noise = torch.zeros_like(rendering).normal_(mean=0.0, std=std_deviation_noise)
+
+            # clipping
+            post_noise = torch.clamp(rendering + noise, min=0.0, max=1.0)
+
+            # gamma encoding
+            post_gamma = gamma_encode(post_noise)
+
+            renderings.append(post_gamma)
+
+        renderings = torch.cat(renderings, dim=0)
+
+        if self.return_params:
+            return renderings, (view_poses, light_poses, light_colors)
+        return renderings

capture/render/scene.py

@@ -0,0 +1,105 @@
+import math
+import torch
+
+
+def encode_as_unit_interval(tensor):
+    """
+        Maps range [-1, 1] to [0, 1]
+    """
+    return (tensor + 1) / 2
+
+def decode_from_unit_interval(tensor):
+    """
+        Maps range [0, 1] to [-1, 1]
+    """
+    return tensor * 2 - 1
+
+def gamma_decode(images):
+    return torch.pow(images, 2.2)
+
+def gamma_encode(images):
+    return torch.pow(images, 1.0/2.2)
+
+def dot_product(a, b):
+    return torch.sum(torch.mul(a, b), dim=-3, keepdim=True)
+
+def normalize(a):
+    return torch.div(a, torch.sqrt(dot_product(a, a)))
+
+def generate_normalized_random_direction(count, min_eps = 0.001, max_eps = 0.05):
+    r1 = torch.Tensor(count, 1).uniform_(0.0 + min_eps, 1.0 - max_eps)
+    r2 = torch.Tensor(count, 1).uniform_(0.0, 1.0)
+
+    r   = torch.sqrt(r1)
+    phi = 2 * math.pi * r2
+
+    x = r * torch.cos(phi)
+    y = r * torch.sin(phi)
+    z = torch.sqrt(1.0 - r**2)
+
+    return torch.cat([x, y, z], axis=-1)
+
+def generate_random_scenes(count):
+    # Randomly distribute both, view and light positions
+    view_positions  = generate_normalized_random_direction(count, 0.001, 0.1) # shape = [count, 3]
+    light_positions = generate_normalized_random_direction(count, 0.001, 0.1)
+
+    scenes = []
+    for i in range(count):
+        c = Camera(view_positions[i])
+        # Light has lower power as the distance to the material plane is not as large
+        l = Light(light_positions[i], [20.]*3)
+        scenes.append(Scene(c, l))
+
+    return scenes
+
+def generate_specular_scenes(count):
+    # Only randomly distribute view positions and place lights in a perfect mirror configuration
+    view_positions  = generate_normalized_random_direction(count, 0.001, 0.1) # shape = [count, 3]
+    light_positions = view_positions * torch.Tensor([-1.0, -1.0, 1.0]).unsqueeze(0)
+
+    # Reference: "parameters chosen empirically to have a nice distance from a -1;1 surface.""
+    distance_view  = torch.exp(torch.Tensor(count, 1).normal_(mean=0.5, std=0.75))
+    distance_light = torch.exp(torch.Tensor(count, 1).normal_(mean=0.5, std=0.75))
+
+    # Reference: "Shift position to have highlight elsewhere than in the center."
+    # NOTE: This code only creates guaranteed specular highlights in the orthographic rendering, not in the perspective one.
+    #       This is because the camera is -looking- at the center of the patch.
+    shift = torch.cat([torch.Tensor(count, 2).uniform_(-1.0, 1.0), torch.zeros((count, 1)) + 0.0001], dim=-1)
+
+    view_positions  = view_positions  * distance_view  + shift
+    light_positions = light_positions * distance_light + shift
+
+    scenes = []
+    for i in range(count):
+        c = Camera(view_positions[i])
+        l = Light(light_positions[i], [20, 20.0, 20.0])
+        scenes.append(Scene(c, l))
+
+    return scenes
+
+class Camera:
+    def __init__(self, pos):
+        self.pos = pos
+    def __str__(self):
+        return f'Camera({self.pos.tolist()})'
+
+class Light:
+    def __init__(self, pos, color):
+        self.pos   = pos
+        self.color = color
+    def __str__(self):
+        return f'Light({self.pos.tolist()}, {self.color})'
+
+class Scene:
+    def __init__(self, camera, light):
+        self.camera = camera
+        self.light  = light
+    def __str__(self):
+        return f'Scene({self.camera}, {self.light})'
+    @classmethod
+    def load(cls, o):
+        cam, light, color = o
+        return Scene(Camera(cam), Light(light, color))
+    def export(self):
+        return [self.camera.pos, self.light.pos, self.light.color]

capture/source/__init__.py

@@ -0,0 +1,5 @@
+from .model import ResnetEncoder, MultiHeadDecoder, DenseMTL
+from .loss import DenseReg, RenderingLoss
+from .routine import Vanilla
+
+__all__ = ['ResnetEncoder', 'MultiHeadDecoder', 'DenseMTL', 'DenseReg', 'RenderingLoss', 'Vanilla']

capture/source/loss.py

@@ -0,0 +1,147 @@
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from easydict import EasyDict
+import torch.nn.functional as F
+import torchvision.transforms.functional as tf
+
+from ..render import Renderer, Scene, generate_random_scenes, generate_specular_scenes
+
+
+class RenderingLoss(nn.Module):
+    def __init__(self, renderer, n_random_configs=0, n_symmetric_configs=0):
+        super().__init__()
+        self.eps = 0.1
+        self.renderer = renderer
+        self.n_random_configs = n_random_configs
+        self.n_symmetric_configs = n_symmetric_configs
+        self.n_renders = n_random_configs + n_symmetric_configs
+
+    def generate_scenes(self):
+        return generate_random_scenes(self.n_random_configs) + generate_specular_scenes(self.n_symmetric_configs)
+
+    def multiview_render(self, y, x):
+        X_renders, Y_renders = [], []
+
+        #x_svBRDFs = zip(x.normals, x.albedo, x.roughness, x.displacement)
+        #y_svBRDFs = zip(y.normals, y.albedo, y.roughness, x.displacement)
+
+        x_svBRDFs = zip(x.normals, x.albedo, x.roughness)
+        y_svBRDFs = zip(y.normals, y.albedo, y.roughness)
+        for x_svBRDF, y_svBRDF in zip(x_svBRDFs, y_svBRDFs):
+            x_renders, y_renders = [], []
+            for scene in self.generate_scenes():
+                x_renders.append(self.renderer.render(scene, x_svBRDF))
+                y_renders.append(self.renderer.render(scene, y_svBRDF))
+            X_renders.append(torch.cat(x_renders))
+            Y_renders.append(torch.cat(y_renders))
+
+        out = torch.stack(X_renders), torch.stack(Y_renders)
+        return out
+
+    def reconstruction(self, y, theta):
+        views = []
+        #for *svBRDF, t in zip(y.normals, y.albedo, y.roughness, y.displacement, theta):
+        for *svBRDF, t in zip(y.normals, y.albedo, y.roughness, theta):
+            render = self.renderer.render(Scene.load(t), svBRDF)
+            views.append(render)
+        return torch.cat(views)
+
+    def __call__(self, y, x, **kargs):
+        loss = F.l1_loss(torch.log(y + self.eps), torch.log(x + self.eps), **kargs)
+        return loss
+
+class DenseReg(nn.Module):
+    def __init__(
+        self,
+        reg_weight: float,
+        render_weight: float,
+        pl_reg_weight: float = 0.,
+        pl_render_weight: float = 0.,
+        use_source: bool = True,
+        use_target: bool = True,
+        n_random_configs= 3,
+        n_symmetric_configs = 6,
+    ):
+        super().__init__()
+
+        self.weights = [('albedo', reg_weight, self.log_l1),
+                        ('roughness', reg_weight, self.log_l1),
+                        ('normals', reg_weight, F.l1_loss)]
+
+        self.reg_weight = reg_weight
+        self.render_weight = render_weight
+        self.pl_reg_weight = pl_reg_weight
+        self.pl_render_weight = pl_render_weight
+        self.use_source = use_source
+        self.use_target = use_target
+
+        self.renderer = Renderer()
+        self.n_random_configs = n_random_configs
+        self.n_symmetric_configs = n_symmetric_configs
+        self.loss = RenderingLoss(self.renderer, n_random_configs=n_random_configs, n_symmetric_configs=n_symmetric_configs)
+
+    def log_l1(self, x, y, **kwargs):
+        return F.l1_loss(torch.log(x + 0.01), torch.log(y + 0.01), **kwargs)
+
+    def forward(self, x, y):
+        loss = EasyDict()
+        x_src, x_tgt = x
+        y_src, y_tgt = y
+
+        if self.use_source:
+            # acg regression loss
+            for k, w, loss_fn in self.weights:
+                loss[k] = w*loss_fn(y_src[k], x_src[k])
+
+            # rendering loss
+            x_src.image, y_src.image = self.loss.multiview_render(y_src, x_src)
+            loss.render = self.render_weight*self.loss(y_src.image, x_src.image)
+
+            # reconstruction
+            y_src.reco = self.loss.reconstruction(y_src, x_src.input_params)
+
+        if self.use_target:
+            for k, w, loss_fn in self.weights:
+                loss[f'tgt_{k}'] = self.pl_reg_weight*loss_fn(y_tgt[k], x_tgt[k])
+
+            # rendering loss w/ pseudo label
+            y_tgt.image, x_tgt.image = self.loss.multiview_render(y_tgt, x_tgt)
+            loss.sd_render = self.pl_render_weight*self.loss(y_tgt.image, x_tgt.image)
+
+            # reconstruction
+            y_tgt.reco = self.loss.reconstruction(y_tgt, x_tgt.input_params)
+
+        loss.total = torch.stack(list(loss.values())).sum()
+        return loss
+
+    @torch.no_grad()
+    def test(self, x, y, batch_idx, epoch, dl_id):
+        assert len(x.name) == 1
+        y.reco = self.loss.reconstruction(y, x.input_params)
+        return EasyDict(total=0)
+
+    @torch.no_grad()
+    def predict(self, x_tgt, y_tgt, batch_idx, split, epoch):
+        assert len(x_tgt.name) == 1
+
+        # gt components
+        I = x_tgt.input[0]
+        name = x_tgt.name[0]
+
+        # get the predicted maps
+        N_pred = y_tgt.normals[0]
+        A_pred = y_tgt.albedo[0]
+        R_pred = y_tgt.roughness[0]
+
+        # A_name = pl_path/f'{name}_albedo.png'
+        # save_image(A_pred, A_name)
+
+        # N_name = pl_path/f'{name}_normals.png'
+        # save_image(encode_as_unit_interval(N_pred), N_name)
+
+        # R_name = pl_path/f'{name}_roughness.png'
+        # save_image(R_pred, R_name)
+
+        return EasyDict(total=0)

capture/source/model.py

@@ -0,0 +1,233 @@
+# Adapted from monodepth2
+# https://github.com/nianticlabs/monodepth2/blob/master/networks/depth_decoder.py
+#
+# Copyright Niantic 2019. Patent Pending. All rights reserved.
+#
+# This software is licensed under the terms of the Monodepth2 licence
+# which allows for non-commercial use only, the full terms of which are made
+# available in the LICENSE file.
+
+from __future__ import absolute_import, division, print_function
+from collections import OrderedDict
+from easydict import EasyDict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import torch.utils.model_zoo as model_zoo
+
+
+class ConvBlock(torch.nn.Module):
+    """Layer to perform a convolution followed by ELU."""
+    def __init__(self, in_channels, out_channels, bn=False, dropout=0.0):
+        super(ConvBlock, self).__init__()
+
+        self.block = nn.Sequential(
+            Conv3x3(in_channels, out_channels),
+            nn.BatchNorm2d(out_channels) if bn else nn.Identity(),
+            nn.ELU(inplace=True),
+            # Pay attention: 2d version of dropout is used
+            nn.Dropout2d(dropout) if dropout > 0 else nn.Identity())
+
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+
+class Conv3x3(nn.Module):
+    """Layer to pad and convolve input with 3x3 kernels."""
+    def __init__(self, in_channels, out_channels, use_refl=True):
+        super(Conv3x3, self).__init__()
+
+        if use_refl:
+            self.pad = nn.ReflectionPad2d(1)
+        else:
+            self.pad = nn.ZeroPad2d(1)
+        self.conv = nn.Conv2d(int(in_channels), int(out_channels), 3)
+
+    def forward(self, x):
+        out = self.pad(x)
+        out = self.conv(out)
+        return out
+
+def upsample(x):
+    """Upsample input tensor by a factor of 2."""
+    return F.interpolate(x, scale_factor=2, mode="nearest")
+
+
+class ResNetMultiImageInput(models.ResNet):
+    """Constructs a resnet model with varying number of input images.
+    Adapted from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
+    """
+    def __init__(self, block, layers, num_classes=1000, in_channels=3):
+        super(ResNetMultiImageInput, self).__init__(block, layers)
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(
+            in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+
+def resnet_multiimage_input(num_layers, pretrained=False, in_channels=3):
+    """Constructs a ResNet model.
+    Args:
+        num_layers (int): Number of resnet layers. Must be 18 or 50
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        in_channels (int): Number of input channels
+    """
+    assert num_layers in [18, 50], "Can only run with 18 or 50 layer resnet"
+    blocks = {18: [2, 2, 2, 2], 50: [3, 4, 6, 3]}[num_layers]
+    block_type = {18: models.resnet.BasicBlock, 50: models.resnet.Bottleneck}[num_layers]
+    model = ResNetMultiImageInput(block_type, blocks, in_channels=in_channels)
+
+    if pretrained:
+        print('loading imagnet weights on resnet...')
+        loaded = model_zoo.load_url(models.resnet.model_urls['resnet{}'.format(num_layers)])
+        # loaded['conv1.weight'] = torch.cat(
+        #     (loaded['conv1.weight'], loaded['conv1.weight']), 1)
+        # diff = model.load_state_dict(loaded, strict=False)
+    return model
+
+
+class ResnetEncoder(nn.Module):
+    """Pytorch module for a resnet encoder
+    """
+    def __init__(self, num_layers, pretrained, in_channels=3):
+        super(ResnetEncoder, self).__init__()
+
+        self.num_ch_enc = np.array([64, 64, 128, 256, 512])
+
+        resnets = {18: models.resnet18,
+                   34: models.resnet34,
+                   50: models.resnet50,
+                   101: models.resnet101,
+                   152: models.resnet152}
+
+        if num_layers not in resnets:
+            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))
+
+        if in_channels > 3:
+            self.encoder = resnet_multiimage_input(num_layers, pretrained, in_channels)
+        else:
+            weights = models.ResNet101_Weights.IMAGENET1K_V1 if pretrained else None
+            self.encoder = resnets[num_layers](weights=weights)
+
+        if num_layers > 34:
+            self.num_ch_enc[1:] *= 4
+
+    def forward(self, x):
+        self.features = []
+
+        # input_image, normals = xx
+        # x = (input_image - 0.45) / 0.225
+        # x = torch.cat((input_image, normals),1)
+        x = self.encoder.conv1(x)
+        x = self.encoder.bn1(x)
+        self.features.append(self.encoder.relu(x))
+        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))
+        self.features.append(self.encoder.layer2(self.features[-1]))
+        self.features.append(self.encoder.layer3(self.features[-1]))
+        self.features.append(self.encoder.layer4(self.features[-1]))
+
+        return self.features
+
+
+class Decoder(nn.Module):
+    def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True,
+        kaiming_init=False, return_feats=False):
+        super().__init__()
+
+        self.num_output_channels = num_output_channels
+        self.use_skips = use_skips
+        self.upsample_mode = 'nearest'
+        self.scales = scales
+
+        self.return_feats = return_feats
+
+        self.num_ch_enc = num_ch_enc
+        self.num_ch_dec = np.array([16, 32, 64, 128, 256])
+
+        # decoder
+        self.convs = OrderedDict()
+        for i in range(4, -1, -1):
+            # upconv_0
+            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
+            num_ch_out = self.num_ch_dec[i]
+            self.convs[("upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out)
+
+            # upconv_1
+            num_ch_in = self.num_ch_dec[i]
+            if self.use_skips and i > 0:
+                num_ch_in += self.num_ch_enc[i - 1]
+            num_ch_out = self.num_ch_dec[i]
+            self.convs[("upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out)
+
+        # for s in self.scales:
+        self.convs[("dispconv", 0)] = Conv3x3(self.num_ch_dec[0], self.num_output_channels)
+
+        self.decoder = nn.ModuleList(list(self.convs.values()))
+        # self.sigmoid = nn.Sigmoid()
+
+        if kaiming_init:
+            print('init weights of decoder')
+            for m in self.children():
+                if isinstance(m, nn.Conv2d):
+                    nn.init.kaiming_normal_(m.weight)
+                    if m.bias is not None:
+                        m.bias.data.fill_(0.01)
+
+    def forward(self, input_features):
+        x = input_features[-1]
+        for i in range(4, -1, -1):
+            x = self.convs[("upconv", i, 0)](x)
+            x = [upsample(x)]
+            if self.use_skips and i > 0:
+                x += [input_features[i - 1]]
+            x = torch.cat(x, 1)
+            x = self.convs[("upconv", i, 1)](x)
+
+        # assert self.scales[0] == 0
+        final_conv = self.convs[("dispconv", 0)]
+        out = final_conv(x)
+
+        if self.return_feats:
+            return out, input_features[-1]
+        return out
+
+class MultiHeadDecoder(nn.Module):
+    def __init__(self, num_ch_enc, tasks, return_feats, use_skips):
+        super().__init__()
+        self.decoders = nn.ModuleDict({k:
+            Decoder(num_ch_enc=num_ch_enc,
+                    num_output_channels=num_ch,
+                    scales=[0],
+                    kaiming_init=False,
+                    use_skips=use_skips,
+                    return_feats=return_feats)
+            for k, num_ch in tasks.items()})
+
+    def forward(self, x):
+        y = EasyDict({k: v(x) for k, v in self.decoders.items()})
+        return y
+
+class DenseMTL(nn.Module):
+    def __init__(self, encoder, decoder):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+    def forward(self, x):
+        return self.decoder(self.encoder(x))

capture/source/routine.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+from torch import optim
+from pathlib import Path
+from easydict import EasyDict
+import pytorch_lightning as pl
+import torch.nn.functional as F
+import torchvision.transforms as T
+from torchvision.utils import save_image
+from torchmetrics import MeanSquaredError, StructuralSimilarityIndexMeasure
+#from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
+
+from . import DenseReg, RenderingLoss
+from ..render import Renderer, encode_as_unit_interval, gamma_decode, gamma_encode
+
+class Vanilla(pl.LightningModule):
+    metrics = ['I_mse','N_mse','A_mse','R_mse','I_ssim','N_ssim','A_ssim','R_ssim']
+    maps = {'I': 'reco', 'N': 'normals', 'R': 'roughness', 'A': 'albedo'}
+
+    def __init__(self, model: nn.Module, loss: DenseReg = None, lr: float = 0, batch_size: int = 0, max_images: int = 10):
+        super().__init__()
+        self.model = model
+        self.loss = loss
+        self.lr = lr
+        self.batch_size = batch_size
+        self.tanh = nn.Tanh()
+        self.norm = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        self.val_images = []
+        self.max_images = max_images
+
+        self.save_hyperparameters(ignore=['model', 'loss'])
+
+    def training_step(self, x):
+        y = self(*x)
+        loss = self.loss(x, y)
+        self.log_to('train', loss)
+        return dict(loss=loss.total, y=y)
+
+    def forward(self, src, tgt):
+        src_out, tgt_out = None, None
+
+        if None not in src:
+            src_out = self.model(self.norm(src.input))
+            self.post_process_(src_out)
+
+        if None not in tgt:
+            tgt_out = self.model(self.norm(tgt.input))
+            self.post_process_(tgt_out)
+
+        return src_out, tgt_out
+
+    def post_process_(self, o: EasyDict):
+        # (1) activation function, (2) concat unit z, (3) normalize to unit vector
+        nxy = self.tanh(o.normals)
+        nx, ny = torch.split(nxy*3, split_size_or_sections=1, dim=1)
+        n = torch.cat([nx, ny, torch.ones_like(nx)], dim=1)
+        o.normals = F.normalize(n, dim=1)
+
+        # (1) activation function, (2) mapping [-1,1]->[0,1]
+        a = self.tanh(o.albedo)
+        o.albedo = encode_as_unit_interval(a)
+
+        # (1) activation function, (2) mapping [-1,1]->[0,1], (3) channel repeat x3
+        r = self.tanh(o.roughness)
+        o.roughness = encode_as_unit_interval(r.repeat(1,3,1,1))
+
+    def validation_step(self, x, *_):
+        y = self(*x)
+        loss = self.loss(x, y)
+        self.log_to('val', loss)
+
+        if len(self.val_images) * self.batch_size < self.max_images and self.logger:
+            self.val_images.append((x, y))
+
+        return dict(loss=loss.total, y=y)
+
+    def on_validation_epoch_end(self):
+        cur_ind = 0
+        for ind, (x, y) in enumerate(self.val_images):
+            for key in ['normals', 'albedo', 'roughness']:
+                for i in range(len(x[0][key])):
+                    pred_image = y[0][key][i]
+                    gt_image = x[0][key][i]
+
+                    self.logger.experiment.add_image(f'val/{cur_ind}_{key}_pred', pred_image, self.global_step)
+                    self.logger.experiment.add_image(f'val/{cur_ind}_{key}_gt', gt_image, self.global_step)
+                cur_ind += 1
+                if cur_ind >= self.max_images:
+                    break
+
+        self.val_images.clear()
+
+    def log_to(self, split, loss):
+        self.log_dict({f'{split}/{k}': v for k, v in loss.items()}, batch_size=self.batch_size)
+
+    def log_images(self, x, y, split, max_images=5):
+        # Log predicted and ground truth images
+        for key in ['normals', 'albedo', 'roughness']:
+            for i in range(min(len(x[0][key]), max_images)):
+                self.logger.experiment.add_image(f'{split}/{key}_pred_{i}', y[0][key][i], self.global_step)
+                self.logger.experiment.add_image(f'{split}/{key}_gt_{i}', x[0][key][i], self.global_step)
+
+
+    def on_test_start(self):
+        self.renderer = RenderingLoss(Renderer())
+
+        for m in Vanilla.metrics:
+            if 'mse' in m:
+                setattr(self, m, MeanSquaredError().to(self.device))
+            elif 'ssim' in m:
+                setattr(self, m, StructuralSimilarityIndexMeasure(data_range=1).to(self.device))
+
+    def test_step(self, x, batch_idx, dl_id=0):
+        y = self.model(self.norm(x.input))
+        self.post_process_(y)
+
+        # image reconstruction
+        y.reco = self.renderer.reconstruction(y, x.input_params)
+        x.reco = gamma_decode(x.input)
+
+        for m in Vanilla.metrics:
+            mapid, *_ = m
+            k = Vanilla.maps[mapid]
+            meter = getattr(self, m)
+            meter(y[k], x[k].to(y[k].dtype))
+            self.log(m, getattr(self, m), on_epoch=True)
+
+        if self.logger:
+            self.log_images(x, y, split='test')
+
+    def predict_step(self, x, batch_idx):
+        y = self.model(self.norm(x.input))
+        self.post_process_(y)
+
+        I, name, outdir = x.input[0], x.name[0], Path(x.path[0]).parent
+        N_pred, A_pred, R_pred  = y.normals[0], y.albedo[0], y.roughness[0]
+
+        save_image(gamma_encode(A_pred), outdir/f'{name}_albedo.png')
+        save_image(encode_as_unit_interval(N_pred), outdir/f'{name}_normals.png')
+        save_image(R_pred, outdir/f'{name}_roughness.png')
+
+    def configure_optimizers(self):
+        optimizer = optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-4)
+        return dict(optimizer=optimizer)

capture/utils/__init__.py

@@ -0,0 +1,8 @@
+
+from .model import get_module
+from .cli import get_args
+from .exp import Trainer, get_name, get_callbacks, get_data
+from .log import get_logger
+
+
+__all__ = ['get_model', 'get_module', 'get_args', 'get_name', 'get_logger', 'get_data', 'get_callbacks', 'Trainer']

capture/utils/cli.py

@@ -0,0 +1,76 @@
+
+from pathlib import Path
+
+import jsonargparse
+import torch.nn as nn
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import WandbLogger
+
+from ..source import Vanilla, DenseReg
+from ..callbacks import VisualizeCallback
+from ..data.module import DataModule
+
+
+#! refactor this simplification required
+class LightningArgumentParser(jsonargparse.ArgumentParser):
+    """
+    Extension of jsonargparse.ArgumentParser to parse pl.classes and more.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def add_datamodule(self, datamodule_obj: pl.LightningDataModule):
+        self.add_method_arguments(datamodule_obj, '__init__', 'data', as_group=True)
+
+    def add_lossmodule(self, lossmodule_obj: nn.Module):
+        self.add_class(lossmodule_obj, 'loss')
+
+    def add_routine(self, model_obj: pl.LightningModule):
+        skip = {'ae', 'decoder', 'loss', 'transnet', 'model', 'discr', 'adv_loss', 'stage'}
+        self.add_class_arguments(model_obj, 'routine', as_group=True, skip=skip)
+
+    def add_logger(self, logger_obj):
+        skip = {'version', 'config', 'name', 'save_dir'}
+        self.add_class_arguments(logger_obj, 'logger', as_group=True, skip=skip)
+
+    def add_class(self, cls, group, **kwargs):
+        self.add_class_arguments(cls, group, as_group=True, **kwargs)
+
+    def add_trainer(self):
+        skip = {'default_root_dir', 'logger', 'callbacks'}
+        self.add_class_arguments(pl.Trainer, 'trainer', as_group=True, skip=skip)
+
+def get_args(datamodule=DataModule, loss=DenseReg, routine=Vanilla, viz=VisualizeCallback):
+    parser = LightningArgumentParser()
+
+    parser.add_argument('--config', action=jsonargparse.ActionConfigFile, required=True)
+    parser.add_argument('--archi', type=str, required=True)
+    parser.add_argument('--out_dir', type=lambda x: Path(x), required=True)
+
+    parser.add_argument('--seed', default=666, type=int)
+    parser.add_argument('--load_weights_from', type=lambda x: Path(x))
+    parser.add_argument('--save_ckpt_every', default=10, type=int)
+    parser.add_argument('--wandb', action='store_true', default=False)
+    parser.add_argument('--mode', choices=['train', 'eval', 'test', 'predict'], default='train', type=str)
+    parser.add_argument('--resume_from', default=None, type=str)
+
+    if datamodule is not None:
+        parser.add_datamodule(datamodule)
+
+    if loss is not None:
+        parser.add_lossmodule(loss)
+
+    if routine is not None:
+        parser.add_routine(routine)
+
+    if viz is not None:
+        parser.add_class_arguments(viz, 'viz', skip={'out_dir', 'exist_ok'})
+
+    # bindings between modules (data/routine/loss)
+    parser.link_arguments('data.batch_size', 'routine.batch_size')
+
+    parser.add_logger(WandbLogger)
+    parser.add_trainer()
+
+    args = parser.parse_args()
+    return args

capture/utils/exp.py

@@ -0,0 +1,77 @@
+import os
+
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning import Trainer as plTrainer
+from pathlib import Path
+
+from ..callbacks import VisualizeCallback, MetricLogging
+from ..data.module import DataModule
+from .log import get_info
+
+
+def get_data(args=None, **kwargs):
+    if args is None:
+        return DataModule(**kwargs)
+    else:
+        return DataModule(**args.data)
+
+def get_name(args) -> str:
+    name = ''#f'{args.mode}'
+
+    src_ds_verbose = str(args.data.source_list).split(os.sep)[-1].replace("_","-").upper()
+
+    if args.mode == 'train':
+        if args.loss.use_source and not args.loss.use_target:
+            name += f'pretrain_ds{args.data.source_ds.upper()}_lr{args.routine.lr}_x{args.data.input_size}_bs{args.data.batch_size}_reg{args.loss.reg_weight}_rend{args.loss.render_weight}_ds{str(args.data.source_list).split(os.sep)[-1].replace("_","-").upper()}'
+
+        elif args.loss.use_target:
+            name += f'_F_{args.data.target_ds.upper()}_lr{args.routine.lr}_x{args.data.input_size}_bs{args.data.tgt_bs}_aug{int(args.data.transform)}_reg{args.loss.pl_reg_weight}_rend{args.loss.pl_render_weight}_ds{str(args.data.target_list).split(os.sep)[-1].replace("_","-").upper()}'
+        else:
+            name += f'_T_{args.data.source_ds.upper()}_lr{args.routine.lr}_x{args.data.input_size}_aug{int(args.data.transform)}_reg{args.loss.render_weight}_rend{args.loss.render_weight}_ds{str(args.data.source_list).split(os.sep)[-1].replace("_","-").upper()}'
+
+        #if args.loss.adv_weight:
+        #    name += f'_ADV{args.loss.adv_weight}'
+        if args.data.source_ds == 'acg':
+            name += f'_mixbs{args.data.batch_size}'
+        if args.loss.reg_weight != 0.1:
+            name += f'_regSRC{args.loss.reg_weight}'
+        if args.loss.render_weight != 1:
+            name += f'_rendSRC{args.loss.render_weight}'
+        #if args.data.use_ref:
+        #    name += '_useRef'
+        if args.load_weights_from:
+            wname, epoch = get_info(str(args.load_weights_from))
+            assert wname and epoch
+            name += f'_init{wname.replace("_", "-")}-{epoch}ep'
+
+    name += f'_s{args.seed}'
+    return name
+    # name += args.load_weights_from.split(os.sep)[-1][:-5]
+
+def get_callbacks(args):
+    callbacks = [
+        VisualizeCallback(out_dir=args.out_dir, exist_ok=bool(args.resume_from), **args.viz),
+        ModelCheckpoint(
+            dirpath=args.out_dir/'ckpt_1',
+            filename='{name}_{epoch}-{step}',
+            save_weights_only=False,
+            save_top_k=-1,
+            every_n_epochs=args.save_ckpt_every),
+        MetricLogging(args.load_weights_from, args.data.test_list, outdir=Path('./logs')),
+    ]
+    return callbacks
+
+class Trainer(plTrainer):
+    def __init__(self, o_args, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ckpt_path = o_args.resume_from
+
+    def __call__(self, mode, module, data) -> None:
+        if mode == 'test':
+            self.test(module, data)
+        elif mode == 'eval':
+            self.validate(module, data)
+        elif mode == 'predict':
+            self.predict(module, data)
+        elif mode == 'train':
+            self.fit(module, data, ckpt_path=self.ckpt_path)

capture/utils/log.py

@@ -0,0 +1,47 @@
+import csv, re, os
+
+from pytorch_lightning.loggers import TensorBoardLogger
+
+
+def read_csv(fname):
+    with open(fname, 'r') as f:
+        reader = csv.DictReader(f)
+        return list(reader)
+
+def append_csv(fname, dicts):
+    if isinstance(dicts, dict):
+        dicts = [dicts]
+
+    if os.path.isfile(fname):
+        dicts = read_csv(fname) + dicts
+
+    write_csv(fname, dicts)
+
+def write_csv(fname, dicts):
+    assert len(dicts) > 0
+    with open(fname, 'w', newline='') as f:
+        writer = csv.DictWriter(f, fieldnames=dicts[0].keys())
+        writer.writeheader()
+        for d in dicts:
+            writer.writerow(d)
+
+def now():
+    from datetime import datetime
+    return datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+
+def get_info(weights: str):
+    search = re.search(r"(.*)_epoch=(\d+)-step", weights)
+    if search:
+        name, epoch = search.groups()
+        return str(name).split(os.sep)[-1], str(epoch)
+    return None, None
+
+def get_matlist(cache_dir, dir):
+    with open(cache_dir, 'r') as f:
+        content = f.readlines()
+    files = [dir/f.strip() for f in content]
+    return files
+
+def get_logger(args):
+    logger = TensorBoardLogger(save_dir=args.out_dir)
+    logger.log_hyperparams(args)

capture/utils/model.py

@@ -0,0 +1,43 @@
+import torch.nn as nn
+
+from pathlib import Path
+
+from ..source import ResnetEncoder, MultiHeadDecoder, DenseMTL, DenseReg, Vanilla
+
+
+def replace_batchnorm_(module: nn.Module):
+    for name, child in module.named_children():
+        if isinstance(child, nn.BatchNorm2d):
+            setattr(module, name, nn.InstanceNorm2d(child.num_features))
+        else:
+            replace_batchnorm_(child)
+
+def get_model(archi):
+    assert archi == 'densemtl'
+
+    encoder = ResnetEncoder(num_layers=101, pretrained=True, in_channels=3)
+    decoder = MultiHeadDecoder(
+        num_ch_enc=encoder.num_ch_enc,
+        tasks=dict(albedo=3, roughness=1, normals=2),
+        return_feats=False,
+        use_skips=True)
+
+    model = nn.Sequential(encoder, decoder)
+    replace_batchnorm_(model)
+    return model
+
+def get_module(args):
+    loss = DenseReg(**args.loss)
+    model = get_model(args.archi)
+
+    weights = args.load_weights_from
+    if weights:
+        assert weights.is_file()
+        return Vanilla.load_from_checkpoint(str(weights), model=model, loss=loss, strict=False, **args.routine)
+
+    return Vanilla(model, loss, **args.routine)
+
+def get_inference_module(pt):
+    assert Path(pt).exists()
+    model = get_model('densemtl')
+    return Vanilla.load_from_checkpoint(str(pt), model=model, strict=False)

fabric_diffusion.py

@@ -0,0 +1,139 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import os
+from diffusers import StableDiffusionInstructPix2PixPipeline
+from PIL import Image
+import random
+
+SEED=42
+torch.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+np.random.seed(SEED)
+random.seed(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+generator=torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(SEED)
+
+class FabricDiffusionPipeline():
+    def __init__(self, device, texture_checkpoint, print_checkpoint):
+        
+        self.device = device
+        self.texture_checkpoint = texture_checkpoint
+        self.print_base_model = print_checkpoint
+
+        if texture_checkpoint:
+            self.texture_model = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+                texture_checkpoint,
+                torch_dtype=torch.float16,
+                safety_checker=None
+            )
+            # with open(os.path.join(texture_checkpoint, "unet", "diffusion_pytorch_model.safetensors"), "rb") as f:
+            #     data = f.read()
+            # loaded = load(data)
+            # self.texture_pipeline.unet.load_state_dict(loaded)
+            self.texture_model = self.texture_model.to(device)
+        else:
+            self.texture_model = None
+
+        # set circular convolution for the texture model
+        if self.texture_model:
+            for a, b in self.texture_model.unet.named_modules():
+                if isinstance(b, nn.Conv2d):
+                    setattr(b, 'padding_mode', 'circular')
+            for a, b in self.texture_model.vae.named_modules():
+                if isinstance(b, nn.Conv2d):
+                    setattr(b, 'padding_mode', 'circular')
+
+        if print_checkpoint:
+            self.print_model = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+                print_checkpoint,
+                torch_dtype=torch.float16,
+                safety_checker=None
+            )
+            self.print_model = self.print_model.to(device)
+        else:
+            self.print_model = None
+
+    def load_real_data_with_mask(self, dataset_path, image_name):
+        image = np.array(Image.open(os.path.join(dataset_path, 'images', image_name)).convert('RGB'))
+        seg_mask = np.array(Image.open(os.path.join(dataset_path, 'seg_mask', image_name)).convert('L'))[..., None]
+        texture_mask = np.array(Image.open(os.path.join(dataset_path, 'texture_mask', image_name)).convert('L'))[
+            ..., None]
+        # crop the image based on texture_mask
+        x1, y1, x2, y2 = np.where(texture_mask > 0)[1].min(), np.where(texture_mask > 0)[0].min(), \
+        np.where(texture_mask > 0)[1].max(), np.where(texture_mask > 0)[0].max()
+        texture_patch = image[y1:y2, x1:x2]
+        # resize the texture_patch to 256x256
+        texture_patch = Image.fromarray(texture_patch.astype(np.uint8)).resize((256, 256))
+
+        return image, seg_mask, texture_patch
+
+    def load_patch_data(self, patch_path):
+        texture_patch = Image.open(patch_path).convert('RGB').resize((256, 256))
+        return texture_patch
+
+    def flatten_texture(self, texture_patch, n_samples=3, use_inversion=True):
+        num_inference_steps = 20
+        self.texture_model.scheduler.set_timesteps(num_inference_steps)
+        timesteps = self.texture_model.scheduler.timesteps
+
+        # convert image to latent using vae
+        image = self.texture_model.image_processor.preprocess(texture_patch)
+        if use_inversion:
+            image_latents = self.texture_model.prepare_image_latents(image, batch_size=1,
+                                                                     num_images_per_prompt=1,
+                                                                     device=self.device,
+                                                                     dtype=torch.float16,
+                                                                     do_classifier_free_guidance=False)
+
+            image_latents = (image_latents - torch.mean(image_latents)) / torch.std(image_latents)
+
+            # forward noising process
+
+            noise = torch.randn_like(image_latents)
+            noisy_image_latents = self.texture_model.scheduler.add_noise(image_latents, noise, timesteps[0:1])
+
+            noisy_image_latents /= self.texture_model.scheduler.init_noise_sigma
+            noisy_image_latents = torch.tile(noisy_image_latents, (n_samples, 1, 1, 1))
+        else:
+            noisy_image_latents = None
+
+        image = torch.tile(image, (n_samples, 1, 1, 1))
+        gen_imgs = self.texture_model(
+            "",
+            image=image,
+            num_inference_steps=20,
+            image_guidance_scale=1.5,
+            guidance_scale=7.,
+            latents=noisy_image_latents,
+            num_images_per_prompt=n_samples,
+            generator=generator
+        ).images
+
+        return gen_imgs
+
+    def flatten_print(self, print_patch, n_samples=3):
+        image = self.print_model.image_processor.preprocess(print_patch)
+        gen_imgs = []
+        for i in range(n_samples):
+            gen_img = self.print_model(
+                "",
+                image=image,
+                num_inference_steps=20,
+                image_guidance_scale=1.5,
+                guidance_scale=7.,
+                generator=generator
+            ).images[0]
+            gen_img = np.asarray(gen_img) / 255.
+            alpha_map = np.clip(gen_img / 0.1 * 1.2 - 0.2, 0., 1).mean(axis=-1, keepdims=True)
+            gen_img = np.clip((gen_img - 0.1) / 0.9, 0., 1.)
+            gen_img = np.concatenate([gen_img, alpha_map], axis=-1)
+            gen_img = (gen_img * 255).astype(np.uint8)
+            gen_img = Image.fromarray(gen_img)
+            gen_imgs.append(gen_img)
+
+        return gen_imgs
+

requirements.txt

@@ -0,0 +1,22 @@
+accelerate==1.3.0
+bitsandbytes==0.45.2
+clearml==1.17.1
+datasets==3.2.0
+diffusers==0.19.3
+easydict==1.13
+gradio==5.16.0
+huggingface_hub==0.25.2
+jsonargparse==4.36.0
+numpy==1.23.5
+opencv_python==4.11.0.86
+peft==0.5.0
+Pillow==11.1.0
+pytorch_lightning==2.4.0
+PyYAML==6.0.2
+Requests==2.32.3
+torch==2.2.0+cu121
+torchmetrics==1.6.1
+torchvision==0.17.0+cu121
+tqdm==4.67.1
+transformers==4.48.1
+wandb==0.19.6