add save_obj with textures

photometric_fitting.py

@@ -153,11 +153,13 @@ def optimize(self, images, landmarks, image_masks, savefolder=None):
                 shape_images = self.render.render_shape(vertices, trans_vertices, images)
                 grids['shape'] = torchvision.utils.make_grid(
                     F.interpolate(shape_images[visind], [224, 224])).detach().float().cpu()
-                # grids['tex'] = torchvision.utils.make_grid(F.interpolate(albedos[visind], [224, 224])).detach().cpu()
 
+
+                # grids['tex'] = torchvision.utils.make_grid(F.interpolate(albedos[visind], [224, 224])).detach().cpu()
                 grid = torch.cat(list(grids.values()), 1)
                 grid_image = (grid.numpy().transpose(1, 2, 0).copy() * 255)[:, :, [2, 1, 0]]
                 grid_image = np.minimum(np.maximum(grid_image, 0), 255).astype(np.uint8)
+
                 cv2.imwrite('{}/{}.jpg'.format(savefolder, k), grid_image)
 
         single_params = {
@@ -166,6 +168,7 @@ def optimize(self, images, landmarks, image_masks, savefolder=None):
             'pose': pose.detach().cpu().numpy(),
             'cam': cam.detach().cpu().numpy(),
             'verts': trans_vertices.detach().cpu().numpy(),
+            'albedos':albedos.detach().cpu().numpy(),
             'tex': tex.detach().cpu().numpy(),
             'lit': lights.detach().cpu().numpy()
         }
@@ -213,6 +216,10 @@ def run(self, imagepath, landmarkpath):
         util.check_mkdir(savefolder)
         # optimize
         single_params = self.optimize(images, landmarks, image_masks, savefolder)
+        self.render.save_obj(filename=savefile[:-4]+'.obj',
+                             vertices=torch.from_numpy(single_params['verts'][0]).to(self.device),
+                             textures=torch.from_numpy(single_params['albedos'][0]).to(self.device)
+                             )
         np.save(savefile, single_params)
 
 

util.py

@@ -86,7 +86,6 @@ def batch_orth_proj(X, camera):
 
 def batch_persp_proj(vertices, cam, f, t, orig_size=256, eps=1e-9):
     '''
-    full pytorch version
     Calculate projective transformation of vertices given a projection matrix
     Input parameters:
     f: torch tensor of focal length
@@ -254,4 +253,61 @@ def plot_kpts(image, kpts, color = 'r'):
         ed = kpts[i + 1, :2]
         image = cv2.line(image, (st[0], st[1]), (ed[0], ed[1]), (255, 255, 255), 1)
 
-    return image
+    return image
+
+
+def save_obj(filename, vertices, faces, textures=None, uvcoords=None, uvfaces=None, texture_type='surface'):
+    assert vertices.ndimension() == 2
+    assert faces.ndimension() == 2
+    assert texture_type in ['surface', 'vertex']
+    # assert texture_res >= 2
+
+    if textures is not None and texture_type == 'surface':
+        textures =textures.detach().cpu().numpy().transpose(1,2,0)
+        filename_mtl = filename[:-4] + '.mtl'
+        filename_texture = filename[:-4] + '.png'
+        material_name = 'material_1'
+        # texture_image, vertices_textures = create_texture_image(textures, texture_res)
+        texture_image = textures
+        texture_image = texture_image.clip(0, 1)
+        texture_image = (texture_image * 255).astype('uint8')
+        imsave(filename_texture, texture_image)
+
+    faces = faces.detach().cpu().numpy()
+
+    with open(filename, 'w') as f:
+        f.write('# %s\n' % os.path.basename(filename))
+        f.write('#\n')
+        f.write('\n')
+
+        if textures is not None:
+            f.write('mtllib %s\n\n' % os.path.basename(filename_mtl))
+
+        if textures is not None and texture_type == 'vertex':
+            for vertex, color in zip(vertices, textures):
+                f.write('v %.8f %.8f %.8f %.8f %.8f %.8f\n' % (vertex[0], vertex[1], vertex[2],
+                                                               color[0], color[1], color[2]))
+            f.write('\n')
+        else:
+            for vertex in vertices:
+                f.write('v %.8f %.8f %.8f\n' % (vertex[0], vertex[1], vertex[2]))
+            f.write('\n')
+
+        if textures is not None and texture_type == 'surface':
+            for vertex in uvcoords.reshape((-1, 2)):
+                f.write('vt %.8f %.8f\n' % (vertex[0], vertex[1]))
+            f.write('\n')
+
+            f.write('usemtl %s\n' % material_name)
+            for i, face in enumerate(faces):
+                f.write('f %d/%d %d/%d %d/%d\n' % (
+                    face[0] + 1, uvfaces[i,0]+1, face[1] + 1, uvfaces[i,1]+1, face[2] + 1, uvfaces[i,2]+1))
+            f.write('\n')
+        else:
+            for face in faces:
+                f.write('f %d %d %d\n' % (face[0] + 1, face[1] + 1, face[2] + 1))
+
+    if textures is not None and texture_type == 'surface':
+        with open(filename_mtl, 'w') as f:
+            f.write('newmtl %s\n' % material_name)
+            f.write('map_Kd %s\n' % os.path.basename(filename_texture))