Shading Residual: Add Albedo residual 2 and read_csv optional data in data loading

Author: bhushan23

3_Shading_Residual/data_loading.py

@@ -119,25 +119,7 @@ def get_sfsnet_dataset(syn_dir=None, read_from_csv=None, read_celeba_csv=None, r
     face    = []
     depth   = []
 
-    if read_from_csv is None:
-        for img in sorted(glob.glob(syn_dir + '*/*_albedo_*')):
-            albedo.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_face_*')):
-            face.append(img)    
-
-        for img in sorted(glob.glob(syn_dir + '*/*_normal_*')):
-            normal.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_depth_*')):
-            depth.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_mask_*')):
-            mask.append(img)
-
-        for img in sorted(glob.glob(syn_dir + '*/*_light_*.txt')):
-            sh.append(img)
-    else:
+    if read_from_csv is not None:
         df = pd.read_csv(read_from_csv)
         if read_first is not None and len(df) > read_first:
             df = df.sample(read_first)
@@ -155,17 +137,17 @@ def get_sfsnet_dataset(syn_dir=None, read_from_csv=None, read_celeba_csv=None, r
         for _, v in name_to_list.items():
             v[:] = [syn_dir + el for el in v]
 
-        # Merge Synthesized Celeba dataset for Psedo-Supervised training
-        if read_celeba_csv is not None:
-            df = pd.read_csv(read_celeba_csv)
-            if read_first is not None and len(df) > read_first:
-                df = df.sample(read_first)
-            albedo += list(df['albedo'])
-            face   += list(df['face'])
-            normal += list(df['normal'])
-            depth  += list(df['depth'])
-            mask   += list(df['mask'])
-            sh     += list(df['light'])
+    # Merge Synthesized Celeba dataset for Psedo-Supervised training
+    if read_celeba_csv is not None:
+        df = pd.read_csv(read_celeba_csv)
+        if read_first is not None and len(df) > read_first:
+            df = df.sample(read_first)
+        albedo += list(df['albedo'])
+        face   += list(df['face'])
+        normal += list(df['normal'])
+        depth  += list(df['depth'])
+        mask   += list(df['mask'])
+        sh     += list(df['light'])
 
     assert(len(albedo) == len(face) == len(normal) == len(depth) == len(mask) == len(sh))
     dataset_size = len(albedo)

3_Shading_Residual/models.py

@@ -278,6 +278,7 @@ def __init__(self):
         self.normal_gen_model      = NormalGenerationNet()
         self.albedo_residual_model = AlbedoResidualBlock()
         self.albedo_gen_model      = AlbedoGenerationNet()
+        self.albedo_residual_model_2 = AlbedoResidualBlock()
         self.light_estimator_model = LightEstimator()
         self.shading_residual_model = ShadingResidualEstimator()
 
@@ -298,22 +299,25 @@ def forward(self, face):
 
         # 3 a. Generate Normal
         predicted_normal = self.normal_gen_model(out_normal_features)
-        # 3 b. Generate Albedo
-        predicted_albedo = self.albedo_gen_model(out_albedo_features)
-        # 3 c. Estimate lighting
+        # 3 b. Estimate lighting
         # First, concat conv, normal and albedo features over channels dimension
         all_features = torch.cat((out_features, out_normal_features, out_albedo_features), dim=1)
         # Predict SH
         predicted_sh = self.light_estimator_model(all_features)
 
-        # 4. Generate shading
+        # 4. Collect new albedo features
+        out_albedo_features_2 = self.albedo_residual_model_2(out_features)
+        new_features = torch.cat((out_features, out_normal_features, out_albedo_features_2), dim=1)
+        predicted_albedo = self.albedo_gen_model(out_albedo_features_2)
+        
+        # 5. Generate shading
         out_shading = get_shading(predicted_normal, predicted_sh)
 
-        # 5. Get Shading Residual
-        shading_residual = self.shading_residual_model(all_features)
+        # 6. Get Shading Residual
+        shading_residual = self.shading_residual_model(new_features)
         updated_shading = out_shading + shading_residual
 
-        # 6. Reconstruction of image
+        # 7. Reconstruction of image
         out_recon = reconstruct_image(updated_shading, predicted_albedo)
 
         return predicted_normal, predicted_albedo, predicted_sh, out_shading, shading_residual, updated_shading, out_recon
@@ -509,6 +513,71 @@ def load_model_from_pretrained(src_model, dst_model):
     dst_model['albedo_gen_model.conv2.1.running_var'] = src_model['aconv2.conv.1.running_var']
     dst_model['albedo_gen_model.conv3.weight'] = src_model['aout.weight']
     dst_model['albedo_gen_model.conv3.bias'] = src_model['aout.bias']
+    dst_model['albedo_residual_model_2.block1.res.0.weight'] = src_model['ares1.res.0.weight']
+    dst_model['albedo_residual_model_2.block1.res.0.bias'] = src_model['ares1.res.0.bias']
+    dst_model['albedo_residual_model_2.block1.res.0.running_mean'] = src_model['ares1.res.0.running_mean']
+    dst_model['albedo_residual_model_2.block1.res.0.running_var'] = src_model['ares1.res.0.running_var']
+    dst_model['albedo_residual_model_2.block1.res.2.weight'] = src_model['ares1.res.2.weight']
+    dst_model['albedo_residual_model_2.block1.res.2.bias'] = src_model['ares1.res.2.bias']
+    dst_model['albedo_residual_model_2.block1.res.3.weight'] = src_model['ares1.res.3.weight']
+    dst_model['albedo_residual_model_2.block1.res.3.bias'] = src_model['ares1.res.3.bias']
+    dst_model['albedo_residual_model_2.block1.res.3.running_mean'] = src_model['ares1.res.3.running_mean']
+    dst_model['albedo_residual_model_2.block1.res.3.running_var'] = src_model['ares1.res.3.running_var']
+    dst_model['albedo_residual_model_2.block1.res.5.weight'] = src_model['ares1.res.5.weight']
+    dst_model['albedo_residual_model_2.block1.res.5.bias'] = src_model['ares1.res.5.bias']
+    dst_model['albedo_residual_model_2.block2.res.0.weight'] = src_model['ares2.res.0.weight']
+    dst_model['albedo_residual_model_2.block2.res.0.bias'] = src_model['ares2.res.0.bias']
+    dst_model['albedo_residual_model_2.block2.res.0.running_mean'] = src_model['ares2.res.0.running_mean']
+    dst_model['albedo_residual_model_2.block2.res.0.running_var'] = src_model['ares2.res.0.running_var']
+    dst_model['albedo_residual_model_2.block2.res.2.weight'] = src_model['ares2.res.2.weight']
+    dst_model['albedo_residual_model_2.block2.res.2.bias'] = src_model['ares2.res.2.bias']
+    dst_model['albedo_residual_model_2.block2.res.3.weight'] = src_model['ares2.res.3.weight']
+    dst_model['albedo_residual_model_2.block2.res.3.bias'] = src_model['ares2.res.3.bias']
+    dst_model['albedo_residual_model_2.block2.res.3.running_mean'] = src_model['ares2.res.3.running_mean']
+    dst_model['albedo_residual_model_2.block2.res.3.running_var'] = src_model['ares2.res.3.running_var']
+    dst_model['albedo_residual_model_2.block2.res.5.weight'] = src_model['ares2.res.5.weight']
+    dst_model['albedo_residual_model_2.block2.res.5.bias'] = src_model['ares2.res.5.bias']
+    dst_model['albedo_residual_model_2.block3.res.0.weight'] = src_model['ares3.res.0.weight']
+    dst_model['albedo_residual_model_2.block3.res.0.bias'] = src_model['ares3.res.0.bias']
+    dst_model['albedo_residual_model_2.block3.res.0.running_mean'] = src_model['ares3.res.0.running_mean']
+    dst_model['albedo_residual_model_2.block3.res.0.running_var'] = src_model['ares3.res.0.running_var']
+    dst_model['albedo_residual_model_2.block3.res.2.weight'] = src_model['ares3.res.2.weight']
+    dst_model['albedo_residual_model_2.block3.res.2.bias'] = src_model['ares3.res.2.bias']
+    dst_model['albedo_residual_model_2.block3.res.3.weight'] = src_model['ares3.res.3.weight']
+    dst_model['albedo_residual_model_2.block3.res.3.bias'] = src_model['ares3.res.3.bias']
+    dst_model['albedo_residual_model_2.block3.res.3.running_mean'] = src_model['ares3.res.3.running_mean']
+    dst_model['albedo_residual_model_2.block3.res.3.running_var'] = src_model['ares3.res.3.running_var']
+    dst_model['albedo_residual_model_2.block3.res.5.weight'] = src_model['ares3.res.5.weight']
+    dst_model['albedo_residual_model_2.block3.res.5.bias'] = src_model['ares3.res.5.bias']
+    dst_model['albedo_residual_model_2.block4.res.0.weight'] = src_model['ares4.res.0.weight']
+    dst_model['albedo_residual_model_2.block4.res.0.bias'] = src_model['ares4.res.0.bias']
+    dst_model['albedo_residual_model_2.block4.res.0.running_mean'] = src_model['ares4.res.0.running_mean']
+    dst_model['albedo_residual_model_2.block4.res.0.running_var'] = src_model['ares4.res.0.running_var']
+    dst_model['albedo_residual_model_2.block4.res.2.weight'] = src_model['ares4.res.2.weight']
+    dst_model['albedo_residual_model_2.block4.res.2.bias'] = src_model['ares4.res.2.bias']
+    dst_model['albedo_residual_model_2.block4.res.3.weight'] = src_model['ares4.res.3.weight']
+    dst_model['albedo_residual_model_2.block4.res.3.bias'] = src_model['ares4.res.3.bias']
+    dst_model['albedo_residual_model_2.block4.res.3.running_mean'] = src_model['ares4.res.3.running_mean']
+    dst_model['albedo_residual_model_2.block4.res.3.running_var'] = src_model['ares4.res.3.running_var']
+    dst_model['albedo_residual_model_2.block4.res.5.weight'] = src_model['ares4.res.5.weight']
+    dst_model['albedo_residual_model_2.block4.res.5.bias'] = src_model['ares4.res.5.bias']
+    dst_model['albedo_residual_model_2.block5.res.0.weight'] = src_model['ares5.res.0.weight']
+    dst_model['albedo_residual_model_2.block5.res.0.bias'] = src_model['ares5.res.0.bias']
+    dst_model['albedo_residual_model_2.block5.res.0.running_mean'] = src_model['ares5.res.0.running_mean']
+    dst_model['albedo_residual_model_2.block5.res.0.running_var'] = src_model['ares5.res.0.running_var']
+    dst_model['albedo_residual_model_2.block5.res.2.weight'] = src_model['ares5.res.2.weight']
+    dst_model['albedo_residual_model_2.block5.res.2.bias'] = src_model['ares5.res.2.bias']
+    dst_model['albedo_residual_model_2.block5.res.3.weight'] = src_model['ares5.res.3.weight']
+    dst_model['albedo_residual_model_2.block5.res.3.bias'] = src_model['ares5.res.3.bias']
+    dst_model['albedo_residual_model_2.block5.res.3.running_mean'] = src_model['ares5.res.3.running_mean']
+    dst_model['albedo_residual_model_2.block5.res.3.running_var'] = src_model['ares5.res.3.running_var']
+    dst_model['albedo_residual_model_2.block5.res.5.weight'] = src_model['ares5.res.5.weight']
+    dst_model['albedo_residual_model_2.block5.res.5.bias'] = src_model['ares5.res.5.bias']
+    dst_model['albedo_residual_model_2.bn1.weight'] = src_model['areso.0.weight']
+    dst_model['albedo_residual_model_2.bn1.bias'] = src_model['areso.0.bias']
+    dst_model['albedo_residual_model_2.bn1.running_mean'] = src_model['areso.0.running_mean']
+    dst_model['albedo_residual_model_2.bn1.running_var'] = src_model['areso.0.running_var']
+    
     dst_model['light_estimator_model.conv1.0.weight'] = src_model['lconv.conv.0.weight']
     dst_model['light_estimator_model.conv1.0.bias'] = src_model['lconv.conv.0.bias']
     dst_model['light_estimator_model.conv1.1.weight'] = src_model['lconv.conv.1.weight']