render_gt_with_mask.py

import os
import pyredner
import torch
import imageio
import cv2
from difflight.tools.render_envmap import render_envmap

root_folder = '/gel/usr/heweb4/datasets/LavalIndoor/params'
masks_folder = '/gel/usr/heweb4/datasets/LavalIndoor/maskPanoramas/'
input_folder_exr = '/gel/usr/heweb4/datasets/LavalIndoor/1942x971/'

files = os.listdir(root_folder)

camera = pyredner.Camera(position=torch.tensor([0,.75,.0]),
            look_at=torch.tensor([.0,.0,0.0]),
            up=torch.tensor([.0,.0,-1.0]),
            resolution = (128, 128), 
            fov=torch.tensor([35.0]))

axis_correction = [1,1,-1]

objects = pyredner.load_obj('difflight/objs/scene.obj', return_objects=True)[0]
objects_shape = pyredner.Shape(vertices = objects.vertices, indices = objects.indices, material_id = 0)

mat_diff_01 = pyredner.Material(diffuse_reflectance = torch.tensor([1.,1.,1.], device = pyredner.get_device()))
materials = [mat_diff_01]

for file in files:
    print(file[:-4])

    masks = (imageio.imread(os.path.join(masks_folder, file[:-3]+'png'))[:,:,1]).astype('uint8')
    masks_count = masks.max()+1
    gt_filename = os.path.join(input_folder_exr, file[:-3]+'exr')
    envmap_exr = pyredner.imread(gt_filename)

    # render each light individually
    for i in range(masks_count):
        mask = (masks == i).astype('uint8')
        mask = torch.from_numpy(cv2.merge([mask, mask, mask]))
        
        # render
        target = render_envmap(camera, [objects_shape], materials=materials, envmap=envmap_exr*mask, max_bounces=1, num_samples=128) 
        print('this light has energy of: ', target.cpu().numpy().sum())