Upload 115 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+PixHtLab-Src/Demo/PixhtLab/Examples/c82c09fc01e84282bc8870c263dcf81b_bg.jpg filter=lfs diff=lfs merge=lfs -text

PixHtLab-Src/.gitignore

@@ -0,0 +1,5 @@
+.DS_Store
+.idea
+*.log
+tmp/
+*__pycache__*

PixHtLab-Src/Data/data_prepare.py

@@ -0,0 +1,21 @@
+import objaverse
+import multiprocessing
+import random
+
+processes = 8
+
+random.seed(0)
+uids = objaverse.load_uids()
+random_object_uids = random.sample(uids, 100)
+objects = objaverse.load_objects(
+    uids=random_object_uids,
+    download_processes=processes
+)
+
+
+objects = objaverse.load_objects(
+    uids=random_object_uids,
+    download_processes=processes
+)
+
+import pdb; pdb.set_trace()

PixHtLab-Src/Demo/PixhtLab/Examples/009_depth.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17338c5be1e9aa07c8d68ed26076af44f433374a92a90130b3f803d66f8d2442
+size 1048704

PixHtLab-Src/Demo/PixhtLab/Examples/009_depth_valid_mask.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85f888fa80c5da36acf97938cda8f55323758aefda263fc859ab83b2955c92c1
+size 262272

PixHtLab-Src/Demo/PixhtLab/Examples/009_pixht_new.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25dfedce61e31739fec2164898d9975ab9224c02aea1540029eb7b04d0505118
+size 2097280

PixHtLab-Src/Demo/PixhtLab/Examples/010_depth.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:872277399788bc63b43a1fc63e817d8c12d18e6960b6d7e51a41d2ea2fae21f6
+size 1048704

PixHtLab-Src/Demo/PixhtLab/Examples/010_depth_valid_mask.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ceb9b34576604a240e557f2bb2fa87b7328e30a28a44357d799c6a3271c6bf6
+size 262272

PixHtLab-Src/Demo/PixhtLab/Examples/010_pixht_new.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb2031249ece53825e983e62e1a695c53bdfcd5ae12d3a1350dff80e710f0aa3
+size 2097280

PixHtLab-Src/Demo/PixhtLab/Examples/011_depth.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfdabce8a328c02e7e52d156a2f2f03e1302fb695e9e70f35355d1aced7de24f
+size 1048704

PixHtLab-Src/Demo/PixhtLab/Examples/011_depth_valid_mask.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12eaae90b70229aa83450d8ea63fd77c7e8fdc6240e473aafab77021c2891efe
+size 262272

PixHtLab-Src/Demo/PixhtLab/Examples/011_pixht_new.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ef1a1f21f4ac53dc680fc6a965cce4f8043bd313a89e52ef733c5986be16595
+size 2097280

PixHtLab-Src/Demo/PixhtLab/GSSN/inference_shadow.py

@@ -0,0 +1,70 @@
+import os
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as T
+import argparse
+import time
+from tqdm import tqdm
+import numpy as np
+import os
+from os.path import join
+
+import math
+
+import cv2
+import random
+
+import sys
+# sys.path.insert(0, '../../Training/app/models')
+sys.path.insert(0, '/home/ysheng/Documents/Research/GSSN/Training/app/models')
+
+from SSN_v1 import SSN_v1
+from SSN import SSN
+
+
+class SSN_Infernece():
+    def __init__(self, ckpt, device=torch.device('cuda:0')):
+        self.device = device
+        self.model  = SSN(3, 1,  mid_act='gelu', out_act='null', resnet=False)
+
+        weight = torch.load(ckpt)
+        self.model.to(device)
+        self.model.load_state_dict(weight['model'])
+
+        # inference related
+        BINs    = 100
+        MAX_RAD = 20
+        self.size_interval     = MAX_RAD / BINs
+        self.soft_distribution = [[np.exp(-0.2 * (i - j) ** 2) for i in np.arange(BINs)] for j in np.arange(BINs)]
+
+
+    def render_ss(self, input_np, softness):
+        """ input_np:
+              H x W x C
+        """
+        input_tensor = torch.tensor(input_np.transpose((2, 0, 1)))[None, ...].float().to(self.device)
+        transform    = T.Resize((256, 256))
+
+        c = input_tensor.shape[1]
+        # for i in range(c):
+        #     print(input_tensor[:, i].min(), input_tensor[:, i].max())
+
+        # print('softness: ', softness)
+        l = torch.from_numpy(np.array(self.soft_distribution[int(softness/self.size_interval)]).astype(np.float32)).unsqueeze(dim=0).to(self.device)
+
+        input_tensor  = transform(input_tensor)
+        output_tensor = self.model(input_tensor, l)
+        output_np     = output_tensor[0].detach().cpu().numpy().transpose((1,2,0))
+
+        return output_np
+
+
+if __name__ == '__main__':
+    model = SSN_Infernece('weights/0000000700.pt')

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/inference_shadow.py

@@ -0,0 +1,70 @@
+import os
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+import argparse
+import time
+from tqdm import tqdm
+import numpy as np
+import os
+from os.path import join
+
+import math
+
+import cv2
+import random
+from .ssn.ssn import Relight_SSN
+device = torch.device('cuda:0')
+
+def net_render_np(model, mask_np, hard_shadow_np, size, orientation):
+    """
+    input:
+        mask_np shape: b x c x h x w
+        ibl_np shape: 1 x 16 x 32
+    output:
+        shadow_predict shape: b x c x h x w
+    """
+
+    size_interval = 0.5 / 100
+    ori_interval = np.pi / 100
+
+    soft_distribution = [[np.exp(-0.2 * (i - j) ** 2) for i in np.arange(0.5 / size_interval)]
+                         for j in np.arange(0.5 / size_interval)]
+
+    # print('mask_np: {}, hard_shadow_np: {}'.format(mask_np.shape, hard_shadow_np.shape))
+    s = time.time()
+    if mask_np.dtype == np.uint8:
+        mask_np = mask_np / 255.0
+
+    mask, h_shadow = torch.Tensor(mask_np), torch.Tensor(hard_shadow_np)
+    size_soft = torch.Tensor(np.array(soft_distribution[int(size / size_interval)])).unsqueeze(0)
+    ori_soft = torch.Tensor(np.array(soft_distribution[int(orientation / ori_interval)])).unsqueeze(0)
+
+    with torch.no_grad():
+        I_m, I_h, size_t, ori = mask.to(device), h_shadow.to(device), size_soft.to(device), ori_soft.to(device)
+        # print('I_m: {}, I_h: {}'.format(I_m.shape, I_h.shape))
+        predicted_img = model(I_h, I_m, size_t, ori)
+
+    # print('net predict finished, time: {}s'.format(time.time() - s))
+
+    return predicted_img.detach().cpu().numpy()
+
+def init_models(ckpt):
+    baseline_model = Relight_SSN(1, 1, is_training=False)
+    baseline_checkpoint = torch.load(ckpt)
+    baseline_model.to(device)
+    baseline_model.load_state_dict(baseline_checkpoint['model_state_dict'])
+    return baseline_model
+
+if __name__ == '__main__':
+    softness = [0.02, 0.2, 0.3, 0.4]
+    model = init_models('weights/human_baseline123.pt')
+    mask, hard_shadow, size, orientation = np.random.randn(1,1,256,256), np.random.randn(1,1,256,256), softness[0], 0 
+    shadow = net_render_np(model, mask, hard_shadow, size, orientation)

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/params.py

@@ -0,0 +1,92 @@
+import argparse
+
+class params():
+    """ Singleton class for doing experiments """
+    
+    class __params():
+        def __init__(self):
+            self.norm = 'group_norm'
+            self.prelu = False
+            self.weight_decay = 5e-4
+            self.small_ds = False
+            self.multi_gpu = False
+            self.log = False
+            self.input_channel = 1
+            self.vis_port = 8002
+            self.cpu = False
+            self.pred_touch = False
+            self.tbaseline = False
+            self.touch_loss = False
+            self.input_channel = 1
+
+        def set_params(self, options):
+            self.options = options
+            self.norm = options.norm
+            self.prelu = options.prelu
+            self.weight_decay = options.weight_decay
+            self.small_ds = options.small_ds
+            self.multi_gpu = options.multi_gpu
+            self.log = options.log
+            self.input_channel = options.input_channel
+            self.vis_port = options.vis_port
+            self.cpu = options.cpu
+            self.ds_folder = options.ds_folder
+            self.pred_touch = options.pred_touch
+            self.tbaseline = options.tbaseline
+            self.touch_loss = options.touch_loss
+            
+        def __str__(self):
+            return 'norm: {}  prelu: {} weight decay: {} small ds: {}'.format(self.norm, self.prelu, self.weight_decay, self.small_ds)
+
+    # private static variable
+    param_instance = None
+    
+    def __init__(self):
+        if not params.param_instance:
+            params.param_instance = params.__params()
+    
+    def get_params(self):
+        return params.param_instance
+    
+    def set_params(self, options):
+        params.param_instance.set_params(options)
+
+def parse_params():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--workers', type=int, help='number of data loading workers', default=16)
+    parser.add_argument('--batch_size', type=int, default=28, help='input batch size during training')
+    parser.add_argument('--epochs', type=int, default=10000, help='number of epochs to train for')
+    parser.add_argument('--lr', type=float, default=0.003, help='learning rate, default=0.005')
+    parser.add_argument('--beta1', type=float, default=0.9, help='momentum for SGD, default=0.9')
+    parser.add_argument('--resume', action='store_true', help='resume training')
+    parser.add_argument('--relearn', action='store_true', help='forget previous best validation loss')
+    parser.add_argument('--weight_file',type=str,  help='weight file')
+    parser.add_argument('--multi_gpu', action='store_true', help='use multiple GPU training')
+    parser.add_argument('--timers', type=int, default=1, help='number of epochs to train for')
+    parser.add_argument('--use_schedule', action='store_true',help='use automatic schedule')
+    parser.add_argument('--patience', type=int, default=2, help='use automatic schedule')
+    parser.add_argument('--exp_name', type=str, default='l1 loss',help='experiment name')    
+    parser.add_argument('--norm', type=str, default='group_norm', help='use group norm')
+    parser.add_argument('--ds_folder', type=str, default='./dataset/general_dataset', help='Dataset folder')
+    parser.add_argument('--hd_dir', type=str, default='/mnt/yifan/data/Adobe/HD_styleshadow/', help='Dataset folder')
+    parser.add_argument('--prelu', action='store_true', help='use prelu')
+    parser.add_argument('--small_ds', action='store_true', help='small dataset')
+    parser.add_argument('--log', action='store_true', help='log information')
+    parser.add_argument('--vis_port', default=8002,type=int, help='visdom port')
+    parser.add_argument('--weight_decay', type=float, default=4e-5, help='weight decay for model weight')
+    parser.add_argument('--save', action='store_true', help='save batch results?')
+    parser.add_argument('--cpu', action='store_true', help='Force training on CPU')
+    parser.add_argument('--pred_touch', action='store_true', help='Use touching surface')
+    parser.add_argument('--input_channel', type=int, default=1, help='how many input channels')
+    
+    # based on baseline method, for fine tuning
+    parser.add_argument('--from_baseline', action='store_true', help='training from baseline')
+    parser.add_argument('--tbaseline', action='store_true', help='T-baseline, input two channels')
+    parser.add_argument('--touch_loss', action='store_true', help='Use touching loss')
+    
+
+    arguments = parser.parse_args()
+    parameter = params()
+    parameter.set_params(arguments)
+    
+    return arguments

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/perturb_touch.py

@@ -0,0 +1,24 @@
+import numpy as np
+import cv2
+import random
+
+random.seed(19920208)
+
+def random_kernel():
+    ksize = random.randint(1,3)
+    kernel = np.ones((ksize, ksize))
+    return kernel
+    
+def random_perturb(img):
+    return img
+#     perturbed = img.copy()
+#     if random.random() < 0.5:
+#         perturbed = cv2.erode(perturbed, random_kernel(), iterations = 1)
+    
+#     if random.random() < 0.5:
+#         perturbed = cv2.dilate(perturbed, random_kernel(), iterations = 1)
+    
+#     cv2.normalize(perturbed, perturbed, 0.0,1.0, cv2.NORM_MINMAX)
+#     if len(perturbed.shape) == 2:
+#         perturbed = perturbed[:,:,np.newaxis]
+#     return perturbed

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/random_pattern.py

@@ -0,0 +1,92 @@
+import random
+import time
+import numbergen as ng
+import imagen as ig
+import numpy as np
+import cv2
+from param.parameterized import get_logger
+import logging
+
+get_logger().setLevel(logging.ERROR)
+
+class random_pattern():
+    def __init__(self, maximum_blob=50):
+        #         self.generator_list = []
+
+        #         start = time.time()
+        #         for i in range(maximum_blob):
+        #             self.generator_list.append(ig.Gaussian(size=))
+        #         print('random pattern init time: {}s'.format(time.time()-start))
+
+        pass
+
+    def y_transform(self, y):
+        # y = []
+        pass
+
+    def get_pattern(self, w, h, x_density=512, y_density=128, num=50, scale=3.0, size=0.1, energy=3500,
+                    mitsuba=False, seed=None, dataset=False):
+        if seed is None:
+            seed = random.randint(0, 19920208)
+        else:
+            seed = seed + int(time.time())
+
+        if num == 0:
+            ibl = np.zeros((y_density, x_density))
+            orientation = np.pi * ng.UniformRandom(seed=seed + 3)()
+        else:
+            y_fact = y_density / 256
+            num = 1
+            size = size * ng.UniformRandom(seed=seed + 4)()
+            orientation = np.pi * ng.UniformRandom(seed=seed + 3)()
+            gs = ig.Composite(operator=np.add,
+                              generators=[ig.Gaussian(
+                                  size=size,
+                                  scale=1.0,
+                                  x=ng.UniformRandom(seed=seed + i + 1) - 0.5,
+                                  y=((1.0 - ng.UniformRandom(seed=seed + i + 2) * y_fact) - 0.5),
+                                  aspect_ratio=0.7,
+                                  orientation=orientation,
+                              ) for i in range(num)],
+                              position=(0, 0),
+                              xdensity=512)
+
+            # gs = ig.Composite(operator=np.add,
+            #                   generators=[ig.Gaussian(
+            #                       size=size * ng.UniformRandom(seed=seed + i + 4),
+            #                       scale=scale * (ng.UniformRandom(seed=seed + i + 5) + 1e-3),
+            #                       x=int(ind / h),
+            #                       y=ind % h,
+            #                       aspect_ratio=0.7,
+            #                       orientation=np.pi * ng.UniformRandom(seed=seed + i + 3),
+            #                   ) for i in range(num)],
+            #                   position=(0, 0),
+            #                   xdensity=512)
+            ibl = gs()[:y_density, :]
+
+        # prepare to fix energy inconsistent
+        if dataset:
+            ibl = self.to_dataset(ibl, w, h)
+
+        if mitsuba:
+            return ibl, size, orientation
+        else:
+            return ibl, size, orientation
+
+    def to_mts_ibl(self, ibl):
+        """ Input: 256 x 512 pattern generated ibl 
+            Output: the ibl in mitsuba ibl
+        """
+        return np.repeat(ibl[:, :, np.newaxis], 3, axis=2)
+
+    def normalize(self, ibl, energy=30.0):
+        total_energy = np.sum(ibl)
+        if total_energy < 1e-3:
+            print('small energy: ', total_energy)
+            h, w = ibl.shape
+            return np.zeros((h, w))
+
+        return ibl * energy / total_energy
+
+    def to_dataset(self, ibl, w, h):
+        return self.normalize(cv2.flip(cv2.resize(ibl, (w, h)), 0), 30)

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/ssn.py

