Merge pull request #54 from NVlabs/FFHQ_preprocess

Ffhq preprocess

dataset_preprocessing/ffhq/3dface2idr_mat.py

@@ -21,6 +21,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--in_root', type=str, default="", help='process folder')
+parser.add_argument('--out_path', type=str, default="cameras.json", help='output filename')
 args = parser.parse_args()
 in_root = args.in_root
 
@@ -73,9 +74,8 @@
     out = {}
     out["intrinsics"] = K
     out["pose"] = pose
-    outAll[src_filename.replace(".mat", ".jpg")] = out
+    outAll[src_filename.replace(".mat", ".png")] = out
 
 
-dst = os.path.join(in_root, "cameras.json")
-with open(dst, "w") as outfile:
+with open(args.out_path, "w") as outfile:
     json.dump(outAll, outfile)

dataset_preprocessing/ffhq/crop_images_in_the_wild.py

@@ -10,7 +10,6 @@
 
 import argparse
 import os
-import json
 from preprocess import align_img
 from PIL import Image
 import numpy as np
@@ -58,7 +57,4 @@
         im_cropped = im_high.crop((left, upper, right,lower))
         im_cropped = im_cropped.resize((output_size, output_size), resample=Image.LANCZOS)
         out_path = os.path.join(out_dir, img_file.split(".")[0] + ".png")
-        im_cropped.save(out_path)
-
-
-
+        im_cropped.save(out_path)

dataset_preprocessing/ffhq/preprocess_face_cameras.py

@@ -20,6 +20,12 @@
 import os
 from tqdm import tqdm
 import argparse
+import torch
+import sys
+sys.path.append('../../eg3d')
+from camera_utils import create_cam2world_matrix
+
+COMPRESS_LEVEL=0
 
 def fix_intrinsics(intrinsics):
     intrinsics = np.array(intrinsics).copy()
@@ -35,6 +41,7 @@ def fix_intrinsics(intrinsics):
     assert intrinsics[2,1] == 0
     return intrinsics
 
+# For our recropped images, with correction
 def fix_pose(pose):
     COR = np.array([0, 0, 0.175])
     pose = np.array(pose).copy()
@@ -43,13 +50,23 @@ def fix_pose(pose):
     pose[:3, 3] = direction * 2.7 + COR
     return pose
 
+# Used in original submission
 def fix_pose_orig(pose):
     pose = np.array(pose).copy()
     location = pose[:3, 3]
     radius = np.linalg.norm(location)
     pose[:3, 3] = pose[:3, 3]/radius * 2.7
     return pose
 
+# Used for original crop images
+def fix_pose_simplify(pose):
+    cam_location = torch.tensor(pose).clone()[:3, 3]
+    normalized_cam_location = torch.nn.functional.normalize(cam_location - torch.tensor([0, 0, 0.175]), dim=0)
+    camera_view_dir = - normalized_cam_location
+    camera_pos = 2.7 * normalized_cam_location + np.array([0, 0, 0.175])
+    simple_pose_matrix = create_cam2world_matrix(camera_view_dir.unsqueeze(0), camera_pos.unsqueeze(0))[0]
+    return simple_pose_matrix.numpy()
+
 def flip_yaw(pose_matrix):
     flipped = pose_matrix.copy()
     flipped[0, 1] *= -1
@@ -85,6 +102,8 @@ def flip_yaw(pose_matrix):
             pose = fix_pose(pose)
         elif args.mode == 'orig':
             pose = fix_pose_orig(pose)
+        elif args.mode == 'simplify':
+            pose = fix_pose_simplify(pose)
         else:
             assert False, "invalid mode"
         intrinsics = fix_intrinsics(intrinsics)

dataset_preprocessing/ffhq/preprocess_in_the_wild.py

@@ -28,17 +28,17 @@
 os.system(command)
 os.chdir('..')
 
-# convert the pose to our format
-command = "python 3dface2idr_mat.py --in_root Deep3DFaceRecon_pytorch/checkpoints/pretrained/results/epoch_20_000000"
+# crop out the input image
+command = "python crop_images_in_the_wild.py --indir=" + args.indir
 print(command)
 os.system(command)
-# additional correction to match the submission version
-command = "python preprocess_cameras.py --source Deep3DFaceRecon_pytorch/checkpoints/pretrained/results/epoch_20_000000 --mode orig"
+
+# convert the pose to our format
+command = f"python 3dface2idr_mat.py --in_root Deep3DFaceRecon_pytorch/checkpoints/pretrained/results/{out_folder}/epoch_20_000000 --out_path {os.path.join(args.indir, 'crop', 'cameras.json')}"
 print(command)
 os.system(command)
 
-
-# crop out the input image
-command = "python crop_images_in_the_wild.py --indir=" + args.indir
+# additional correction to match the submission version
+command = f"python preprocess_face_cameras.py --source {os.path.join(args.indir, 'crop')} --dest {out_folder} --mode orig"
 print(command)
-os.system(command)
+os.system(command)