@@ -0,0 +1,146 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .ssn_submodule import Conv, Up, Up_Stream, get_layer_info, add_coords
+import copy
+
+class Relight_SSN(nn.Module):
+    """ Implementation of Relighting Net """
+
+    def __init__(self, n_channels=3, out_channels=3, is_training=True, activation_func = 'relu'):
+        super(Relight_SSN, self).__init__()
+        self.is_training = is_training 
+
+        norm_layer1, activation_func1 = get_layer_info(16, activation_func)
+        norm_layer2, activation_func2 = get_layer_info(16, activation_func)
+        if norm_layer1 is not None:
+            self.in_conv1 = nn.Sequential(
+                nn.Conv2d(n_channels, 16, kernel_size=7, padding=3, bias=True),
+                norm_layer1,
+                activation_func1
+            )
+        elif norm_layer1 is None:
+            self.in_conv1 = nn.Sequential(
+                nn.Conv2d(n_channels, 16, kernel_size=7, padding=3, bias=True),
+                activation_func1
+            )
+
+        if norm_layer2 is not None:
+            self.in_conv2 = nn.Sequential(
+                nn.Conv2d(n_channels, 16, kernel_size=7, padding=3, bias=True),
+                norm_layer2,
+                activation_func2
+            )
+        elif norm_layer2 is None:
+            self.in_conv2 = nn.Sequential(
+                nn.Conv2d(n_channels, 16, kernel_size=7, padding=3, bias=True),
+                activation_func2
+            )
+
+        self.down_256_128 = Conv(32, 64, conv_stride=2)
+        self.down_128_128 = Conv(64, 64, conv_stride=1)
+        self.down_128_64 = Conv(64, 128, conv_stride=2)
+        self.down_64_64 = Conv(128, 128, conv_stride=1)
+        self.down_64_32 = Conv(128, 256, conv_stride=2)
+        self.down_32_32 = Conv(256, 256, conv_stride=1)
+        self.down_32_16 = Conv(256, 512, conv_stride=2)
+        self.down_16_16_1 = Conv(512, 512, conv_stride=1)
+        self.down_16_16_2 = Conv(512, 512, conv_stride=1)
+        self.down_16_16_3 = Conv(512, 512, conv_stride=1)
+        self.to_bottleneck = Conv(512, 2, conv_stride=1)
+
+        self.up_stream = Up_Stream(out_channels)
+
+    """
+        Input is (source image, target light, source light, )
+        Output is: predicted new image, predicted source light, self-supervision image
+    """
+
+    def forward(self, I_h, I_m, size, angle):
+        if self.is_training:
+            bs, fake_bs, dim = size.shape
+            bs, fake_bs, ch, w, h = I_h.shape
+            size = size.reshape(bs * fake_bs, dim)
+            angle = angle.reshape(bs * fake_bs, dim)
+            I_h = I_h.reshape(bs * fake_bs, ch, w, h)
+            I_m = I_m.reshape(bs * fake_bs, ch, w, h)
+        else:
+            size = size
+            angle = angle
+            I_h = I_h
+            I_m = I_m
+        style = torch.cat((size, angle), dim=1)
+        x1 = self.in_conv1(I_m)  # 29 x 256 x 256
+        x2 = self.in_conv2(I_h)
+
+        x1 = torch.cat((x1, x2), dim=1)  # 32 x 256 x 256
+
+        x2 = self.down_256_128(x1, x1)  # 64 x 128 x 128
+
+        x3 = self.down_128_128(x2, x1)  # 64 x 128 x 128
+
+        x4 = self.down_128_64(x3, x1)  # 128 x 64 x 64
+
+        x5 = self.down_64_64(x4, x1)  # 128 x 64 x 64
+
+        x6 = self.down_64_32(x5, x1)  # 256 x 32 x 32
+
+        x7 = self.down_32_32(x6, x1)  # 256 x 32 x 32
+
+        x8 = self.down_32_16(x7, x1)  # 512 x 16 x 16
+
+        x9 = self.down_16_16_1(x8, x1)  # 512 x 16 x 16
+
+        x10 = self.down_16_16_2(x9, x1)  # 512 x 16 x 16
+
+        x11 = self.down_16_16_3(x10, x1)  # 512 x 16 x 16
+        ty = self.up_stream(x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, style)
+
+        return ty
+
+
+def baseline_2_tbaseline(model):
+    """ change input layer to be two channels
+    """
+    input_channel = 2
+    tbase_inlayer = nn.Sequential(
+        nn.Conv2d(input_channel, 32 - input_channel, kernel_size=7, padding=3, bias=True),
+        nn.GroupNorm(1, 32 - input_channel),
+        nn.ReLU()
+    )
+    model.in_conv = tbase_inlayer
+    return model
+
+
+def baseline_2_touchloss(model):
+    """ change output layer to be two channels
+    """
+    touchless_outlayer = nn.Sequential(
+        nn.Conv2d(64, 2, stride=1, kernel_size=3, padding=1, bias=True),
+        nn.GroupNorm(1, 2),
+        nn.ReLU()
+    )
+    model.up_stream.out_conv = touchless_outlayer
+    return model
+
+
+if __name__ == '__main__':
+    mask_test, touch_test = torch.zeros((1, 1, 256, 256)), torch.zeros((1, 1, 256, 256))
+    ibl = torch.zeros((1, 1, 16, 32))
+
+    I_s = mask_test
+    baseline = Relight_SSN(1, 1)
+    baseline_output, _ = baseline(I_s, ibl)
+
+    tbaseline = baseline_2_tbaseline(copy.deepcopy(baseline))
+    I_s = torch.cat((mask_test, touch_test), axis=1)
+    tbaseline_output, _ = tbaseline(I_s, ibl)
+
+    t_loss_baseline = baseline_2_touchloss(copy.deepcopy(baseline))
+    I_s = mask_test
+    tloss_output, _ = t_loss_baseline(I_s, ibl)
+
+    print('baseline output: ', baseline_output.shape)
+    print('tbaseline output: ', tbaseline_output.shape)
+    print('tloss output: ', tloss_output.shape)

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/ssn_dataset.py

@@ -0,0 +1,290 @@
+import sys
+
+sys.path.append("..")
+
+import os
+from os.path import join
+import torch
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, utils
+import time
+import random
+# import matplotlib.pyplot as plt
+import cv2
+from params import params
+from .random_pattern import random_pattern
+from .perturb_touch import random_perturb
+
+
+class ToTensor(object):
+    """Convert ndarrays in sample to Tensors."""
+
+    def __call__(self, img, is_transpose=True):
+        # swap color axis because
+        # numpy image: H x W x C
+        # torch image: C X H X W
+        if is_transpose:
+            img = img.transpose((0, 3, 1, 2))
+        return torch.Tensor(img)
+
+
+class SSN_Dataset(Dataset):
+    def __init__(self, ds_dir, hd_dir, is_training, fake_batch_size=8):
+        start = time.time()
+        self.fake_batch_size = fake_batch_size
+        # # of samples in each group
+        # magic number here
+        self.ibl_group_size = 16
+
+        parameter = params().get_params()
+
+        # (shadow_path, mask_path)
+        self.meta_data = self.init_meta(ds_dir, hd_dir)
+
+        self.is_training = is_training
+        self.to_tensor = ToTensor()
+
+        end = time.time()
+        print("Dataset initialize spent: {} ms".format(end - start))
+
+        # fake random
+        np.random.seed(19950220)
+        np.random.shuffle(self.meta_data)
+
+        self.valid_divide = 10
+        if parameter.small_ds:
+            self.meta_data = self.meta_data[:len(self.meta_data) // self.valid_divide]
+
+        self.training_num = len(self.meta_data) - len(self.meta_data) // self.valid_divide
+        print('training: {}, validation: {}'.format(self.training_num, len(self.meta_data) // self.valid_divide))
+
+        self.random_pattern_generator = random_pattern()
+
+        self.thread_id = os.getpid()
+        self.seed = os.getpid()
+        self.perturb = not parameter.pred_touch and not parameter.touch_loss
+        self.size_interval = 0.5 / 100
+        self.ori_interval = np.pi / 100
+
+        self.soft_distribution = [[np.exp(-0.4 * (i - j) ** 2) for i in np.arange(0.5 / self.size_interval)]
+                                  for j in np.arange(0.5 / self.size_interval)]
+
+    def __len__(self):
+        if self.is_training:
+            return self.training_num
+        else:
+            # return len(self.meta_data) - self.training_num
+            return len(self.meta_data) // self.valid_divide
+
+    def __getitem__(self, idx):
+        if self.is_training and idx > self.training_num:
+            print("error")
+        # offset to validation set
+        if not self.is_training:
+            idx = self.training_num + idx
+
+        cur_seed = idx * 1234 + os.getpid() + time.time()
+        random.seed(cur_seed)
+
+        # random ibls
+        shadow_path, mask_path, hard_path, touch_path = self.meta_data[idx]
+        hard_folder = hard_path.replace(hard_path.split('/')[-1], '')
+        if os.path.exists(hard_folder):
+            # hard_shadow = cv2.imread(hard_folder)
+
+            mask_img = cv2.imread(mask_path)
+            mask_img = mask_img[:, :, 0]
+            if mask_img.dtype == np.uint8:
+                mask_img = mask_img / 255.0
+            mask_img, shadow_bases = np.expand_dims(mask_img, axis=2), np.load(shadow_path)
+
+            w, h, c, m = shadow_bases.shape
+            shadow_soft_list = []
+            shadow_hard_list = []
+            size_list = []
+            orientation_list = []
+            mask_img_list = []
+            for i in range(int(self.fake_batch_size)):
+                shadow_img, light_img, size, orientation = self.render_new_shadow(shadow_bases)
+
+                h, w = mask_img.shape[0], mask_img.shape[1]
+                hi, wi = np.where(light_img == light_img.max())
+
+                while len(hi) > 1:
+                    shadow_img, light_img, size, orientation = self.render_new_shadow(shadow_bases)
+                    hi, wi = np.where(light_img == light_img.max())
+                size_soft = np.array(self.soft_distribution[int(size / self.size_interval)])
+                ori_soft = np.array(self.soft_distribution[int(orientation / self.ori_interval)])
+                prefix = '_ibli_' + str(int(wi * 8)) + '_iblj_' + str(int(hi * 8) + 128) + '_shadow.png'
+                shadow_hard_path = hard_path.replace('_shadow.png', prefix)
+                shadow_base = cv2.imread(shadow_hard_path, -1)[:, :, 0] / 255.0
+                shadow_base = np.expand_dims(shadow_base, axis=2)
+                shadow_base = self.line_aug(shadow_base)
+                shadow_soft_list.append(shadow_img)
+                shadow_hard_list.append(shadow_base)
+                size_list.append(size_soft)
+                orientation_list.append(ori_soft)
+                mask_img_list.append(mask_img)
+            shadow_softs = np.array(shadow_soft_list)
+            shadow_hards = np.array(shadow_hard_list)
+            sizes = np.array(size_list)
+            orientations = np.array(orientation_list)
+            mask_imgs = np.array(mask_img_list)
+
+            # touch_img = self.read_img(touch_path)
+            # touch_img = touch_img[:, :, 0:1]
+
+            #         if self.perturb:
+            #             touch_img = random_perturb(touch_img)
+
+            # input_img = np.concatenate((mask_img, touch_img), axis=2)
+            size = torch.Tensor(sizes)
+            ori = torch.Tensor(orientations)
+            hard_shadow, soft_shadow, mask_img = self.to_tensor(shadow_hards), self.to_tensor(
+                shadow_softs), self.to_tensor(
+                mask_imgs)
+            return {"hard_shadow": hard_shadow, "soft_shadow": soft_shadow, "mask_img": mask_img, "size": size,
+                    "angle": ori}
+        else:
+            mask_img = cv2.imread(mask_path)
+            mask_img = mask_img[:, :, 0]
+            if mask_img.dtype == np.uint8:
+                mask_img = mask_img / 255.0
+            mask_img, shadow_bases = np.expand_dims(mask_img, axis=2), 1.0 - np.load(shadow_path)
+
+            w, h, c, m = shadow_bases.shape
+            shadow_soft_list = []
+            shadow_hard_list = []
+            size_list = []
+            orientation_list = []
+            mask_img_list = []
+            for i in range(int(self.fake_batch_size)):
+                shadow_img, light_img, size, orientation = self.render_new_shadow(shadow_bases)
+
+                h, w = mask_img.shape[0], mask_img.shape[1]
+                hi, wi = np.where(light_img == light_img.max())
+
+                while len(hi) > 1:
+                    shadow_img, light_img, size, orientation = self.render_new_shadow(shadow_bases)
+                    hi, wi, _ = np.where(light_img == light_img[:, :, :].max())
+                size_soft = np.array(self.soft_distribution[int(size / self.size_interval)])
+                ori_soft = np.array(self.soft_distribution[int(orientation / self.ori_interval)])
+
+                shadow_base = shadow_bases[:, :, wi, hi]
+                shadow_base[shadow_base > 0.3] = 1
+                shadow_base[shadow_base < 0.4] = 0
+                shadow_base = self.line_aug(shadow_base)
+                mask_img = np.expand_dims(cv2.resize(mask_img, (512, 512)), axis=2)
+                shadow_base = np.expand_dims(cv2.resize(shadow_base, (512, 512)), axis=2)
+                shadow_img = np.expand_dims(cv2.resize(shadow_img, (512, 512)), axis=2)
+                shadow_soft_list.append(shadow_img)
+                shadow_hard_list.append(shadow_base)
+                size_list.append(size_soft)
+                orientation_list.append(ori_soft)
+                mask_img_list.append(mask_img)
+            shadow_softs = np.array(shadow_soft_list)
+            shadow_hards = np.array(shadow_hard_list)
+            sizes = np.array(size_list)
+            orientations = np.array(orientation_list)
+            mask_imgs = np.array(mask_img_list)
+
+            # touch_img = self.read_img(touch_path)
+            # touch_img = touch_img[:, :, 0:1]
+
+            #         if self.perturb:
+            #             touch_img = random_perturb(touch_img)
+
+            # input_img = np.concatenate((mask_img, touch_img), axis=2)
+            size = torch.Tensor(sizes)
+            ori = torch.Tensor(orientations)
+
+            hard_shadow, soft_shadow, mask_img = self.to_tensor(shadow_hards), self.to_tensor(
+                shadow_softs), self.to_tensor(
+                mask_imgs)
+
+            return {"hard_shadow": hard_shadow, "soft_shadow": soft_shadow, "mask_img": mask_img, "size": size,
+                    "angle": ori}
+
+    def init_meta(self, ds_dir, hd_dir):
+        metadata = []
+        # base_folder = join(ds_dir, 'base')
+        # mask_folder = join(ds_dir, 'mask')
+        # hard_folder = join(ds_dir, 'hard')
+        # touch_folder = join(ds_dir, 'touch')
+        # model_list = [f for f in os.listdir(base_folder) if os.path.isdir(join(base_folder, f))]
+        # for m in model_list:
+        #     shadow_folder, cur_mask_folder = join(base_folder, m), join(mask_folder, m)
+        #     shadows = [f for f in os.listdir(shadow_folder) if f.find('_shadow.npy') != -1]
+        #     for s in shadows:
+        #         prefix = s[:s.find('_shadow')]
+        #         metadata.append((join(shadow_folder, s),
+        #                          join(cur_mask_folder, prefix + '_mask.png'),
+        #                          join(join(hard_folder, m), prefix + '_shadow.png'),
+        #                          join(join(touch_folder, m), prefix + '_touch.png')))
+
+        base_folder = join(hd_dir, 'base')
+        mask_folder = join(hd_dir, 'mask')
+        hard_folder = join(hd_dir, 'hard')
+        touch_folder = join(hd_dir, 'touch')
+        model_list = [f for f in os.listdir(base_folder) if os.path.isdir(join(base_folder, f))]
+        for m in model_list:
+            shadow_folder, cur_mask_folder = join(base_folder, m), join(mask_folder, m)
+            shadows = [f for f in os.listdir(shadow_folder) if f.find('_shadow.npy') != -1]
+            for s in shadows:
+                prefix = s[:s.find('_shadow')]
+                metadata.append((join(shadow_folder, s),
+                                 join(cur_mask_folder, prefix + '_mask.png'),
+                                 join(join(hard_folder, m), prefix + '_shadow.png'),
+                                 join(join(touch_folder, m), prefix + '_touch.png')))
+
+        return metadata
+
+    def line_aug(self, shadow):
+        p = np.random.random()
+        if p > 0.6:
+            k = np.tan(min((np.random.random() + 0.000000001), 0.999) * np.pi - np.pi / 2)
+            x, y, c = shadow.shape
+            b_max = y - x * k
+            line_num = np.random.randint(1, 20)
+            b_list = np.random.random(line_num) * b_max
+            x_coord = np.tile(np.arange(shadow.shape[1])[None, :], (shadow.shape[0], 1))
+            y_coord = np.tile(np.arange(shadow.shape[0])[:, None], (1, shadow.shape[1]))
+
+            for b in b_list:
+                mask_res = y_coord - k * x_coord - b
+                shadow[np.abs(mask_res) < 1] = 0
+        return shadow
+
+    def get_prefix(self, path):
+        folder = os.path.dirname(path)
+        basename = os.path.basename(path)
+        return os.path.join(folder, basename[:basename.find('_')])
+
+    def render_new_shadow(self, shadow_bases):
+        shadow_bases = shadow_bases[:, :, :, :]
+        h, w, iw, ih = shadow_bases.shape
+
+        num = random.randint(0, 50)
+        pattern_img, size, orientation = self.random_pattern_generator.get_pattern(iw, ih, num=num, size=0.5,
+                                                                                   mitsuba=False)
+
+        # flip to mitsuba ibl
+        pattern_img = self.normalize_energy(cv2.flip(cv2.resize(pattern_img, (iw, ih)), 0))
+        shadow = np.tensordot(shadow_bases, pattern_img, axes=([2, 3], [1, 0]))
+        # pattern_img = np.expand_dims(cv2.resize(pattern_img, (iw, 16)), 2)
+
+        return np.expand_dims(shadow, 2), pattern_img, size, orientation
+
+    def get_min_max(self, batch_data, name):
+        print('{} min: {}, max: {}'.format(name, np.min(batch_data), np.max(batch_data)))
+
+    def log(self, log_info):
+        with open('log.txt', 'a+') as f:
+            f.write(log_info)
+
+    def normalize_energy(self, ibl, energy=30.0):
+        if np.sum(ibl) < 1e-3:
+            return ibl
+        return ibl * energy / np.sum(ibl)

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/ssn_submodule.py

@@ -0,0 +1,282 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def get_layer_info(out_channels, activation_func='relu'):
+    if out_channels >= 32:
+        group_num = 32
+    else:
+        group_num = 1
+
+    norm_layer = nn.GroupNorm(group_num, out_channels)
+
+    if activation_func == 'relu':
+        activation_func = nn.ReLU()
+    elif activation_func == 'prelu':
+        activation_func = nn.PReLU(out_channels)
+
+    return norm_layer, activation_func
+
+
+# add coord_conv
+class add_coords(nn.Module):
+    def __init__(self, use_cuda=True):
+        super(add_coords, self).__init__()
+        self.use_cuda = use_cuda
+
+    def forward(self, input_tensor):
+        b, c, dim_y, dim_x = input_tensor.shape
+        xx_ones = torch.ones([1, 1, 1, dim_x], dtype=torch.int32)
+        yy_ones = torch.ones([1, 1, 1, dim_y], dtype=torch.int32)
+
+        xx_range = torch.arange(dim_y, dtype=torch.int32)
+        yy_range = torch.arange(dim_x, dtype=torch.int32)
+        xx_range = xx_range[None, None, :, None]
+        yy_range = yy_range[None, None, :, None]
+
+        xx_channel = torch.matmul(xx_range, xx_ones)
+        yy_channel = torch.matmul(yy_range, yy_ones)
+
+        # transpose y
+        yy_channel = yy_channel.permute(0, 1, 3, 2)
+
+        xx_channel = xx_channel.float() / (dim_y - 1)
+        yy_channel = yy_channel.float() / (dim_x - 1)
+
+        xx_channel = xx_channel * 2 - 1
+        yy_channel = yy_channel * 2 - 1
+
+        xx_channel = xx_channel.repeat(b, 1, 1, 1)
+        yy_channel = yy_channel.repeat(b, 1, 1, 1)
+
+        if torch.cuda.is_available and self.use_cuda:
+            input_tensor = input_tensor.cuda()
+            xx_channel = xx_channel.cuda()
+            yy_channel = yy_channel.cuda()
+        out = torch.cat([input_tensor, xx_channel, yy_channel], dim=1)
+        return out
+
+
+class Conv(nn.Module):
+    """ (convolution => [BN] => ReLU) """
+
+    def __init__(self, in_channels, out_channels, kernel_size=3, conv_stride=1, padding=1, bias=True,
+                 activation_func='relu', style=False):
+        super().__init__()
+
+        self.style = style
+        norm_layer, activation_func = get_layer_info(out_channels, activation_func)
+        if style:
+            self.styleconv = Conv2DMod(in_channels, out_channels, kernel_size)
+            self.relu = nn.LeakyReLU(0.2, inplace=True)
+        else:
+            if norm_layer is not None:
+                self.conv = nn.Sequential(
+                    nn.Conv2d(in_channels, out_channels, stride=conv_stride, kernel_size=kernel_size, padding=padding,
+                              bias=bias),
+                    norm_layer,
+                    activation_func)
+            else:
+                self.conv = nn.Sequential(
+                    nn.Conv2d(in_channels, out_channels, stride=conv_stride, kernel_size=kernel_size, padding=padding,
+                              bias=bias),
+                    activation_func)
+
+    def forward(self, x, style_fea):
+        if self.style:
+            res = self.styleconv(x, style_fea)
+            res = self.relu(res)
+            return res
+        else:
+            return self.conv(x)
+
+
+class Conv2DMod(nn.Module):
+    def __init__(self, in_chan, out_chan, kernel, demod=True, stride=1, dilation=1, eps=1e-8, **kwargs):
+        super().__init__()
+        self.filters = out_chan
+        self.demod = demod
+        self.kernel = kernel
+        self.stride = stride
+        self.dilation = dilation
+        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))
+        self.eps = eps
+        nn.init.kaiming_normal_(self.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
+
+    def _get_same_padding(self, size, kernel, dilation, stride):
+        return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2
+
+    def forward(self, x, y):
+        b, c, h, w = x.shape
+
+        w1 = y[:, None, :, None, None]
+        w2 = self.weight[None, :, :, :, :]
+        weights = w2 * (w1 + 1)
+
+        if self.demod:
+            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps)
+            weights = weights * d
+
+        x = x.reshape(1, -1, h, w)
+
+        _, _, *ws = weights.shape
+        weights = weights.reshape(b * self.filters, *ws)
+
+        padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)
+        x = F.conv2d(x, weights, padding=padding, groups=b)
+
+        x = x.reshape(-1, self.filters, h, w)
+        return x
+
+class Up(nn.Module):
+    """ Upscaling then conv """
+
+    def __init__(self, in_channels, out_channels, activation_func='relu', style=False):
+        super().__init__()
+        self.style = style
+        activation_func = 'relu'
+        norm_layer, activation_func = get_layer_info(out_channels, activation_func)
+
+        self.up_layer = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.up = Conv(in_channels, in_channels // 4, activation_func=activation_func, style=style)
+
+    def forward(self, x, style_fea):
+        if self.style:
+            x = self.up_layer(x)
+            return self.up(x, style_fea)
+        else:
+            x = self.up_layer(x)
+            return self.up(x, style_fea)
+
+
+class PSPUpsample(nn.Module):
+    def __init__(self, in_channels, out_channels, scale_x, scale_y):
+        super().__init__()
+        self.conv = Conv(in_channels, out_channels)
+        self.scale_x = scale_x
+        self.scale_y = scale_y
+
+    def forward(self, x):
+        h, w = self.scale_y * x.size(2), self.scale_x * x.size(3)
+        p = F.upsample(input=x, size=(h, w), mode='bilinear')
+        return self.conv(p)
+
+
+class PSP(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        # pooling
+        self.pool2 = nn.AdaptiveAvgPool2d((2, 2))
+        self.pool4 = nn.AdaptiveAvgPool2d((4, 4))
+        self.pool8 = nn.AdaptiveAvgPool2d((8, 8))
+
+        # conv -> compress channels
+        avg_channel = in_channels // 4
+        self.conv2 = Conv(in_channels, avg_channel)
+        self.conv4 = Conv(in_channels, avg_channel)
+        self.conv8 = Conv(in_channels, avg_channel)
+        self.conv16 = Conv(in_channels, avg_channel)
+
+        # up sapmle -> match dimension
+        self.up2 = PSPUpsample(avg_channel, avg_channel, 16 // 2, 16 // 2)
+        self.up4 = PSPUpsample(avg_channel, avg_channel, 16 // 4, 16 // 4)
+        self.up8 = PSPUpsample(avg_channel, avg_channel, 16 // 8, 16 // 8)
+
+    def forward(self, x):
+        x2 = self.up2(self.conv2(self.pool2(x)))
+        x4 = self.up4(self.conv4(self.pool4(x)))
+        x8 = self.up8(self.conv8(self.pool8(x)))
+        x16 = self.conv16(x)
+        return torch.cat((x2, x4, x8, x16), dim=1)
+
+
+class Up_Stream(nn.Module):
+    """ Up Stream Sequence """
+
+    def __init__(self, out_channels=3, activation_func = 'relu'):
+        super(Up_Stream, self).__init__()
+
+        input_channel = 512
+        fea_dim = 200
+        norm_layer, activation_func = get_layer_info(input_channel, activation_func)
+        self.to_style1 = nn.Linear(in_features=fea_dim, out_features=input_channel)
+        self.up_16_16_1 = Conv(input_channel, 256, activation_func=activation_func, style=True)
+        self.up_16_16_2 = Conv(768, 512, activation_func=activation_func)
+        self.up_16_16_3 = Conv(1024, 512, activation_func=activation_func)
+
+        self.up_16_32 = Up(1024, 256, activation_func=activation_func)
+        self.to_style2 = nn.Linear(in_features=fea_dim, out_features=512)
+        self.up_32_32_1 = Conv(512, 256, activation_func=activation_func, style=True)
+
+        self.up_32_64 = Up(512, 128, activation_func=activation_func)
+        self.to_style3 = nn.Linear(in_features=fea_dim, out_features=256)
+        self.up_64_64_1 = Conv(256, 128, activation_func=activation_func, style=True)
+
+        self.up_64_128 = Up(256, 64, activation_func=activation_func)
+        self.to_style4 = nn.Linear(in_features=fea_dim, out_features=128)
+        self.up_128_128_1 = Conv(128, 64, activation_func=activation_func, style=True)
+
+        self.up_128_256 = Up(128, 32, activation_func=activation_func)
+        self.out_conv = Conv(64, out_channels, activation_func='relu')
+
+    def forward(self, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, style):
+        batch_size, c, h, w = x1.size()
+
+        # import pdb; pdb.set_trace()
+        # multiple channel ibl
+        # style = torch.zeros(batch_size, 500).to(x11.device)
+
+        # y = l.view(-1, 512, 1, 1).repeat(1, 1, 16, 16)
+        style1 = self.to_style1(style)
+        y = self.up_16_16_1(x11, style1)  # 256 x 16 x 16
+
+        y = torch.cat((x10, y), dim=1)  # 768 x 16 x 16
+        # print(y.size())
+
+        y = self.up_16_16_2(y, y)  # 512 x 16 x 16
+        # print(y.size())
+
+        y = torch.cat((x9, y), dim=1)  # 1024 x 16 x 16
+        # print(y.size())
+
+        # import pdb; pdb.set_trace()
+        y = self.up_16_16_3(y, y)  # 512 x 16 x 16
+        # print(y.size())
+
+        y = torch.cat((x8, y), dim=1)  # 1024 x 16 x 16
+        # print(y.size())
+
+        # import pdb; pdb.set_trace()
+        y = self.up_16_32(y, y)  # 256 x 32 x 32
+        # print(y.size())
+
+        y = torch.cat((x7, y), dim=1)
+        style2 = self.to_style2(style)
+        y = self.up_32_32_1(y, style2)  # 256 x 32 x 32
+        # print(y.size())
+
+        y = torch.cat((x6, y), dim=1)
+        y = self.up_32_64(y, y)
+        # print(y.size())
+        y = torch.cat((x5, y), dim=1)
+        style3 = self.to_style3(style)
+        y = self.up_64_64_1(y, style3)  # 128 x 64 x 64
+        # print(y.size())
+
+        y = torch.cat((x4, y), dim=1)
+        y = self.up_64_128(y, y)
+        # print(y.size())
+        y = torch.cat((x3, y), dim=1)
+        style4 = self.to_style4(style)
+        y = self.up_128_128_1(y, style4)  # 64 x 128 x 128
+        # print(y.size())
+
+        y = torch.cat((x2, y), dim=1)
+        y = self.up_128_256(y, y)  # 32 x 256 x 256
+        # print(y.size())
+
+        y = torch.cat((x1, y), dim=1)
+
+        y = self.out_conv(y, y)  # 3 x 256 x 256
+
+        return y

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/test.py

@@ -0,0 +1,24 @@
+import torch
+import ssn_dataset
+from torchvision import transforms, utils
+import numpy as np
+
+csv_file = "~/Dataset/soft_shadow/train/metadata.csv"
+# compose_transform = None
+training_dataset = ssn_dataset.SSN_Dataset(csv_file, is_training = True)
+testing_dataset = ssn_dataset.SSN_Dataset(csv_file, is_training = False)
+
+print('training set size: ', len(training_dataset))
+print('testing set size: ',len(testing_dataset))
+
+print(len(training_dataset.meta_data))
+print(training_dataset.meta_data[0])
+
+# for j in range(10):
+#     for i in range(len(training_dataset)):
+#         data = training_dataset[i]
+# #         print("{} \r".format(i), flush=True, end="")
+#         print("{} ".format(i))
+
+# for i,data in enumerate(testing_dataset):
+#     print("{} \r".format(i), flush=True, end="")

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/ssn/test_dataset.py

@@ -0,0 +1,21 @@
+import ssn_dataset
+import time
+
+if __name__ == '__main__':
+    start = time.time()
+    csv_file = "~/Dataset/soft_shadow/single_human/metadata.csv"
+    training_dataset = ssn_dataset.SSN_Dataset(csv_file, is_training=True)
+    testing_dataset = ssn_dataset.SSN_Dataset(csv_file, is_training=False)
+
+    print("Training dataset num: ", len(training_dataset))
+    print("Testing dataset num: ",len(testing_dataset))
+
+    for i in range(len(training_dataset)):
+        data = training_dataset[i]
+        print('Training set: successfully iterate {} \r'.format(i), flush=True, end='')
+
+    for i in range(len(testing_dataset)):
+        data = testing_dataset[i]
+        print('Validation set: successfully iterate {} \r'.format(i), flush=True, end='')
+
+    end = time.time()

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/html.py

@@ -0,0 +1,61 @@
+import dominate
+from dominate.tags import meta, h3, table, tr, td, p, a, img, br
+import os
+
+class HTML:
+    def __init__(self, web_dir, title, reflesh=0):
+        self.title = title
+        self.web_dir = web_dir
+        if not os.path.exists(self.web_dir):
+            os.makedirs(self.web_dir)
+
+        # print(self.img_dir)
+
+        self.doc = dominate.document(title=title)
+        if reflesh > 0:
+            with self.doc.head:
+                meta(http_equiv="reflesh", content=str(reflesh))
+
+    def get_image_dir(self):
+        return self.img_dir
+
+    def add_header(self, str):
+        with self.doc:
+            h3(str)
+
+    def add_table(self, border=1):
+        self.t = table(border=border, style="table-layout: fixed;")
+        self.doc.add(self.t)
+
+    def add_images(self, ims, txts, links, width=400, height=300):
+        self.add_table()
+        with self.t:
+            with tr():
+                for im, txt, link in zip(ims, txts, links):
+                    with td(style="word-wrap: break-word; height:{}px; width:{}px".format(height + 10,width + 10), halign="center", valign="top"):
+                        with p():
+                            with a(href=os.path.join('/',link)):
+                                img(style="width:{}px;height:{}".format(width, height), src=os.path.join('/',im))
+                            br()
+                            p(txt)
+
+    def save(self):
+        html_file = '%s/index.html' % self.web_dir
+        f = open(html_file, 'wt')
+        f.write(self.doc.render())
+        f.close()
+
+
+if __name__ == '__main__':
+    html = HTML('web/', 'test_html')
+    html.add_header('hello world')
+
+    ims = []
+    txts = []
+    links = []
+    for n in range(4):
+        ims.append('image_%d.png' % n)
+        txts.append('text_%d' % n)
+        links.append('image_%d.png' % n)
+    html.add_images(ims, txts, links)
+    html.save()

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/html_server.py

@@ -0,0 +1,9 @@
+import http.server
+import socketserver
+
+PORT = 8081
+Handler = http.server.SimpleHTTPRequestHandler
+
+with socketserver.TCPServer(("", PORT), Handler) as httpd:
+    print("serving at port", PORT)
+    httpd.serve_forever()

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/imgs

@@ -0,0 +1 @@
+/home/ysheng/Dataset/dropbox/Research/SSN_Training_Share/touch

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/make_html.py

@@ -0,0 +1,133 @@
+import numpy as np
+import json
+import pdb
+import os
+from os.path import join
+import html
+from tqdm import tqdm
+import argparse
+import pandas as pd
+import matplotlib.pyplot as plt
+
+def get_files(folder):
+    return [join(folder, f) for f in os.listdir(folder) if os.path.isfile(join(folder, f))]
+
+def get_folders(folder):
+    return [join(folder, f) for f in os.listdir(folder) if os.path.isdir(join(folder, f))]
+
+vis_img_folder = 'imgs'
+vis_eval_img_folder = 'eval_imgs'
+def eval_gen(webpage, output_folder, is_pattern=True):
+    def get_file(files, key_world):
+        mitsuba_shadow = ''
+        for f in files:
+            if f.find(key_world) != -1:
+                mitsuba_shadow = f
+                break
+        return mitsuba_shadow
+
+    def flip_shadow(img_file):
+        dirname, fname = os.path.dirname(img_file), os.path.splitext(os.path.basename(img_file))[0]
+        if img_file == '':
+            print('find one zero')
+            mts_shadow_np = np.zeros((256,256,3))
+        else:
+            mts_shadow_np = plt.imread(img_file)
+        
+        save_path = join(dirname, fname + '_flip.png')
+        plt.imsave(save_path, 1.0-mts_shadow_np)
+        return save_path
+
+    img_folders = join(output_folder, 'imgs')
+    folders = get_folders(img_folders)
+    print("There are {} folders".format(len(folders)))
+    
+    for model in tqdm(folders):
+        cur_model_relative = join(vis_img_folder, os.path.basename(model))
+        evl_cur_model_relative = join(vis_eval_img_folder, os.path.basename(model))
+
+        if is_pattern:
+            ibl_relative = join(cur_model_relative, 'pattern')
+        else:
+            ibl_relative = join(cur_model_relative, 'real')
+        
+        # import pdb; pdb.set_trace()
+        ibl_folders = get_folders(ibl_relative)
+        ibl_folders.sort()
+        for ibl in ibl_folders:
+            cur_ibl_relative = join(ibl_relative,os.path.basename(ibl))
+            gt_files = get_files(cur_ibl_relative)
+            mts_shadow = get_file(gt_files, '_shadow.png')
+
+            ibl_name = os.path.basename(ibl)
+            ibl = join(cur_ibl_relative, ibl_name + '.png')
+           
+            mitsuba_shadow = flip_shadow(mts_shadow)
+
+            cur_eval_folder = join(evl_cur_model_relative, join('pattern', ibl_name))
+            net_predict = get_file(get_files(cur_eval_folder), 'predict.png')
+
+            # mitsuba_final = join(cur_ibl_relative, 'composite.png')
+            # pred_final = join(cur_ibl_relative, 'composite_pred.png')
+            
+            # print(ibl_name)
+            ims, txts, links = [ibl,mitsuba_shadow, net_predict], ['ibl','mitsuba', 'predict'], [ibl,mitsuba_shadow, net_predict]
+
+            webpage.add_images(ims, txts, links)
+
+vis_pattern_folder = '/home/ysheng/Documents/vis_pattern'
+vis_real_folder = '/home/ysheng/Documents/vis_real'
+def vis_files_in_folder():
+    folder = '/home/ysheng/Documents/vis_models'
+    webpage = html.HTML(folder, 'models', reflesh=1)    
+    img_folders = join(folder, 'imgs')
+    files = get_files(img_folders)
+    print("There are {} files".format(len(files)))
+    
+    prefix_set = set()
+    for cur_file in tqdm(files):
+        cur_name = os.path.splitext(os.path.basename(cur_file))[0]
+        prefix_set.add(cur_name[:-3])
+
+    print('there are {} prefixs'.format(len(prefix_set)))
+    prefix_set = list(prefix_set)
+    prefix_set.sort()
+
+    # import pdb; pdb.set_trace()
+    relative_folder = './imgs'
+    for i, prefix in enumerate(prefix_set):
+        ims = [join(relative_folder, prefix + '{:03d}.png'.format(i)) for i in range(len(files) // len(prefix_set))]
+        txts = [prefix + '{:03d}'.format(i) for i in range(len(files) // len(prefix_set))]
+        links = ims
+        webpage.add_images(ims, txts, links)
+
+    webpage.save()
+    print('finished')
+    
+def vis_files(df_file):
+    """ input is a pandas dataframe
+        format: path, path,..., name,name, ...
+    """
+    folder = '.'
+    webpage = html.HTML(folder, 'benchmark', reflesh=1)    
+
+    relative_folder = './imgs'
+#     for i, prefix in enumerate(prefix_set):
+#         ims = [join(relative_folder, prefix + '{:03d}.png'.format(i)) for i in range(len(files) // len(prefix_set))]
+#         txts = [prefix + '{:03d}'.format(i) for i in range(len(files) // len(prefix_set))]
+#         links = ims
+#         webpage.add_images(ims, txts, links)
+    
+    df = pd.read_csv(df_file)
+    for i,v in tqdm(df.iterrows(), total=len(df)):
+        img_range = len(v)//2+1
+        imgs = [join(relative_folder,v[i]) for i in range(1,img_range)]
+        txts = [v[i] for i in range(img_range, len(v))]
+        links = imgs
+        webpage.add_images(imgs, txts, links)
+        
+    webpage.save()
+    print('finished')
+    
+if __name__ == "__main__":
+    vis_files('/home/ysheng/Documents/paper_project/adobe/soft_shadow/benchmark_results/html.csv')

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/net_utils.py

@@ -0,0 +1,70 @@
+#import matplotlib.pyplot as plt
+import os
+from torchvision import transforms, utils
+import torch
+#import matplotlib.pyplot as plt
+import numpy as np
+from utils.utils_file import get_cur_time_stamp, create_folder
+
+def compute_differentiable_params(net):
+    return sum(p.numel() for p in net.parameters() if p.requires_grad)
+
+def convert_Relight_latent_light(latent_feature):
+    """ Convert n x 6 x 16 x 16 -> n x 3 x 16 x 32 """
+    # torch image: C X H X W
+    batch_size, C, H, W = latent_feature.size()
+    latent_feature = torch.reshape(latent_feature, (batch_size, 3, 16, 32))  # make sure it is right
+    # print(latent_feature.size())
+    return latent_feature
+
+def show_batch(sample_batch, out_file=None):
+    grid = utils.make_grid(sample_batch)
+    plt.figure(figsize=(30,20))
+    plt.imshow(grid.detach().cpu().numpy().transpose((1,2,0)))
+
+    if not out_file is None:
+        print('try save ', out_file)
+        plt.savefig(out_file)
+
+    plt.show()
+
+def show_light_batch(light_batch):
+    light_batch = convert_Relight_latent_light(light_batch)
+    show_batch(light_batch)
+    
+def save_loss(figure_fname, train_loss, valid_loss):
+    plt.plot(train_loss)
+    plt.plot(valid_loss)
+    plt.legend(['train_loss', 'valid_loss'])
+    plt.savefig(figure_fname)
+
+def save_model(output_folder, model, optimizer, epoch, best_loss, fname, hist_train_loss, hist_valid_loss, hist_lr, params):
+    """ Save current best model into some folder """
+    create_folder(output_folder)
+
+    # cur_time_stamp = get_cur_time_stamp()
+    # output_fname = os.path.join(output_folder, exp_name + '_' + cur_time_stamp + ".pt")
+    output_fname = os.path.join(output_folder, fname)
+    tmp_model = model
+    if params.multi_gpu and hasattr(tmp_model, 'module'):
+        tmp_model = model.module
+
+    torch.save({
+        'epoch': epoch,
+        'best_loss': best_loss,
+        'model_state_dict': tmp_model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'hist_train_loss': hist_train_loss,
+        'hist_valid_loss': hist_valid_loss,
+        'hist_lr':hist_lr,
+        'params':str(params)
+        }, output_fname)
+    return output_fname
+    
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+def set_lr(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/tensorboard_utils.py

@@ -0,0 +1,29 @@
+import numpy as np
+import torch
+from torch.utils.tensorboard import SummaryWriter
+
+def tensorboard_plot_loss(win_name, loss, writer):
+    writer.add_scalar("Loss/{}".format(win_name), loss[-1], len(loss))
+    writer.flush()
+
+def normalize_img(imgs):
+    b,c,h,w = imgs.shape
+    gt_batch = b//2
+    for i in range(gt_batch):
+        factor = torch.max(imgs[i])
+        imgs[i] = imgs[i]/factor
+        imgs[gt_batch + i] = imgs[gt_batch + i]/factor
+
+    imgs = torch.clamp(imgs, 0.0,1.0)
+    return imgs
+
+def tensorboard_show_batch(imgs, writer, win_name=None, nrow=2, normalize=True, step=0):
+    if normalize:
+        imgs = normalize_img(imgs)
+
+    writer.add_images('{}'.format(win_name), imgs, step)
+    writer.flush()
+
+def tensorboard_log(log_info, writer, win_name='logger', step=0):
+    writer.add_text(win_name, log_info, step)
+    writer.flush()

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/time_utils.py

@@ -0,0 +1,6 @@
+import datetime
+
+def get_time_stamp():
+    return '{:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now())
+
+

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/utils_file.py

@@ -0,0 +1,59 @@
+from shutil import copyfile
+from os import listdir
+from os.path import isfile, join
+import os
+import datetime
+
+def get_all_folders(folder):
+    if check_file_exists(folder) == False:
+        print("Cannot find the folder ", folder)
+        return []
+    subfolders = [f for f in os.listdir(folder) if not isfile(join(folder,f))]
+    return subfolders
+
+def get_all_files(folder):
+    if check_file_exists(folder) == False:
+        print("Cannot find the folder ", folder)
+        return []
+ 
+    ori_files = [f for f in listdir(folder) if isfile(join(folder, f))]
+    return ori_files
+
+def create_folder(folder):
+    if not os.path.exists(folder):
+        os.mkdir(folder)
+
+def create_folders(folder_list):
+    for f in folder_list:
+        create_folder(f)
+
+def replace_file_ext(fname, new_ext):
+    ext_pos = fname.find(".")
+    if ext_pos != -1:
+        return fname[0:ext_pos] + "."+ new_ext
+    else:
+        print("Please check " + fname)
+
+def check_file_exists(fname, verbose=True):
+    try:
+        if not os.path.exists(fname) or get_file_size(fname) == 0:
+            if verbose:
+                print("file {} does not exists! ".format(fname))
+            return False
+    except:
+        print("File {} has some issue! ".format(fname))
+        return False
+    return True
+
+def delete_file(fname):
+    if check_file_exists(fname):
+        os.remove(fname)
+
+def get_file_size(fname):
+    return os.path.getsize(fname)
+
+def get_folder_size(folder):
+    return sum(os.path.getsize(folder + f) for f in listdir(folder) if isfile(join(folder, f)))
+
+def get_cur_time_stamp():
+    return datetime.datetime.now().strftime("%d-%B-%I-%M-%p")

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/vis_test_results.py

@@ -0,0 +1,21 @@
+import numpy as np
+import json
+import pdb
+import os
+from os.path import join
+import html
+from tqdm import tqdm
+import argparse
+import pandas as pd
+import cv2 as cv
+
+base_softness = "0.1"
+base_exp_name = "fov_results_real_89"
+root_dir = "/mnt/share/yifan/code/soft_shadow-master/vis_res/"
+base_dir = root_dir + base_exp_name + '/' + base_softness
+
+exp_name = ["fov_results_real_hd", "fov_results_real_hd_0713"]
+
+softness = [0.1, 0.2]
+
+case_name = ["case1", "case2", "case7"]

PixHtLab-Src/Demo/PixhtLab/ShadowStyle/inference/utils/visdom_utils.py

@@ -0,0 +1,53 @@
+from visdom import Visdom
+import numpy as np
+import torch
+
+# viz = Visdom(port=8002)
+# viz2 = Visdom(port=8003)
+
+def setup_visdom(port=8002):
+    return Visdom(port=port)
+
+def visdom_plot_loss(win_name, loss, cur_viz):
+    loss_np = np.array(loss)
+    x = np.arange(1, 1 + len(loss))
+    cur_viz.line(win=win_name,
+                 X=x,
+                 Y=loss_np,
+                 opts=dict(showlegend=True, legend=[win_name]))
+
+def guassian_light(light_tensor):
+    light_tensor = light_tensor.detach().cpu()
+    channel = light_tensor.size()[0]
+    tensor_ret = torch.zeros(light_tensor.size())
+    for i in range(channel):
+        light_np = light_tensor[0].numpy() * 100.0
+        light_np = gaussian_filter(light_np, sigma=2)
+        tensor_ret[i] = torch.from_numpy(light_np)
+        tensor_ret[i] = torch.clamp(tensor_ret[i], 0.0, 1.0)
+        
+    return tensor_ret
+
+def normalize_img(imgs):
+    b,c,h,w = imgs.shape
+    gt_batch = b//2
+    for i in range(gt_batch):
+        factor = torch.max(imgs[i])
+        imgs[i] = imgs[i]/factor
+        imgs[gt_batch + i] = imgs[gt_batch + i]/factor
+        # imgs[i] = imgs[i]/3.0
+        
+    imgs = torch.clamp(imgs, 0.0,1.0)
+    return imgs
+
+def visdom_show_batch(imgs, cur_viz, win_name=None, nrow=2, normalize=True):
+    if normalize:
+        imgs = normalize_img(imgs)
+    
+    if win_name is None:
+        cur_viz.images(imgs, win="batch visualize",nrow=nrow)
+    else:
+        cur_viz.images(imgs, win=win_name, opts=dict(title=win_name),nrow=nrow)
+    
+def visdom_log(log_info, viz, win_name='logger'):
+    viz.text(log_info, win=win_name)

PixHtLab-Src/Demo/PixhtLab/Torch_Render/hshadow_cuda.cpp

@@ -0,0 +1,98 @@
+#include <torch/extension.h>
+
+#include <vector>
+#include <stdio.h>
+
+// CUDA forward declarations
+std::vector<torch::Tensor> hshadow_render_cuda_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor mask_bb, torch::Tensor hmap, torch::Tensor rechmap, torch::Tensor light_pos);
+std::vector<torch::Tensor> reflect_render_cuda_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor hmap, torch::Tensor rechmap, torch::Tensor thresholds);
+std::vector<torch::Tensor> glossy_reflect_render_cuda_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor hmap, torch::Tensor rechmap, const int sample_n, const float glossy);
+torch::Tensor ray_intersect_cuda_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor hmap, torch::Tensor rechmap, torch::Tensor rd_map);
+torch::Tensor ray_scene_intersect_cuda_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor hmap, torch::Tensor ro, torch::Tensor rd, float dh);
+
+// C++ interface
+// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4.
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+/*  Heightmap Shadow Rendering 
+    rgb:        B x 3 x H x W
+    mask:       B x 1 x H x W
+    mask:       B x 1 
+    hmap:       B x 1 x H x W
+    rechmap:    B x 1 x H x W
+    light_pos:  B x 1 (x,y,h)
+*/
+std::vector<torch::Tensor> hshadow_render_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor bb, torch::Tensor hmap, torch::Tensor rechmap, torch::Tensor light_pos) {
+    CHECK_INPUT(rgb);
+    CHECK_INPUT(mask);
+    CHECK_INPUT(bb);
+    CHECK_INPUT(hmap);
+    CHECK_INPUT(rechmap);
+    CHECK_INPUT(light_pos);
+
+    return hshadow_render_cuda_forward(rgb, mask, bb, hmap, rechmap, light_pos);
+}
+
+std::vector<torch::Tensor> reflect_render_forward(torch::Tensor rgb, torch::Tensor mask, torch::Tensor hmap, torch::Tensor rechmap, torch::Tensor thresholds) {
+    CHECK_INPUT(rgb);
+    CHECK_INPUT(mask);
+    CHECK_INPUT(hmap);
+    CHECK_INPUT(rechmap);
+    CHECK_INPUT(thresholds);
+
+    return reflect_render_cuda_forward(rgb, mask, hmap, rechmap, thresholds);
+}
+
+
+std::vector<torch::Tensor> glossy_reflect_render_forward(torch::Tensor rgb,
+                                                  torch::Tensor mask,
+                                                  torch::Tensor hmap,
+                                                  torch::Tensor rechmap,
+                                                  int sample_n,
+                                                  float glossy) {
+    CHECK_INPUT(rgb);
+    CHECK_INPUT(mask);
+    CHECK_INPUT(hmap);
+    CHECK_INPUT(rechmap);
+
+    return glossy_reflect_render_cuda_forward(rgb, mask, hmap, rechmap, sample_n, glossy);
+}
+
+
+torch::Tensor ray_intersect_foward(torch::Tensor rgb,
+                            torch::Tensor mask,
+                            torch::Tensor hmap,
+                            torch::Tensor rechmap,
+                            torch::Tensor rd_map) {
+    CHECK_INPUT(rgb);
+    CHECK_INPUT(mask);
+    CHECK_INPUT(hmap);
+    CHECK_INPUT(rechmap);
+
+    return ray_intersect_cuda_forward(rgb, mask, hmap, rechmap, rd_map);
+}
+
+torch::Tensor ray_scene_intersect_foward(torch::Tensor rgb,
+                                         torch::Tensor mask,
+                                         torch::Tensor hmap,
+                                         torch::Tensor ro,
+                                         torch::Tensor rd,
+                                         float dh) {
+    CHECK_INPUT(rgb);
+    CHECK_INPUT(mask);
+    CHECK_INPUT(hmap);
+    CHECK_INPUT(ro);
+    CHECK_INPUT(rd);
+
+    return ray_scene_intersect_cuda_forward(rgb, mask, hmap, ro, rd, dh);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &hshadow_render_forward, "Heightmap Shadow Rendering Forward (CUDA)");
+    m.def("reflection", &reflect_render_forward, "Reflection Rendering Forward (CUDA)");
+    m.def("glossy_reflection", &glossy_reflect_render_forward, "Glossy Reflection Rendering Forward (CUDA)");
+    m.def("ray_intersect", &ray_intersect_foward, "Ray scene intersection");
+    m.def("ray_scene_intersect", &ray_scene_intersect_foward, "Ray scene intersection");
+}

PixHtLab-Src/Demo/PixhtLab/Torch_Render/hshadow_cuda_kernel.cu

@@ -0,0 +1,682 @@
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+#include <stdio.h>
+
+namespace {
+    template <typename scalar_t>
+    __device__
+    scalar_t sign(scalar_t t) {
+        if (t > 0.0) {
+            return (scalar_t)1.0;
+        } else {
+            return -(scalar_t)1.0;
+        }
+    }
+
+    template <typename scalar_t>
+    struct vec2 {
+        scalar_t x, y;
+
+        __device__
+        vec2() { x=0.0, y=0.0;}
+
+        __device__
+        vec2(scalar_t x, scalar_t y):x(x), y(y) {}
+    };
+
+
+    template <typename scalar_t>
+    struct vec3 {
+        scalar_t x, y, z;
+
+        __device__
+        vec3() { x=0.0, y=0.0, z=0.0;}
+
+        __device__
+        vec3(scalar_t x, scalar_t y, scalar_t z):x(x), y(y), z(z) {}
+    };
+
+
+    template <typename scalar_t>
+    __device__
+    scalar_t lerp(scalar_t a, scalar_t b, scalar_t t) {
+        return (1.0-t) * a + t * b;
+    }
+
+    template <typename scalar_t>
+    __device__
+    void proj_ground(
+        scalar_t x0, scalar_t y0, scalar_t h0,
+        scalar_t x1, scalar_t y1, scalar_t h1,
+        scalar_t &x2, scalar_t &y2
+    ) {
+        scalar_t t = (0-h0)/(h1-h0);
+        x2 = lerp(x0, x1, t);
+        y2 = lerp(y0, y1, t);
+    }
+
+
+    // line checking condition with thickness value dh, which is the height difference for double-height map
+    // we can also use dh as a tolerance value
+    template <typename scalar_t>
+    __device__
+    bool check_intersect(
+        scalar_t xa, scalar_t ya, scalar_t ha,
+        scalar_t xb, scalar_t yb, scalar_t hb,
+        scalar_t x, scalar_t y, scalar_t h,
+        scalar_t dh, int& flag) {
+        scalar_t t = xa == xb ? (y-ya)/(yb-ya):(x-xa)/(xb-xa);
+        scalar_t h_ = lerp(ha, hb, t);
+        flag = h_ <= h ? 1:-1;
+        return (h_ <= h) && (h_ >= h-dh);
+    }
+
+    /*
+    * Ray trace in the current scene
+    * Given start point xyh, light point xyh, current receiver's height map,
+    * Return:
+    *      1. if intersect or not
+    *      2. the color for the intersection point
+    * */
+    template <typename scalar_t>
+    __device__
+    bool ray_trace(vec3<scalar_t> s,
+                   vec3<scalar_t> l,
+                   const int bi,
+                   const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+                   const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+                   const torch::PackedTensorAccessor64<scalar_t,4> d_rechmap,
+                   const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+                   vec3<scalar_t> &out) {
+        bool ret = false;
+
+        const int batch_size = d_rgb.size(0), h = d_rgb.size(2), w = d_rgb.size(3);
+        scalar_t lx = l.x;
+        scalar_t ly = l.y;
+        scalar_t lh = l.z;
+
+        scalar_t recx = s.x;
+        scalar_t recy = s.y;
+        scalar_t rech = s.z;
+
+        scalar_t dirx = lx - recx, diry = ly - recy;
+        bool gox = abs(dirx) > abs(diry);
+        int searching_n = gox ? w : h;
+        int starti = 0, endi = searching_n;
+        if (lh > 0) {
+            if (gox) {
+                starti =  recx < lx ? recx:lx;
+
+                endi = recx < lx ? lx:recx;
+            }
+            else {
+                starti =  recy < ly ? recy:ly;
+                endi = recy < ly ? ly:recy;
+            }
+        }
+        if (lh < 0) {
+            if (gox) {
+                starti =  recx < lx ? 0:recx;
+                endi = recx < lx ? recx:endi;
+            }
+            else {
+                starti =  recy < ly ? 0:recy;
+                endi = recy < ly ? recy:endi;
+            }
+        }
+
+        scalar_t sx, sy;
+        int flag = 0, last_flag = 0;
+        for(int si = starti; si < endi; ++si) {
+            /* Searching Point xyh */
+            if (gox) {
+                sx = si;
+                sy = recy + (sx-recx)/dirx * diry;
+            } else {
+                sy = si;
+                sx = recx + (sy-recy)/diry * dirx;
+            }
+
+            if (sx < 0 || sx > w-1 || sy < 0 || sy > h-1 || d_mask[bi][0][sy][sx] < 0.989) {
+                last_flag = 0;
+                continue;
+            }
+
+
+            scalar_t sh0 = d_hmap[bi][0][sy][sx];
+            scalar_t sh1, sh;
+            // do linear interpolation for sh; note that either sy or sx are floating number
+            if (gox) {
+                if ( sy+1 > h-1 || d_mask[bi][0][sy+1][sx] < 0.989)
+                    sh = sh0;
+                else {
+                    sh1 = d_hmap[bi][0][sy+1][sx];
+                    sh = lerp(sh0, sh1, sy - int(sy));
+                }
+            }
+            else {
+                if ( sx+1 > w-1 || d_mask[bi][0][sy][sx+1] < 0.989)
+                    sh = sh0;
+                else {
+                    sh1 = d_hmap[bi][0][sy][sx+1];
+                    sh = lerp(sh0, sh1, sx - int(sx));
+                }
+            }
+
+            scalar_t dh = 1.0; // this controls the thickness; for double height map, dh = h_f - h_b
+            bool intersect = check_intersect(recx, recy, rech, lx, ly, lh, sx, sy, sh, dh, flag);
+            if (intersect) {
+                /* TODO, which sampling? linear interpolation?  */
+                out.x = d_rgb[bi][0][(int)sy][(int)sx];
+                out.y = d_rgb[bi][1][(int)sy][(int)sx];
+                out.z = d_rgb[bi][2][(int)sy][(int)sx];
+
+                ret = true;
+                break;
+            }
+            if (last_flag != 0){
+                if (last_flag != flag) {
+                    out.x = d_rgb[bi][0][(int)sy][(int)sx];
+                    out.y = d_rgb[bi][1][(int)sy][(int)sx];
+                    out.z = d_rgb[bi][2][(int)sy][(int)sx];
+
+                    ret = true;
+                    break;
+                }
+            }
+            last_flag = flag;
+        }
+
+        return ret;
+    }
+
+    /*
+    * Ray trace in the current scene
+    * Given start point xyh, light point xyh, current receiver's height map,
+    * Return:
+    *      1. if intersect or not
+    *      2. the color for the intersection point
+    * */
+    template <typename scalar_t>
+    __device__
+    bool ray_scene_intersect(vec3<scalar_t> ro,
+                            vec3<scalar_t> rd,
+                            const scalar_t dh,
+                            const int bi,
+                            const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+                            const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+                            const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+                            vec3<scalar_t> &out) {
+        bool ret = false;
+        int h = d_mask.size(2);
+        int w = d_mask.size(3);
+        scalar_t dirx = rd.x, diry = rd.y, dirh = rd.z;
+
+        /* Special Case, there's no direction update in x or y, but in h */
+        if (abs(dirx) < 1e-6f && abs(diry) < 1e-6f) {
+            out.x = d_rgb[bi][0][(int)ro.y][(int)ro.x];
+            out.y = d_rgb[bi][1][(int)ro.y][(int)ro.x];
+            out.z = d_rgb[bi][2][(int)ro.y][(int)ro.x];
+            return true;
+        }
+
+        bool gox = abs(dirx) > abs(diry);
+        int searching_n = gox ? w : h;
+
+        scalar_t cur_h;
+        scalar_t sx, sy;
+
+        int prev_sign, cur_sign;
+
+        // for(int si = starti; si < endi; ++si) {
+        for(int si = 0; si < searching_n; ++si) {
+            /* Searching Point XYH */
+            if (gox) {
+                sx = ro.x + si * sign(dirx);
+                sy = ro.y + (sx-ro.x)/dirx * diry;
+            } else {
+                sy = ro.y + si * sign(diry);
+                sx = ro.x + (sy-ro.y)/diry * dirx;
+            }
+
+            if (sx < 0 || sx > w-1 || sy < 0 || sy > h-1 || d_mask[bi][0][sy][sx] < 0.989) {
+                continue;
+            }
+
+            scalar_t sh0 = d_hmap[bi][0][sy][sx];
+            scalar_t sh1, sh;
+            // do linear interpolation for sh; note that either sy or sx are floating number
+            if (gox) {
+                if (sy+1 > h-1 || d_mask[bi][0][sy+1][sx] < 0.989)
+                    sh = sh0;
+                else {
+                    sh1 = d_hmap[bi][0][sy+1][sx];
+                    sh = lerp(sh0, sh1, sy - int(sy));   // Always use 0.5 to do interpolation
+                }
+
+                cur_h = ro.z + (sx - ro.x) / dirx * dirh;
+            }
+            else {
+                if ( sx + 1 > w-1 || d_mask[bi][0][sy][sx+1] < 0.989)
+                    sh = sh0;
+                else {
+                    sh1 = d_hmap[bi][0][sy][sx+1];
+                    sh = lerp(sh0, sh1, sx - int(sx));
+                }
+
+                cur_h = ro.z + (sy - ro.y) / diry * dirh;
+            }
+
+            // collide with the rechmap?
+            if (si == 0) {  /* First sign */
+                cur_sign = cur_h - sh; 
+                continue;
+            } else { 
+                prev_sign = cur_sign;
+            }
+
+            cur_sign = cur_h - sh; 
+            // if (cur_sign * prev_sign < 0.0 || abs(cur_sign) < dh) { /* pass through some objects */
+            if (abs(cur_sign) < dh) { /* pass through some objects */
+                out.x = d_rgb[bi][0][(int)sy][(int)sx];
+                out.y = d_rgb[bi][1][(int)sy][(int)sx];
+                out.z = d_rgb[bi][2][(int)sy][(int)sx];
+                ret = true;
+                break;
+            }
+
+        }
+
+        return ret;
+    }
+
+
+    template <typename scalar_t>
+    __global__ void hshadow_render_cuda_forward(
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+        const torch::PackedTensorAccessor64<scalar_t,2> d_bb,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rechmap,
+        const torch::PackedTensorAccessor64<scalar_t,2> d_lightpos,
+        torch::PackedTensorAccessor64<scalar_t,4> d_shadow) {
+        const int wstride = gridDim.x * blockDim.x, hstride = gridDim.y * blockDim.y, bstride = gridDim.z * blockDim.z;
+        const int batch_size = d_rgb.size(0), h = d_rgb.size(2), w = d_rgb.size(3);
+
+        for (int bi = blockIdx.z; bi < batch_size; bi += bstride) {
+            /* light xyh */
+            scalar_t lx = d_lightpos[bi][0], ly = d_lightpos[bi][1], lh = d_lightpos[bi][2];
+            int minh = max((int)d_bb[bi][0], 0), maxh = min((int)d_bb[bi][1], h-1), minw = max((int)d_bb[bi][2], 0), maxw = min((int)d_bb[bi][3], w-1);
+
+            vec3<scalar_t> light(lx, ly, lh);
+            for (int wi = blockIdx.x * blockDim.x + threadIdx.x;  wi < w; wi += wstride) for(int hi = blockIdx.y * blockDim.y + threadIdx.y;  hi < h; hi += hstride) {
+                scalar_t shadow(1.0), mask_alpha(0.0);
+                scalar_t recx = wi + 0.5, recy = hi+0.5, rech = d_rechmap[bi][0][hi][wi];
+
+                vec3<scalar_t> start(recx, recy, rech);
+                vec3<scalar_t> intersect_color;
+
+                /* Searching Potentials */
+                if (ray_trace(start, light, bi, d_mask, d_hmap, d_rechmap, d_rgb, intersect_color)) {
+                    shadow = 0.0;
+                }
+
+                d_shadow[bi][0][hi][wi] = shadow;
+                d_shadow[bi][1][hi][wi] = shadow;
+                d_shadow[bi][2][hi][wi] = shadow;
+            }
+        }
+    }
+
+
+    template <typename scalar_t>
+    __global__ void ray_intersect_cuda_forward(
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rechmap,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rd_map,
+        torch::PackedTensorAccessor64<scalar_t,4> d_intersect) {
+
+        const int wstride = gridDim.x * blockDim.x, hstride = gridDim.y * blockDim.y, bstride = gridDim.z * blockDim.z;
+        const int batch_size = d_rgb.size(0), h = d_rgb.size(2), w = d_rgb.size(3);
+        const scalar_t default_value = 0.0;
+
+        for (int bi = blockIdx.z; bi < batch_size; bi += bstride) {
+            for (int wi = blockIdx.x * blockDim.x + threadIdx.x;  wi < w; wi += wstride) for(int hi = blockIdx.y * blockDim.y + threadIdx.y;  hi < h; hi += hstride) {
+                scalar_t shadow(1.0), mask_alpha(0.0);
+                scalar_t recx = wi + 0.5, recy = hi+0.5, rech = d_rechmap[bi][0][hi][wi];
+
+                scalar_t lx = d_rd_map[bi][0][hi][wi];
+                scalar_t ly = d_rd_map[bi][1][hi][wi];
+                scalar_t lh = d_rd_map[bi][2][hi][wi];
+
+                vec3<scalar_t> start(recx, recy, rech);
+                vec3<scalar_t> rd(lx, ly, lh);
+
+                vec3<scalar_t> intersect_color;
+                if (ray_trace(start, rd, bi, d_mask, d_hmap, d_rechmap, d_rgb, intersect_color)) {
+                    d_intersect[bi][0][hi][wi] = intersect_color.x;
+                    d_intersect[bi][1][hi][wi] = intersect_color.y;
+                    d_intersect[bi][2][hi][wi] = intersect_color.z;
+                    d_intersect[bi][3][hi][wi] = 1.0;
+                } else {
+                    d_intersect[bi][0][hi][wi] = default_value;
+                    d_intersect[bi][1][hi][wi] = default_value;
+                    d_intersect[bi][2][hi][wi] = default_value;
+                    d_intersect[bi][3][hi][wi] = 0.0;
+                }
+            }
+        }
+    }
+
+    template <typename scalar_t>
+    __global__ void ray_scene_intersect_cuda_forward(
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_ro,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rd,
+        const scalar_t dh,
+        torch::PackedTensorAccessor64<scalar_t,4> d_intersect) {
+
+        const int wstride = gridDim.x * blockDim.x, hstride = gridDim.y * blockDim.y, bstride = gridDim.z * blockDim.z;
+        const int batch_size = d_rgb.size(0), h = d_rgb.size(2), w = d_rgb.size(3);
+        const scalar_t default_value = 0.0;
+
+        for (int bi = blockIdx.z; bi < batch_size; bi += bstride) {
+            for (int wi = blockIdx.x * blockDim.x + threadIdx.x;  wi < w; wi += wstride)  {
+                for(int hi = blockIdx.y * blockDim.y + threadIdx.y;  hi < h; hi += hstride) {
+                    // scalar_t rox = wi + 0.5, roy = hi+0.5, roh = d_ro[bi][0][hi][wi];
+                    scalar_t rox = d_ro[bi][0][hi][wi];
+                    scalar_t roy = d_ro[bi][1][hi][wi];
+                    scalar_t roh = d_ro[bi][2][hi][wi];
+
+                    scalar_t rdx = d_rd[bi][0][hi][wi];
+                    scalar_t rdy = d_rd[bi][1][hi][wi];
+                    scalar_t rdh = d_rd[bi][2][hi][wi];
+
+                    vec3<scalar_t> ro(rox, roy, roh);
+                    vec3<scalar_t> rd(rdx, rdy, rdh);
+
+                    vec3<scalar_t> intersect_color;
+                    if (ray_scene_intersect(ro, rd, dh, bi, d_mask, d_hmap, d_rgb, intersect_color)) {
+                        d_intersect[bi][0][hi][wi] = intersect_color.x;
+                        d_intersect[bi][1][hi][wi] = intersect_color.y;
+                        d_intersect[bi][2][hi][wi] = intersect_color.z;
+                        d_intersect[bi][3][hi][wi] = 1.0;
+                    } else {
+                        d_intersect[bi][0][hi][wi] = default_value;
+                        d_intersect[bi][1][hi][wi] = default_value;
+                        d_intersect[bi][2][hi][wi] = default_value;
+                        d_intersect[bi][3][hi][wi] = 0.0;
+                    }
+                }
+            }
+        }
+    }
+
+    template <typename scalar_t>
+    __global__ void reflect_render_cuda_forward(
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rechmap,
+        const torch::PackedTensorAccessor64<scalar_t,2> d_thresholds,
+        torch::PackedTensorAccessor64<scalar_t,4> d_reflect,
+        torch::PackedTensorAccessor64<scalar_t,4> d_reflect_height,
+        torch::PackedTensorAccessor64<scalar_t,4> d_reflect_mask) {
+        const int wstride = gridDim.x * blockDim.x, hstride = gridDim.y * blockDim.y, bstride = gridDim.z * blockDim.z;
+        const int batch_size = d_rgb.size(0), h = d_rgb.size(2), w = d_rgb.size(3);
+
+        for (int bi = blockIdx.z; bi < batch_size; bi += bstride) {
+            for (int wi = blockIdx.x * blockDim.x + threadIdx.x;  wi < w; wi += wstride) for(int hi = blockIdx.y * blockDim.y + threadIdx.y;  hi < h; hi += hstride) {
+                /*  Back tracing along the height
+                    Find the closest point, filter the closest point
+                */
+                scalar_t min_dis = FLT_MAX;
+                scalar_t min_r, min_g, min_b, min_height, min_mask;
+                for(int ti = hi-1; ti >= 0; --ti) {
+                    if (d_mask[bi][0][ti][wi] < 0.45)
+                        continue;
+
+                    scalar_t dis = abs(d_hmap[bi][0][ti][wi] * 2 + ti - hi);
+                    if (dis < min_dis) {
+                        min_dis = dis;
+                        min_r = d_rgb[bi][0][ti][wi];
+                        min_g = d_rgb[bi][1][ti][wi];
+                        min_b = d_rgb[bi][2][ti][wi];
+
+                        min_height = d_hmap[bi][0][ti][wi];
+                        min_mask = d_mask[bi][0][ti][wi];
+                    }
+                }
+
+                /* Check Condition */
+                scalar_t cur_thresholds = d_thresholds[bi][0];
+                if (min_dis < cur_thresholds) {
+                    /* Let Use Nearest Neighbor First */
+                    d_reflect[bi][0][hi][wi] = min_r;
+                    d_reflect[bi][1][hi][wi] = min_g;
+                    d_reflect[bi][2][hi][wi] = min_b;
+                    d_reflect_height[bi][0][hi][wi] = min_height;
+                    d_reflect_mask[bi][0][hi][wi] = 1.0;
+                }
+
+                // } else {
+                //     scalar_t fadding = 1.0-(min_dis-cur_thresholds);
+                //     if (fadding < 0.0) fadding = 0.0;
+                //     d_reflect[bi][0][hi][wi] = min_r * fadding + (1.0-fadding);
+                //     d_reflect[bi][1][hi][wi] = min_g * fadding + (1.0-fadding);
+                //     d_reflect[bi][2][hi][wi] = min_b * fadding + (1.0-fadding);
+                //     d_reflect_height[bi][0][hi][wi] = 0.0;
+                //     d_reflect_mask[bi][0][hi][wi] = fadding;
+                // }
+            }
+        }
+    }
+
+    template <typename scalar_t>
+    __global__ void glossy_reflect_render_cuda_forward(
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rgb,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_mask,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_hmap,
+        const torch::PackedTensorAccessor64<scalar_t,4> d_rechmap,
+        const int sample_n,
+        const float glossy,
+        torch::PackedTensorAccessor64<scalar_t,4> d_reflect) {
+
+        const int wstride = gridDim.x * blockDim.x, hstride = gridDim.y * blockDim.y, bstride = gridDim.z * blockDim.z;
+        const int batch_size = d_rgb.size(0), h = d_rgb.size(2), w = d_rgb.size(3);
+
+        for (int bi = blockIdx.z; bi < batch_size; bi += bstride) {
+            for (int wi = blockIdx.x * blockDim.x + threadIdx.x;  wi < w; wi += wstride) for(int hi = blockIdx.y * blockDim.y + threadIdx.y;  hi < h; hi += hstride) {
+                /*  Back tracing along the height
+                    Find the closest point, filter the closest point
+                */
+                scalar_t min_dis = FLT_MAX;
+                scalar_t min_r, min_g, min_b, min_height, min_mask;
+                for(int ti = hi-1; ti >= 0; --ti) {
+                    if (d_mask[bi][0][ti][wi] < 0.45)
+                        continue;
+
+                    scalar_t dis = abs(d_hmap[bi][0][ti][wi] * 2 + ti - hi);
+                    if (dis < min_dis) {
+                        min_dis = dis;
+                        min_r = d_rgb[bi][0][ti][wi];
+                        min_g = d_rgb[bi][1][ti][wi];
+                        min_b = d_rgb[bi][2][ti][wi];
+
+                        min_height = d_hmap[bi][0][ti][wi];
+                        min_mask = d_mask[bi][0][ti][wi];
+                    }
+                }
+
+                /* Check Condition */
+                float cur_thresholds = 1e-1;
+                if (min_dis < cur_thresholds) {
+                    /* Let Use Nearest Neighbor First */
+                    d_reflect[bi][0][hi][wi] = min_r;
+                    d_reflect[bi][1][hi][wi] = min_g;
+                    d_reflect[bi][2][hi][wi] = min_b;
+                }
+            }
+        }
+    }
+
+} // namespace
+
+std::vector<torch::Tensor> hshadow_render_cuda_forward(
+    torch::Tensor rgb,
+    torch::Tensor mask,
+    torch::Tensor bb,
+    torch::Tensor hmap,
+    torch::Tensor rechmap,
+    torch::Tensor light_pos) {
+    const auto batch_size = rgb.size(0);
+    const auto channel_size = rgb.size(1);
+    const auto h = rgb.size(2);
+    const auto w = rgb.size(3);
+    const dim3 threads(16, 16, 1);
+    const dim3 blocks((w + threads.x - 1) / threads.x, (h+threads.y-1)/threads.y, batch_size);
+    torch::Tensor shadow_tensor = torch::ones({batch_size, 3, h, w}).to(rgb);
+
+    AT_DISPATCH_FLOATING_TYPES(rgb.type(), "hshadow_render_cuda_forward", ([&] {
+        hshadow_render_cuda_forward<scalar_t><<<blocks, threads>>>(
+            rgb.packed_accessor64<scalar_t,4>(),
+            mask.packed_accessor64<scalar_t,4>(),
+            bb.packed_accessor64<scalar_t,2>(),
+            hmap.packed_accessor64<scalar_t,4>(),
+            rechmap.packed_accessor64<scalar_t,4>(),
+            light_pos.packed_accessor64<scalar_t,2>(),
+            shadow_tensor.packed_accessor64<scalar_t,4>());
+    }));
+
+    return {shadow_tensor};
+}
+
+std::vector<torch::Tensor> reflect_render_cuda_forward(
+    torch::Tensor rgb,
+    torch::Tensor mask,
+    torch::Tensor hmap,
+    torch::Tensor rechmap,
+    torch::Tensor thresholds) {
+    const auto batch_size = rgb.size(0);
+    const auto channel_size = rgb.size(1);
+    const auto h = rgb.size(2);
+    const auto w = rgb.size(3);
+    const dim3 threads(16, 16, 1);
+    const dim3 blocks((w + threads.x - 1) / threads.x, (h+threads.y-1)/threads.y, batch_size);
+    torch::Tensor reflection_tensor = torch::ones({batch_size, 3, h, w}).to(rgb);
+    torch::Tensor reflection_mask_tensor = torch::zeros({batch_size, 1, h, w}).to(rgb);
+    torch::Tensor reflection_height_tensor = torch::zeros({batch_size, 1, h, w}).to(rgb);
+
+    AT_DISPATCH_FLOATING_TYPES(rgb.type(), "reflect_render_cuda_forward", ([&] {
+        reflect_render_cuda_forward<scalar_t><<<blocks, threads>>>(
+            rgb.packed_accessor64<scalar_t,4>(),
+            mask.packed_accessor64<scalar_t,4>(),
+            hmap.packed_accessor64<scalar_t,4>(),
+            rechmap.packed_accessor64<scalar_t,4>(),
+            thresholds.packed_accessor64<scalar_t,2>(),
+            reflection_tensor.packed_accessor64<scalar_t,4>(),
+            reflection_height_tensor.packed_accessor64<scalar_t,4>(),
+            reflection_mask_tensor.packed_accessor64<scalar_t,4>());
+    }));
+
+    return {reflection_tensor, reflection_height_tensor,reflection_mask_tensor};
+}
+
+
+std::vector<torch::Tensor> glossy_reflect_render_cuda_forward(torch::Tensor rgb,
+                                                              torch::Tensor mask,
+                                                              torch::Tensor hmap,
+                                                              torch::Tensor rechmap,
+                                                              const int sample_n,
+                                                              const float glossy) {
+    const auto batch_size = rgb.size(0);
+    const auto channel_size = rgb.size(1);
+    const auto h = rgb.size(2);
+    const auto w = rgb.size(3);
+    const dim3 threads(16, 16, 1);
+    const dim3 blocks((w + threads.x - 1) / threads.x, (h+threads.y-1)/threads.y, batch_size);
+
+    torch::Tensor reflection_tensor = torch::ones({batch_size, 3, h, w}).to(rgb);
+
+    AT_DISPATCH_FLOATING_TYPES(rgb.type(), "reflect_render_cuda_forward", ([&] {
+        glossy_reflect_render_cuda_forward<scalar_t><<<blocks, threads>>>(
+            rgb.packed_accessor64<scalar_t,4>(),
+            mask.packed_accessor64<scalar_t,4>(),
+            hmap.packed_accessor64<scalar_t,4>(),
+            rechmap.packed_accessor64<scalar_t,4>(),
+            sample_n,
+            glossy,
+            reflection_tensor.packed_accessor64<scalar_t,4>());
+    }));
+
+    return {reflection_tensor};
+
+}
+
+
+torch::Tensor ray_intersect_cuda_forward(torch::Tensor rgb,
+                                        torch::Tensor mask,
+                                        torch::Tensor hmap,
+                                        torch::Tensor rechmap,
+                                        torch::Tensor rd_map){
+    const auto batch_size = rgb.size(0);
+    const auto channel_size = rgb.size(1);
+    const auto h = rgb.size(2);
+    const auto w = rgb.size(3);
+    const dim3 threads(16, 16, 1);
+    const dim3 blocks((w + threads.x - 1) / threads.x, (h+threads.y-1)/threads.y, batch_size);
+
+    torch::Tensor intersect_tensor = torch::ones({batch_size, 4, h, w}).to(rgb);
+
+    AT_DISPATCH_FLOATING_TYPES(rgb.type(), "reflect_render_cuda_forward", ([&] {
+        ray_intersect_cuda_forward<scalar_t><<<blocks, threads>>>(
+            rgb.packed_accessor64<scalar_t,4>(),
+            mask.packed_accessor64<scalar_t,4>(),
+            hmap.packed_accessor64<scalar_t,4>(),
+            rechmap.packed_accessor64<scalar_t,4>(),
+            rd_map.packed_accessor64<scalar_t,4>(),
+            intersect_tensor.packed_accessor64<scalar_t,4>());
+    }));
+
+    return intersect_tensor;
+
+}
+
+
+torch::Tensor ray_scene_intersect_cuda_forward(torch::Tensor rgb,
+                                               torch::Tensor mask,
+                                               torch::Tensor hmap,
+                                               torch::Tensor ro,
+                                               torch::Tensor rd,
+                                               float dh){
+    const auto batch_size = rgb.size(0);
+    const auto channel_size = rgb.size(1);
+    const auto h = rgb.size(2);
+    const auto w = rgb.size(3);
+    const dim3 threads(16, 16, 1);
+    const dim3 blocks((w + threads.x - 1) / threads.x, (h+threads.y-1)/threads.y, batch_size);
+
+    torch::Tensor intersect_tensor = torch::ones({batch_size, 4, h, w}).to(rgb);
+
+    AT_DISPATCH_FLOATING_TYPES(rgb.type(), "reflect_render_cuda_forward", ([&] {
+        ray_scene_intersect_cuda_forward<scalar_t><<<blocks, threads>>>(
+            rgb.packed_accessor64<scalar_t,4>(),
+            mask.packed_accessor64<scalar_t,4>(),
+            hmap.packed_accessor64<scalar_t,4>(),
+            ro.packed_accessor64<scalar_t,4>(),
+            rd.packed_accessor64<scalar_t,4>(),
+            dh,
+            intersect_tensor.packed_accessor64<scalar_t,4>());
+    }));
+
+    return intersect_tensor;
+
+}

PixHtLab-Src/Demo/PixhtLab/Torch_Render/plane_visualize.cpp

@@ -0,0 +1,26 @@
+
+#include <torch/extension.h>
+
+#include <vector>
+#include <stdio.h>
+
+// CUDA forward declarations
+std::vector<torch::Tensor> plane_visualize_cuda(torch::Tensor planes, torch::Tensor camera, int h, int w);
+
+// C++ interface
+// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4.
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> plane_visualize(torch::Tensor planes, torch::Tensor camera, int h, int w) {
+    CHECK_INPUT(planes);
+    CHECK_INPUT(camera);
+
+    return plane_visualize_cuda(planes, camera, h, w);
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+	m.def("forward", &plane_visualize, "Plane Visualization (CUDA)");
+}

PixHtLab-Src/Demo/PixhtLab/Torch_Render/plane_visualize_cuda.cu

@@ -0,0 +1,237 @@
+#include <torch/extension.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+
+namespace {
+	template <typename scalar_t>
+	struct vec3 {
+		scalar_t x, y, z;
+
+		__device__ __host__
+		vec3<scalar_t>():x(0),y(0),z(0) {}
+
+		__device__ __host__
+		vec3<scalar_t>(scalar_t a):x(a),y(a),z(a) {}
+
+		__device__ __host__
+		vec3<scalar_t>(scalar_t xx, scalar_t yy, scalar_t zz):x(xx),y(yy),z(zz) {}
+
+		__device__ __host__
+		vec3<scalar_t> operator*(const scalar_t &rhs) const {
+			vec3<scalar_t> ret(x,y,z);
+			ret.x = x * rhs;
+			ret.y = y * rhs;
+			ret.z = z * rhs;
+			return ret;
+		}
+
+		__device__ __host__
+		vec3<scalar_t> operator/(const scalar_t &rhs) const {
+			vec3<scalar_t> ret(x,y,z);
+			ret.x = x / rhs;
+			ret.y = y / rhs;
+			ret.z = z / rhs;
+			return ret;
+		}
+
+		__device__ __host__
+		vec3<scalar_t> operator+(const vec3<scalar_t> &rhs) const {
+			vec3<scalar_t> ret(x,y,z);
+			ret.x = x + rhs.x;
+			ret.y = y + rhs.y;
+			ret.z = z + rhs.z;
+			return ret;
+		}
+
+		__device__ __host__
+		vec3<scalar_t> operator-(const vec3<scalar_t> &rhs) const {
+			vec3<scalar_t> ret(x,y,z);
+			ret.x = x - rhs.x;
+			ret.y = y - rhs.y;
+			ret.z = z - rhs.z;
+			return ret;
+		}
+	};
+
+	template <typename scalar_t>
+	struct Ray {
+		vec3<scalar_t> ro, rd;
+	};
+
+	template <typename scalar_t>
+	struct Scene {
+		vec3<scalar_t> pp, pn;
+	};
+
+	template <typename scalar_t>
+	__device__
+	float deg2rad(scalar_t d) {
+		return d/180.0 * 3.1415926f;
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	scalar_t dot(vec3<scalar_t> a, vec3<scalar_t> b) {
+		return a.x * b.x + a.y * b.y + a.z * b.z;
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	vec3<scalar_t> cross(vec3<scalar_t> a, vec3<scalar_t> b) {
+		vec3<scalar_t> ret(0.0f);
+		ret.x = a.y * b.z - a.z * b.y;
+		ret.y = a.z * b.x - a.x * b.z;
+		ret.z = a.x * b.y - a.y * b.x;
+		return ret;
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	scalar_t length(vec3<scalar_t> a) {
+		return sqrt(dot(a, a));
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	vec3<scalar_t> normalize(vec3<scalar_t> a) {
+		return a/length(a);
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	scalar_t get_focal(int w, scalar_t fov) {
+		return 0.5 * w / tan(deg2rad(fov * 0.5));
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	vec3<scalar_t> get_rd(vec3<scalar_t> right, vec3<scalar_t> front, vec3<scalar_t> up, int h, int w, scalar_t focal, scalar_t x, scalar_t y) {
+		/* x, y in [-1, 1] */
+		right = normalize(right);
+		front = normalize(front); 
+		up = normalize(up);
+
+		return front * focal + right * x * (float)w * 0.5f + up * y * (float)h * 0.5f;
+	}
+
+	template <typename scalar_t>
+	__device__ 
+	Ray<scalar_t> get_ray(int h, int w, float hi, float wi, vec3<scalar_t> right, vec3<scalar_t>front, vec3<scalar_t> up, vec3<scalar_t> cam_pos, float fov) {
+		/* Note, wi/hi is in [-1.0, 1.0] */
+		Ray<scalar_t> ray;
+
+		float focal = 0.5f * w / tan(deg2rad(fov)); 
+		ray.ro = cam_pos;
+		ray.rd = front * focal + right * 0.5f * w * wi + up * 0.5f * h * hi;
+		return ray;
+	}
+	
+	template <typename scalar_t>
+	__device__
+	bool plane_intersect(Ray<scalar_t> ray, vec3<scalar_t> p, vec3<scalar_t> n, float &t) {
+		vec3<scalar_t> ro = ray.ro, rd = ray.rd;
+		t = dot(p-ro, n)/dot(rd, n);
+		return t >= 0.0;
+	}
+
+	template <typename scalar_t>
+	__device__
+	vec3<scalar_t> horizon2front(scalar_t horizon, int h, int w, float fov) {
+		scalar_t yoffset =  h / 2 - horizon;
+		scalar_t focal = 0.5f * w / tan(deg2rad(fov));
+		vec3<scalar_t> front = vec3<scalar_t>(0.0f,0.0f,-1.0f) * focal + vec3<scalar_t>(0.0f, 1.0f, 0.0f) * yoffset;
+		return normalize(front); 
+	}
+
+	template <typename scalar_t>
+	__device__
+	vec3<scalar_t> plane_texture(vec3<scalar_t> p) {
+		float freq = 6.0f;
+		float u = sin(p.x * freq), v = sin(p.z * freq);
+		vec3<scalar_t> ret(0.0f);
+		float line_width = 0.05f; 
+		if ((abs(u) < line_width || abs(v) < line_width)) {
+			ret = vec3<scalar_t>(1.0f);
+		}
+		return ret;
+	}
+
+	template <typename scalar_t>
+	__device__
+	bool ray_scene_trace(Ray<scalar_t> ray, Scene<scalar_t> scene, vec3<scalar_t> &color) {
+		color = vec3<scalar_t>(0.0f);
+		float t;
+		if (plane_intersect(ray, scene.pp, scene.pn, t)) {
+			vec3<scalar_t> intersect_pos = ray.ro + ray.rd * t;
+			color = plane_texture(intersect_pos);
+			return true;
+		}
+		return false;
+	}
+
+	template <typename scalar_t>
+	__global__ void plane_visualize_foward(
+		const torch::PackedTensorAccessor64<scalar_t,2> d_plane,
+		const torch::PackedTensorAccessor64<scalar_t,2> d_camera,
+        torch::PackedTensorAccessor64<scalar_t,4> d_vis) {
+		const int wstride = gridDim.x * blockDim.x, hstride = gridDim.y * blockDim.y, bstride = gridDim.z * blockDim.z;
+		const int batch_size = d_vis.size(0), h = d_vis.size(2), w = d_vis.size(3);
+		const int samples = 10;
+
+		vec3<scalar_t> cam_pos(0.0f, 1.0f, 1.0f), front(0.0f,0.0f,-1.0f), right(1.0f, 0.0f, 0.0f), up(0.0f, 1.0f, 0.0f);
+		for (int bi = blockIdx.z; bi < batch_size; bi += bstride) {
+			// scalar_t px = d_plane[bi][0], py = d_plane[bi][1], pz = d_plane[bi][2];
+			vec3<scalar_t> plane_pos(d_plane[bi][0], d_plane[bi][1], d_plane[bi][2]);
+			vec3<scalar_t> plane_norm(d_plane[bi][3], d_plane[bi][4], d_plane[bi][5]);
+			Scene<scalar_t> scene = {plane_pos, plane_norm};
+
+			scalar_t fov = d_camera[bi][0], horizon = d_camera[bi][1]; 
+			front = normalize(horizon2front(horizon, h, w, fov));
+			up = normalize(cross(right, front));
+			for (int wi = blockIdx.x * blockDim.x + threadIdx.x; wi < w; wi += wstride) 
+                for(int hi = blockIdx.y * blockDim.y + threadIdx.y; hi < h; hi += hstride) {
+					bool intersect = false;
+					vec3<scalar_t> color(0.0f);
+					for (int si = 0; si < samples * samples; ++si) {
+						float hoffset = (float)(si/samples)/max(samples-1, 1);
+						float woffset = (float)(si%samples)/max(samples-1, 1);
+						float x = (float)(wi + woffset)/w * 2.0 - 1.0;
+						float y = (float)(hi + hoffset)/h * 2.0 - 1.0;
+						Ray<scalar_t> ray = get_ray(h, w, y, x, right, front, up, cam_pos,fov);
+						vec3<scalar_t> tmp_color(0.0f);
+						if(ray_scene_trace(ray, scene, tmp_color)) {
+							color = color + tmp_color;
+							intersect = intersect || true;
+						}
+					}
+					if (intersect) {
+						color = color / float(samples * samples);
+						d_vis[bi][0][hi][wi] = color.x;
+						d_vis[bi][1][hi][wi] = color.y;
+						d_vis[bi][2][hi][wi] = color.z;
+					}
+            }
+        }
+    }
+
+} // namespace
+
+std::vector<torch::Tensor> plane_visualize_cuda(torch::Tensor planes, torch::Tensor camera, int h, int w){
+	const auto batch_size = planes.size(0);
+	const int threads = 512;
+	const dim3 blocks((w + threads - 1) / threads, (h+threads-1)/threads, batch_size);
+
+	torch::Tensor vis_tensor = torch::zeros({batch_size, 3, h, w}).to(planes);
+	AT_DISPATCH_FLOATING_TYPES(planes.type(), "plane_visualize_foward", ([&] {
+		plane_visualize_foward<scalar_t><<<blocks, threads>>>(
+			planes.packed_accessor64<scalar_t,2>(),
+			camera.packed_accessor64<scalar_t,2>(),
+			vis_tensor.packed_accessor64<scalar_t,4>()
+		);
+	}));
+
+	return {vis_tensor};
+}

PixHtLab-Src/Demo/PixhtLab/Torch_Render/setup.py

@@ -0,0 +1,29 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+setup(
+	name='hshadow',
+	ext_modules=[
+		CUDAExtension('hshadow', [
+			'hshadow_cuda.cpp',
+			'hshadow_cuda_kernel.cu',
+		])
+	],
+	cmdclass={
+		'build_ext': BuildExtension
+	}
+)
+
+# setup(
+# 	name='plane_visualize',
+# 	ext_modules=[
+# 		CUDAExtension('plane_visualize', [
+# 			'plane_visualize.cpp',
+# 			'plane_visualize_cuda.cu',
+# 		])
+# 	],
+# 	cmdclass={
+# 		'build_ext': BuildExtension
+# 	}
+# )

PixHtLab-Src/Demo/PixhtLab/Torch_Render/test_ground.py

@@ -0,0 +1,33 @@
+import time
+import torch
+import plane_visualize
+import matplotlib.pyplot as plt
+from PIL import Image
+from torchvision import transforms
+import os
+from os.path import join
+import numpy as np
+
+test_output = 'imgs/output'
+os.makedirs(test_output, exist_ok=True)
+device = torch.device("cuda:0")
+
+def test_ground():
+	fov, horizon = 120, 400 
+	camera = torch.tensor([[fov, horizon]])
+	planes = torch.tensor([[0.0, 0.0, 0.0, 0.0, 1.0, 0.0]])
+
+	camera = camera.repeat(5,1).float().to(device)
+	planes = planes.repeat(5,1).float().to(device) 
+
+	ground_vis = plane_visualize.forward(planes, camera, int(512), int(512))[0]
+	return ground_vis
+
+t = time.time()
+ground_vis = test_ground()
+print('{} s'.format(time.time() - t))
+batch = ground_vis.shape[0]
+for bi in range(batch):
+	img = ground_vis[bi].detach().cpu().numpy().transpose(1,2,0)
+	img = np.clip(img, 0.0, 1.0)
+	plt.imsave(join(test_output, 'ground_{}.png'.format(bi)),img)

PixHtLab-Src/Demo/PixhtLab/Torch_Render/test_hshadow.py

@@ -0,0 +1,130 @@
+import time
+import torch
+import hshadow
+import matplotlib.pyplot as plt
+from PIL import Image
+from torchvision import transforms
+import os
+from os.path import join
+import numpy as np
+from scipy.ndimage import uniform_filter
+
+test_output = 'imgs/output'
+os.makedirs(test_output, exist_ok=True)
+def test_shadow(rgb, mask, hmap, rechmap, light_pos):
+	h,w = rgb.shape[:2]
+	bb = torch.tensor([[0, h-1, 0, w-1]]).float().to(device)
+	start = time.time()
+	shadow = hshadow.forward(rgb, mask, bb, hmap, rechmap, light_pos)[0]
+	end = time.time()
+
+	print('Shadow rendering: {}s'.format(end-start))
+	res = (1.0-mask) * rgb * shadow + mask * rgb
+	return res, shadow
+
+def frenel_reflect(reflect_tensor, reflect_mask, fov, ref_ind):
+	# Use schilik approximation https://en.wikipedia.org/wiki/Schlick%27s_approximation
+	def deg2rad(deg):
+		return deg/180.0 * 3.1415926 
+
+	def img2cos(reflect_img, fov, horizon):
+		# Note, this factor needs calibration if we have camera parameters
+		b, c, h, w = reflect_img.shape
+		focal = 0.5 * h / np.tan(deg2rad(0.5 * fov))
+		fadding_map = torch.arange(0, h).unsqueeze(1).expand(h, w).unsqueeze(0).unsqueeze(0).repeat(b, c, 1, 1)
+		fadding_map = focal / torch.sqrt((fadding_map-horizon)**2 + focal **2)
+		return fadding_map.to(reflect_img)
+
+
+	ind = (1.0-ref_ind)/(1.0+ref_ind)
+	ind = ind ** 2
+	h = reflect_tensor.shape[2]
+	horizon = h * 0.7 
+	cos_map = img2cos(reflect_tensor, fov, horizon)
+	# fadding =  1.0 - (ind + (1.0-ind) * torch.pow(1.0-cos_map, 4))
+	b, c, h, w = reflect_tensor.shape
+	fadding = torch.linspace(3.0,0.0,h)[None, None, ..., None].repeat(b,c,1,w).to(reflect_tensor) ** 4
+	fadding = torch.clip(fadding, 0.0, 1.0)
+	plt.imsave('test_fadding.png', fadding[0].detach().cpu().numpy().transpose(1,2,0))
+	reflect_mask = reflect_mask.repeat(1,3,1,1)
+	return fadding * reflect_tensor * reflect_mask + (1.0-reflect_mask * fadding) * torch.ones_like(reflect_tensor) 
+
+def refine_boundary(output, filter=3):
+	return output
+	h,w,c = output.shape
+	for i in range(c):
+		output[...,i] = uniform_filter(output[...,i], size=filter)
+	return output
+
+def height_fadding(reflect, reflect_height, reflect_mask, fadding_factor):
+	def np_sigmoid(a):
+		return 1.0/(1.0+np.exp(-a))
+
+	h,w,c = reflect.shape
+	reflect_h = reflect_height/h
+	reflect_h = reflect_h/reflect_h.max()
+	fadding = (1.0-(np_sigmoid(reflect_h * fadding_factor)-0.5)* 2.0) * reflect_mask 
+	after_fadding = fadding * reflect + (1.0-fadding) * np.ones_like(reflect)
+	return  after_fadding
+
+def to_numpy(tensor):
+	return tensor[0].detach().cpu().numpy().transpose(1,2,0)
+
+def test_reflect(rgb, mask, hmap, rechmap, thresholds=1.5, fadding_factor=10.0):
+	b, c, h, w = rgb.shape
+	# thresholds = torch.tensor([[1.0 + i/b] for i in range(b)]).float().to(device)
+	thresholds = torch.tensor([[thresholds]]).float().to(device)
+	start = time.time()
+	reflect, reflect_height, reflect_mask = hshadow.reflection(rgb, mask, hmap, rechmap, thresholds)
+	end = time.time()
+	print('Reflection rendering: {}s'.format(end-start))
+
+	reflect, reflect_height, reflect_mask = to_numpy(reflect), to_numpy(reflect_height), to_numpy(reflect_mask)
+	# reflect = frenel_reflect(reflect, reflect_height, reflect_mask, 175, 0.9)
+	refine_reflect, refine_reflect_height = refine_boundary(reflect), refine_boundary(reflect_height)
+	refine_reflect_mask = reflect_mask
+	reflect = height_fadding(refine_reflect, refine_reflect_height, refine_reflect_mask, fadding_factor)
+	rgb, mask = to_numpy(rgb), to_numpy(mask)
+	res = (1.0-mask) * rgb * reflect + mask * rgb
+	return res, reflect
+
+
+def test_glossy_reflect(rgb, mask, hmap, rechmap, sample, glossy, fadding_factor=10.0):
+	b, c, h, w = rgb.shape
+	# thresholds = torch.tensor([[1.0 + i/b] for i in range(b)]).float().to(device)
+	start = time.time()
+	reflect = hshadow.glossy_reflection(rgb, mask, hmap, rechmap, sample, glossy)[0]
+	end = time.time()
+	print('Reflection rendering: {}s'.format(end-start))
+
+	reflect = to_numpy(reflect)
+	refine_reflect = refine_boundary(reflect)
+	rgb, mask = to_numpy(rgb), to_numpy(mask)
+	res = (1.0-mask) * rgb * reflect + mask * rgb
+	return res, reflect
+
+
+device = torch.device("cuda:0")
+to_tensor = transforms.ToTensor()
+# for i in range(1,5):
+i = 2
+if True:
+	prefix = 'canvas{}'.format(i)
+	rgb, mask, hmap = to_tensor(Image.open('imgs/{}_rgb.png'.format(prefix)).convert('RGB')).to(device), to_tensor(Image.open('imgs/{}_mask.png'.format(prefix)).convert('RGB'))[0:1].to(device), to_tensor(Image.open('imgs/{}_height.png'.format(prefix)).convert('RGB'))[0:1].to(device)
+	h, w = hmap.shape[1:]
+	hmap = hmap * h * 0.45
+
+	rechmap = torch.zeros_like(hmap)
+	rgb, mask, hmap, rechmap = rgb.unsqueeze(dim=0), mask.unsqueeze(dim=0), hmap.unsqueeze(dim=0), rechmap.unsqueeze(dim=0)
+	lightpos = torch.tensor([[300, -100, 200.0]]).to(device)
+	shadow_res, shadow = test_shadow(rgb, mask, hmap, rechmap, lightpos)
+	reflect_res, reflect = test_reflect(rgb, mask, hmap, rechmap, thresholds=2.5, fadding_factor=18.0)
+
+	glossy_reflect_res, glossy_reflect = test_glossy_reflect(rgb, mask, hmap, rechmap, sample=10, glossy=0.5, fadding_factor=18.0)
+
+	plt.imsave(join(test_output, prefix + "_shadow_final.png"), shadow_res[0].detach().cpu().numpy().transpose(1,2,0))
+	plt.imsave(join(test_output, prefix + "_shadow.png"), shadow[0].detach().cpu().numpy().transpose(1,2,0))
+	plt.imsave(join(test_output, prefix + "_reflect_final.png"), reflect_res)
+	plt.imsave(join(test_output, prefix + "_reflect.png"), reflect)
+	plt.imsave(join(test_output, prefix + "_glossy_reflect_final.png"), glossy_reflect_res)
+	plt.imsave(join(test_output, prefix + "_glossy_reflect.png"), glossy_reflect)

PixHtLab-Src/Demo/PixhtLab/camera.py

@@ -0,0 +1,246 @@
+import numpy as np
+import math
+from abc import ABC
+import copy
+
+class camera(ABC):
+    def __init__(self, hfov, h, w, height=100.0):
+        self.fov = hfov
+        self.h   = h
+        self.w   = w
+
+        self.ori_height = height
+        self.height     = copy.deepcopy(self.ori_height)
+        self.O          = np.array([0.0, self.height, 0.0]) # ray origianl
+
+
+    ######################################################################################
+    """ Abstraction
+    """
+    def align_horizon(self, cur_horizon):
+        raise NotImplementedError('Not implemented yet')
+
+    def C(self):
+        raise NotImplementedError('Not implemented yet')
+
+    def right(self):
+        raise NotImplementedError('Not implemented yet')
+
+    def up(self):
+        raise NotImplementedError('Not implemented yet')
+
+    ######################################################################################
+
+    def deg2rad(self, d):
+        return d / 180.0 * 3.1415925
+
+
+    def rad2deg(self, d):
+        return d / 3.1415925 * 180.0
+
+
+    def get_ray(self, xy):
+        """ Assume the center is on the top-left corner
+        """
+        u, v = xy
+        mat  = self.get_ABC_mat()
+        r    = np.dot(mat, np.array([u, v, 1.0]).T)
+
+        # r = r/np.sqrt(r @ r)
+        return r
+
+
+    def project(self, xyz):
+        relative = xyz - self.O
+
+        mat   = self.get_ABC_mat()
+        pp    = np.dot(np.linalg.inv(mat), relative)
+        pixel = np.array([pp[0]/pp[2], pp[1]/pp[2]])
+
+        return pixel
+
+
+    def xyh2w(self, xyh):
+        u, v, h = xyh
+
+        foot_xyh    = np.copy(xyh)
+        foot_xyh[1] = foot_xyh[1] + foot_xyh[2]
+        foot_xyh[2] = 0.0
+        fu, fv, fh  = foot_xyh
+
+        a   = self.right()
+        b   = -self.up()
+        c   = self.C()
+        mat = self.get_ABC_mat()
+
+        w = -self.height/(a[1] * fu + b[1] * fv + c[1])
+        return w
+
+
+    def xyh2xyz(self, xyh):
+        u, v, h = xyh
+
+        foot_xyh    = np.copy(xyh)
+        foot_xyh[1] = foot_xyh[1] + foot_xyh[2]
+        foot_xyh[2] = 0.0
+        fu, fv, fh  = foot_xyh
+
+        a   = self.right()
+        b   = -self.up()
+        c   = self.C()
+        mat = self.get_ABC_mat()
+
+        w = -self.height/(a[1] * fu + b[1] * fv + c[1])
+        # print('w: {} a*u + b * v + c: {}, b/c: {}/{}, fv: {}'.format(w, a[1] * fu + b[1] * fv + c[1], b, c, fv))
+        xyz = self.O + np.dot(mat, np.array([u, v, 1.0]).T) * w
+
+        # print('w: {}, -{}/{}'.format(w, self.height, a[1] * fu + b[1] * fv + c[1]))
+
+        return xyz
+
+
+    def xyz2xyh(self, xyz):
+        foot_xyz    = np.copy(xyz)
+        foot_xyz[1] = 0.0
+
+        foot_xy = self.project(foot_xyz)
+        xy      = self.project(xyz)
+
+        ret     = np.copy(xyz)
+        ret[:2] = xy
+        ret[2]  = foot_xy[1] - xy[1]
+
+        return ret
+
+
+    def get_ABC_mat(self):
+        a = self.right()
+        b = -self.up()
+        c = self.C()
+
+        mat = np.concatenate([a[:, None], b[:,None], c[:, None]], axis=1)
+        return mat
+
+
+class pitch_camera(camera):
+    """ Picth alignment camera
+    """
+    def __init__(self, hfov, h, w, height=100.0):
+        """
+           alignment algorithm:
+                  1. pitch alignment
+                  2. axis alignment
+        """
+        super().__init__(hfov, h, w, height)
+
+        self.ori_view   = np.array([0.0, 0.0, -1.0])
+        self.cur_view   = np.copy(self.ori_view)
+
+
+    def align_horizon(self, cur_horizon):
+        """ Given horizon, compute the camera pitch
+        """
+        ref_horizon  = self.h / 2
+        rel_distance = -(ref_horizon - cur_horizon)
+
+        focal = self.focal()
+        pitch = math.atan2(rel_distance, focal)
+
+        # construct a rotation matrix
+        c, s = np.cos(pitch), np.sin(pitch)
+        rot = np.array([[0, 0, 0], [0, c, -s], [0, s, c]])
+
+        # compute the new view vector
+        img_plane_view = self.ori_view * focal
+        img_plane_view = rot @ img_plane_view.T
+
+        self.cur_view = img_plane_view / math.sqrt(np.dot(img_plane_view, img_plane_view))
+
+    def C(self):
+        return self.view() * self.focal() - 0.5 * self.w * self.right() + 0.5 * self.h * self.up()
+
+
+    def right(self):
+        return np.array([1.0, 0.0, 0.0])
+
+
+    def up(self):
+        return np.cross(self.right(), self.view())
+
+
+    def focal(self):
+        focal = self.w * 0.5 / math.tan(self.deg2rad(self.fov * 0.5))
+        return focal
+
+
+    def view(self):
+        return self.cur_view
+
+
+
+class axis_camera(camera):
+    """ Axis alignment camera
+    """
+    def __init__(self, hfov, h, w, height=100.0):
+        super().__init__(hfov, h, w, height)
+
+        focal          = self.w * 0.5 / math.tan(self.deg2rad(self.fov * 0.5))
+        self.up_vec    = np.array([0.0, 1.0, 0.0])
+        self.right_vec = np.array([1.0, 0.0, 0.0])
+
+        self.ori_c = np.array([-0.5 * self.w, 0.5 * self.h, -focal])
+        self.c_vec = np.copy(self.ori_c)
+
+
+    def align_horizon(self, cur_horizon):
+        """ Given horizon, we move the axis to update the horizon
+            i.e. we need to change C
+        """
+        ref_horizon   = self.h // 2
+        delta_horizon = cur_horizon - ref_horizon
+        self.c_vec    = self.ori_c + delta_horizon * self.up()
+        # self.height   = self.ori_height + delta_horizon
+        # self.O        = np.array([0.0, self.height, 0.0])
+
+
+
+    def C(self):
+        return self.c_vec
+
+
+    def right(self):
+        return self.right_vec
+
+
+    def up(self):
+        return self.up_vec
+
+
+def test(ppc):
+    xyh = np.array([500, 500, 100.0])
+
+    proj_xyz = ppc.xyh2xyz(xyh)
+    proj_xyh = ppc.xyz2xyh(proj_xyz)
+
+    print('xyh: {}, proj xyz: {}, proj xyh: {}'.format(xyh, proj_xyz, proj_xyh))
+
+    # import pdb; pdb.set_trace()
+    new_horizon_list = [0, 100, 250, 400, 500]
+    new_horizon_list = [100, 250, 400, 500]
+
+    # import pdb; pdb.set_trace()
+    for cur_horizon in new_horizon_list:
+        ppc.align_horizon(cur_horizon)
+        test_xyh = np.array([500, cur_horizon, 0])
+        test_xyz = ppc.xyh2xyz(test_xyh)
+
+        # print('{} -> {} -> {}'.format(test_xyh, test_xyz, ppc.xyz2xyh(test_xyz)))
+        print('{} \t -> {} \t -> {}'.format(test_xyh, test_xyz, ppc.xyz2xyh(test_xyz)))
+
+
+if __name__ == '__main__':
+    p_camera = pitch_camera(90.0, 500, 500)
+    a_camera = axis_camera(90.0, 500, 500)
+
+    test(p_camera)
+    test(a_camera)

PixHtLab-Src/Demo/PixhtLab/gssn_demo.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+import torch
+import gradio as gr
+from torch import nn
+import numpy as np
+
+
+def render_btn_fn(mask, background, buffers, pitch, roll, softness):
+    print('Pitch and roll: {}, {}'.format(pitch, roll))
+    print('Mask, background, bufferss: {}, {}, {}'.format(mask.shape, background.shape, buffers.shape))
+    pass
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        mask_input = gr.Image(shape=(256, 256), image_mode="L", label="Mask")
+        bg_input   = gr.Image(shape=(256, 256), image_mode="RGB", label="Background")
+        buff_input = gr.Image(shape=(256, 256), image_mode="RGB", label="Buffers")
+
+    with gr.Row():
+        with gr.Column():
+            pitch_input    = gr.Slider(minimum=0, maximum=1, step=0.01, default=0.5, label="Pitch")
+            roll_input     = gr.Slider(minimum=0, maximum=1, step=0.01, default=0.5, label="Roll")
+            softness_input = gr.Slider(minimum=0, maximum=1, step=0.01, default=0.5, label="Softness")
+
+    render_btn = gr.Button(label="Render")
+    output     = gr.Image(shape=(256, 256), image_mode="RGB", label="Output")
+
+    render_btn.click(render_btn_fn, inputs=[mask_input, bg_input, buff_input, pitch_input, roll_input, softness_input], outputs=output)
+
+
+demo.launch()

PixHtLab-Src/Demo/PixhtLab/hshadow_render.py

@@ -0,0 +1,268 @@
+import time
+import torch
+import hshadow
+# import plane_visualize
+import numpy as np
+from torchvision import transforms
+from scipy.ndimage import uniform_filter
+from ShadowStyle.inference import inference_shadow
+import cv2
+import matplotlib.pyplot as plt
+from utils import *
+from GSSN.inference_shadow import SSN_Infernece
+
+device     = torch.device("cuda:0")
+to_tensor  = transforms.ToTensor()
+model      = inference_shadow.init_models('/home/ysheng/Documents/Research/GSSN/HardShadow/qtGUI/weights/human_baseline_all_21-July-04-52-AM.pt')
+# GSSN_model = SSN_Infernece('GSSN/weights/0000000700.pt')
+GSSN_model = SSN_Infernece('/home/ysheng/Documents/Research/GSSN/HardShadow/qtGUI/GSSN/weights/only_shadow/0000000200.pt')
+
+def crop_mask(mask):
+        hnon, wnon = np.nonzero(mask)
+        aabb = (hnon.min(), hnon.max(), wnon.min(), wnon.max())
+        return aabb
+
+def norm_output(np_img):
+        return np.clip(cv2.normalize(np_img, None, 0.0, 1.0, cv2.NORM_MINMAX),0.0,1.0)
+
+def padding(mask, shadow, mask_aabb, shadow_aabb, final_shape=(512, 512)):
+        mh, mhh, mw, mww = mask_aabb
+        sh, shh, sw, sww = shadow_aabb
+        cropped_mask, cropped_shadow = mask[mh:mhh, mw:mww], shadow[sh:shh, sw:sww]
+        global_h, global_w = mask.shape[:2]
+        h, w, c, sc = *cropped_mask.shape, shadow.shape[2]
+        fract = 0.4
+        if h > w:
+                newh = int(final_shape[0]*fract)
+                neww = int(newh/h*w)
+        else:
+                neww = int(final_shape[1]*fract)
+                newh = int(neww/w*h)
+
+        small_mask = cv2.resize(cropped_mask, (neww, newh), interpolation=cv2.INTER_AREA)
+        if len(small_mask.shape) == 2:
+                small_mask = small_mask[...,np.newaxis]
+
+        mask_ret, shadow_ret = np.zeros((final_shape[0], final_shape[1], c)),np.ones((final_shape[0], final_shape[1], sc))
+        paddingh, paddingw = 10, (final_shape[0]-neww)//2
+        mask_lpos = (paddingh, paddingw)
+        mask_ret = overlap_replace(mask_ret, small_mask, mask_lpos)
+
+        # padding shadow
+        hscale, wscale = newh/h, neww/w
+        newsh, newsw = int((shh-sh) * hscale), int((sww-sw) * wscale)
+        small_shadow = cv2.resize(cropped_shadow, (newsw, newsh), interpolation=cv2.INTER_AREA)
+
+        if len(small_shadow.shape) == 2:
+                small_shadow = small_shadow[...,np.newaxis]
+
+
+        loffseth, loffsetw = int((sh-mh)*hscale), int((sw-mw)*wscale)
+        shadow_lpos = (paddingh + loffseth, paddingw + loffsetw)
+        shadow_ret = overlap_replace(shadow_ret, small_shadow, shadow_lpos)
+
+        # return mask_ret, shadow_ret[...,0:1], [mask_aabb, mask_lpos, hscale, wscale, final_shape, mask.shape[0], mask.shape[1]]
+        return mask_ret, shadow_ret, [mask_aabb, mask_lpos, hscale, wscale, final_shape, mask.shape[0], mask.shape[1]]
+
+
+def transform_input(mask, hardshadow):
+        """ Note, trans_info marks the AABBs, and scaling factors
+        """
+        mask_aabb, shadow_aabb = crop_mask(mask[...,0]), crop_mask(hardshadow[...,0])
+        # import pdb; pdb.set_trace()
+        cmask, cshadow, trans_info = padding(mask, hardshadow, mask_aabb, shadow_aabb)
+        return cmask.transpose(2,0,1)[np.newaxis,...], 1.0 - cshadow.transpose(2,0,1)[np.newaxis, ...], trans_info
+
+
+def transform_output(softshadow, trans_info):
+        mask_aabb, mask_lpos, hscale, wscale, final_shape, h, w = trans_info
+        # import pdb; pdb.set_trace()
+        ret, gsh, gsw = np.zeros((h,w,1)), int(final_shape[0]/hscale), int(final_shape[1]/wscale)
+        global_shadow = cv2.resize(softshadow[0,0], (gsw, gsh))
+
+        # global start = global_mask_aabb - (local_mask_start)/scaling
+        mh, mw, mask_lh, mask_lw = mask_aabb[0], mask_aabb[2], mask_lpos[0], mask_lpos[1]
+        starth, startw = int(mh - mask_lh / hscale), int(mw - mask_lw / wscale)
+        ret = norm_output(overlap_replace(ret, global_shadow[...,np.newaxis], (starth, startw)))
+        if len(ret.shape) == 2:
+                ret = ret[..., np.newaxis]
+
+        return 1.0-ret.repeat(3,axis=2)
+
+def style_hardshadow(mask, hardshadow, softness):
+        mask_net, hardshadow_net, trans_info = transform_input(mask, hardshadow)
+        netsoftshadow = inference_shadow.net_render_np(model, mask_net, hardshadow_net, softness, 0.0)
+        softshadow = transform_output(netsoftshadow, trans_info)
+
+        return softshadow, (norm_output(mask_net[0,0]), norm_output(hardshadow_net[0,0]), norm_output(netsoftshadow[0,0]))
+
+def gssn_shadow(mask, pixel_height, shadow_channels, softness):
+        # mask_net, hardshadow_net, trans_info = transform_input(mask, shadow_channels)
+
+        mask_aabb, shadow_aabb                 = crop_mask(mask[...,0]), crop_mask(shadow_channels[...,0])
+        ph_channel, hardshadow_net, trans_info = padding(pixel_height, shadow_channels, mask_aabb, shadow_aabb)
+
+        ph_channel     = ph_channel/512.0
+        hardshadow_net = 1.0-hardshadow_net
+        input_np       = np.concatenate([ph_channel, hardshadow_net], axis=2)
+
+        # import pdb; pdb.set_trace()
+
+        netsoftshadow = np.clip(GSSN_model.render_ss(input_np, softness), 0.0, 1.0)
+        netsoftshadow = netsoftshadow.transpose((2,0,1))[None, ...]
+        softshadow    = transform_output(netsoftshadow, trans_info)
+
+        return softshadow
+
+
+def proj_ground(p, light_pos):
+        tmpp = p.copy()
+
+        t = (0-tmpp[2])/(light_pos[:, 2:3]-tmpp[2]+1e-6)
+        tmpp = (1.0-t) * tmpp[:2] + t * light_pos[:, :2]
+        return tmpp
+
+def proj_bb(mask, hmap, light_pos, mouse_pos):
+        tmp_lights = light_pos.copy()
+        if len(light_pos.shape) == 1:
+                tmp_lights = tmp_lights[..., np.newaxis]
+
+        # bb -> four points
+        highest = hmap.max()
+        highest_h, highest_w = list(np.unravel_index(np.argmax(hmap), hmap.shape))
+        hbb, wbb = np.nonzero(mask)
+        h, hh, w, ww = hbb.min(), hbb.max(), wbb.min(), wbb.max()
+        bb0, bb1, bb2, bb3 = np.array([w, h, hmap.max()]), np.array([ww, h, hmap.max()]), np.array([w, hh, 0]), np.array([ww, hh, 0])
+
+        # compute projection for the four points
+        tmp_lights = tmp_lights.transpose(1,0)
+        bb0, bb1, bb2, bb3 = proj_ground(bb0, tmp_lights), proj_ground(bb1, tmp_lights), proj_ground(bb2, tmp_lights), proj_ground(bb3, tmp_lights)
+
+        batch = len(tmp_lights)
+        new_bb = np.zeros((batch, 4))
+        for i in range(batch):
+                new_bb[i, 0] = min([bb0[i, 1], bb1[i,1], bb2[i, 1], bb3[i, 1], mouse_pos[1], h]) # h
+                new_bb[i, 1] = max([bb0[i, 1], bb1[i,1], bb2[i, 1], bb3[i, 1], mouse_pos[1], hh])
+                new_bb[i, 2] = min([bb0[i, 0], bb1[i,0], bb2[i, 0], bb3[i, 0], mouse_pos[0], w]) # w
+                new_bb[i, 3] = max([bb0[i, 0], bb1[i,0], bb2[i, 0], bb3[i, 0], mouse_pos[0], ww])
+
+        return new_bb
+
+def to_torch_device(np_img):
+        if len(np_img.shape) == 3:
+                return to_tensor(np_img).float().unsqueeze(dim=0).contiguous().to(device)
+        else:
+                return torch.from_numpy(np_img).float().contiguous().to(device)
+
+def hshadow_render(rgb, mask, hmap, rechmap, light_pos, mouse_pos):
+        """ Heightmap Shadow Rendering
+                rgb:            H x W x e
+                mask:           H x W x 1
+                hmap:           H x W x 1
+                rechmap:        H x W x 1
+                light_pos:  (3,B)
+                return:
+                        shadow masking
+        """
+
+        hbb, wbb = np.nonzero(mask[...,0])
+        # speed optimization
+        bb = proj_bb(mask[...,0], hmap[...,0], light_pos, mouse_pos)
+
+        # import pdb; pdb.set_trace()
+        if len(light_pos.shape) == 1:
+                light_pos_d = torch.from_numpy(light_pos).to(device).unsqueeze(dim=0).float()
+                rgb_d, mask_d, hmap_d, rechmap_d = to_torch_device(rgb), to_torch_device(mask), to_torch_device(hmap), to_torch_device(rechmap)
+                bb_d = torch.from_numpy(bb).float().to(device)
+                batch = 1
+        else:
+                light_pos_d = torch.from_numpy(np.ascontiguousarray(light_pos.transpose(1,0))).float().to(device)
+                batch = len(light_pos_d)
+                h,w = rgb.shape[:2]
+                rgb_d = to_torch_device(np.repeat(rgb[np.newaxis,...].transpose(0,3,1,2), batch, axis=0))
+                mask_d = to_torch_device(np.repeat(mask[np.newaxis,...].transpose(0,3,1,2), batch, axis=0))
+                hmap_d = to_torch_device(np.repeat(hmap[np.newaxis,...].transpose(0,3,1,2), batch, axis=0))
+                rechmap_d = to_torch_device(np.repeat(rechmap[np.newaxis,...].transpose(0,3,1,2), batch, axis=0))
+                bb_d = torch.from_numpy(np.ascontiguousarray(bb)).float().to(device)
+
+        shadow = hshadow.forward(rgb_d, mask_d, bb_d, hmap_d, rechmap_d, light_pos_d)
+        # mask_top_pos = list(np.unravel_index(np.argmax(hmap), hmap.shape))
+        # x,y = mask_top_pos[1], mask_top_pos[0]
+        # mh = hmap[y,x,0]
+        # light_top_d = light_pos_d - torch.tensor([[x,y,mh]]).to(light_pos_d)
+        # weights = torch.abs(light_top_d[:,2]/torch.sqrt((light_top_d[:,0] **2 + light_top_d[:,1] **2)))
+        # print('weights: ', weights)
+        # weights = (weights)/weights.sum()
+
+        # print(weights.shape, shadow[0].shape)
+        # flipped = (weights[...,None, None,None] * (1.0-shadow[0])).sum(dim=0, keepdim=True)
+        # shadow = shadow[0].sum(dim=0, keepdim=True)/len(shadow[0])
+        # return (1.0-flipped)[0].detach().cpu().numpy().transpose(1,2,0)
+
+        shadow = shadow[0].sum(dim=0, keepdim=True)/len(shadow[0])
+        return shadow[0].detach().cpu().numpy().transpose(1,2,0)
+
+def refine_shadow(shadow, intensity=0.6, filter=5):
+        shadow[...,0] = uniform_filter(shadow[...,0], size=filter)
+        shadow[...,1] = uniform_filter(shadow[...,1], size=filter)
+        shadow[...,2] = uniform_filter(shadow[...,2], size=filter)
+        return 1.0 - (1.0-shadow) * intensity
+
+def render_ao(rgb, mask, hmap):
+        rechmap = np.zeros_like(hmap)
+        hbb, wbb = np.nonzero(mask[...,0])
+        # light_pos = np.array([hbb.min(), (wbb.min() + wbb.max()) * 0.8, -100000])
+        light_pos = np.array([-1300.10811363, -46999.86253089, 46486.73121776])
+        mouse_pos = light_pos
+
+        shadow = hshadow_render(rgb, mask, hmap, rechmap, light_pos, mouse_pos)
+        softshadow = style_hardshadow(mask, shadow[..., :1], 0.45)[0]
+        softshadow = refine_shadow(softshadow)
+        return softshadow
+
+def ao_composite(rgb, mask, hmap, rechmap, light_pos, mouse_pos):
+        # shadow = hshadow_render(rgb, mask, hmap, rechmap, light_pos, mouse_pos)
+        # softshadow = style_hardshadow(mask, shadow, 0.45)[0]
+        # softshadow = refine_shadow(softshadow)
+
+        softshadow = render_ao(rgb, mask, hmap)
+        mask_ = np.repeat(mask, 3, axis=2)
+        return (1.0-mask_) * softshadow * rgb + mask_ * rgb, softshadow.copy()
+
+
+def render_shadow(rgb, mask, hmap, rechmap, light_pos, mouse_pos, softness, shadow_intensity=0.6):
+        shadow = hshadow_render(rgb, mask, hmap, rechmap, light_pos, mouse_pos)
+
+        if softness is not None:
+                shadow, dbgs = style_hardshadow(mask, shadow[..., :1], softness)
+        else:
+                dbgs = None
+
+        shadow = refine_shadow(shadow, intensity=shadow_intensity)
+        return shadow, dbgs
+
+
+def hshadow_composite(rgb, mask, hmap, rechmap, light_pos, mouse_pos, softness, shadow_intensity=0.6):
+        """ Shadow Rendering and Composition
+                rgb:            H x W x 3
+                mask:           H x W x 1
+                hmap:           H x W x 1
+                rechmap:        H x W x 1
+                light_pos:  [x,y,h]
+                return:
+                        Compositied image
+        """
+        shadow, dbgs = render_shadow(rgb, mask, hmap, rechmap, light_pos, mouse_pos, softness, shadow_intensity)
+        mask_ = np.repeat(mask, 3, axis=2)
+        return (1.0-mask_) * shadow * rgb + mask_ * rgb, shadow.copy(), dbgs
+
+# def vis_horizon(fov, horizon, h, w):
+#         # fov, horizon = 120, 400
+#         camera = torch.tensor([[fov, horizon]])
+#         planes = torch.tensor([[0.0, 0.0, 0.0, 0.0, 1.0, 0.0]])
+
+#         camera = camera.float().to(device)
+#         planes = planes.float().to(device)
+
+#         ground_vis = plane_visualize.forward(planes, camera, h, w)[0]
+#         return 1.0-ground_vis[0].detach().cpu().numpy().transpose(1,2,0)