ac5113 commited on
Commit
63e7312
·
1 Parent(s): cb464db
Files changed (2) hide show
  1. app.py +104 -103
  2. packages.txt +1 -2
app.py CHANGED
@@ -18,7 +18,7 @@ import pyrender
18
  from models.deco import DECO
19
  from common import constants
20
 
21
- os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
22
 
23
  if torch.cuda.is_available():
24
  device = torch.device('cuda')
@@ -75,105 +75,105 @@ def initiate_model(model_path):
75
 
76
  return deco_model
77
 
78
- def render_image(scene, img_res, img=None, viewer=False):
79
- '''
80
- Render the given pyrender scene and return the image. Can also overlay the mesh on an image.
81
- '''
82
- if viewer:
83
- pyrender.Viewer(scene, use_raymond_lighting=True)
84
- return 0
85
- else:
86
- r = pyrender.OffscreenRenderer(viewport_width=img_res,
87
- viewport_height=img_res,
88
- point_size=1.0)
89
- color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
90
- color = color.astype(np.float32) / 255.0
91
-
92
- if img is not None:
93
- valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
94
- input_img = img.detach().cpu().numpy()
95
- output_img = (color[:, :, :-1] * valid_mask +
96
- (1 - valid_mask) * input_img)
97
- else:
98
- output_img = color
99
- return output_img
100
-
101
- def create_scene(mesh, img, focal_length=500, camera_center=250, img_res=500):
102
- # Setup the scene
103
- scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0],
104
- ambient_light=(0.3, 0.3, 0.3))
105
- # add mesh for camera
106
- camera_pose = np.eye(4)
107
- camera_rotation = np.eye(3, 3)
108
- camera_translation = np.array([0., 0, 2.5])
109
- camera_pose[:3, :3] = camera_rotation
110
- camera_pose[:3, 3] = camera_rotation @ camera_translation
111
- pyrencamera = pyrender.camera.IntrinsicsCamera(
112
- fx=focal_length, fy=focal_length,
113
- cx=camera_center, cy=camera_center)
114
- scene.add(pyrencamera, pose=camera_pose)
115
- # create and add light
116
- light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
117
- light_pose = np.eye(4)
118
- for lp in [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [-1, -1, 1]]:
119
- light_pose[:3, 3] = mesh.vertices.mean(0) + np.array(lp)
120
- # out_mesh.vertices.mean(0) + np.array(lp)
121
- scene.add(light, pose=light_pose)
122
- # add body mesh
123
- material = pyrender.MetallicRoughnessMaterial(
124
- metallicFactor=0.0,
125
- alphaMode='OPAQUE',
126
- baseColorFactor=(1.0, 1.0, 0.9, 1.0))
127
- mesh_images = []
128
-
129
- # resize input image to fit the mesh image height
130
- img_height = img_res
131
- img_width = int(img_height * img.shape[1] / img.shape[0])
132
- img = cv2.resize(img, (img_width, img_height))
133
- mesh_images.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
134
-
135
- for sideview_angle in [0, 90, 180, 270]:
136
- out_mesh = mesh.copy()
137
- rot = trimesh.transformations.rotation_matrix(
138
- np.radians(sideview_angle), [0, 1, 0])
139
- out_mesh.apply_transform(rot)
140
- out_mesh = pyrender.Mesh.from_trimesh(
141
- out_mesh,
142
- material=material)
143
- mesh_pose = np.eye(4)
144
- scene.add(out_mesh, pose=mesh_pose, name='mesh')
145
- output_img = render_image(scene, img_res)
146
- output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
147
- output_img = np.asarray(output_img)[:, :, :3]
148
- mesh_images.append(output_img)
149
- # delete the previous mesh
150
- prev_mesh = scene.get_nodes(name='mesh').pop()
151
- scene.remove_node(prev_mesh)
152
-
153
- # show upside down view
154
- for topview_angle in [90, 270]:
155
- out_mesh = mesh.copy()
156
- rot = trimesh.transformations.rotation_matrix(
157
- np.radians(topview_angle), [1, 0, 0])
158
- out_mesh.apply_transform(rot)
159
- out_mesh = pyrender.Mesh.from_trimesh(
160
- out_mesh,
161
- material=material)
162
- mesh_pose = np.eye(4)
163
- scene.add(out_mesh, pose=mesh_pose, name='mesh')
164
- output_img = render_image(scene, img_res)
165
- output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
166
- output_img = np.asarray(output_img)[:, :, :3]
167
- mesh_images.append(output_img)
168
- # delete the previous mesh
169
- prev_mesh = scene.get_nodes(name='mesh').pop()
170
- scene.remove_node(prev_mesh)
171
-
172
- # stack images
173
- IMG = np.hstack(mesh_images)
174
- IMG = pil_img.fromarray(IMG)
175
- IMG.thumbnail((3000, 3000))
176
- return IMG
177
 
178
  def main(pil_img, out_dir='demo_out', model_path='checkpoint/deco_best.pth', mesh_colour=[130, 130, 130, 255], annot_colour=[0, 255, 0, 255]):
179
  deco_model = initiate_model(model_path)
@@ -208,9 +208,10 @@ def main(pil_img, out_dir='demo_out', model_path='checkpoint/deco_best.pth', mes
208
  body_model_smpl.visual.vertex_colors[vert] = mesh_colour
209
  body_model_smpl.visual.vertex_colors[cont_smpl] = annot_colour
210
 
211
- rend = create_scene(body_model_smpl, img)
212
- os.makedirs(os.path.join(out_dir, 'Renders'), exist_ok=True)
213
- rend.save(os.path.join(out_dir, 'Renders', 'pred.png'))
 
214
 
215
  mesh_out_dir = os.path.join(out_dir, 'Preds')
216
  os.makedirs(mesh_out_dir, exist_ok=True)
 
18
  from models.deco import DECO
19
  from common import constants
20
 
21
+ # os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
22
 
23
  if torch.cuda.is_available():
24
  device = torch.device('cuda')
 
75
 
76
  return deco_model
77
 
78
+ # def render_image(scene, img_res, img=None, viewer=False):
79
+ # '''
80
+ # Render the given pyrender scene and return the image. Can also overlay the mesh on an image.
81
+ # '''
82
+ # if viewer:
83
+ # pyrender.Viewer(scene, use_raymond_lighting=True)
84
+ # return 0
85
+ # else:
86
+ # r = pyrender.OffscreenRenderer(viewport_width=img_res,
87
+ # viewport_height=img_res,
88
+ # point_size=1.0)
89
+ # color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
90
+ # color = color.astype(np.float32) / 255.0
91
+
92
+ # if img is not None:
93
+ # valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
94
+ # input_img = img.detach().cpu().numpy()
95
+ # output_img = (color[:, :, :-1] * valid_mask +
96
+ # (1 - valid_mask) * input_img)
97
+ # else:
98
+ # output_img = color
99
+ # return output_img
100
+
101
+ # def create_scene(mesh, img, focal_length=500, camera_center=250, img_res=500):
102
+ # # Setup the scene
103
+ # scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0],
104
+ # ambient_light=(0.3, 0.3, 0.3))
105
+ # # add mesh for camera
106
+ # camera_pose = np.eye(4)
107
+ # camera_rotation = np.eye(3, 3)
108
+ # camera_translation = np.array([0., 0, 2.5])
109
+ # camera_pose[:3, :3] = camera_rotation
110
+ # camera_pose[:3, 3] = camera_rotation @ camera_translation
111
+ # pyrencamera = pyrender.camera.IntrinsicsCamera(
112
+ # fx=focal_length, fy=focal_length,
113
+ # cx=camera_center, cy=camera_center)
114
+ # scene.add(pyrencamera, pose=camera_pose)
115
+ # # create and add light
116
+ # light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
117
+ # light_pose = np.eye(4)
118
+ # for lp in [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [-1, -1, 1]]:
119
+ # light_pose[:3, 3] = mesh.vertices.mean(0) + np.array(lp)
120
+ # # out_mesh.vertices.mean(0) + np.array(lp)
121
+ # scene.add(light, pose=light_pose)
122
+ # # add body mesh
123
+ # material = pyrender.MetallicRoughnessMaterial(
124
+ # metallicFactor=0.0,
125
+ # alphaMode='OPAQUE',
126
+ # baseColorFactor=(1.0, 1.0, 0.9, 1.0))
127
+ # mesh_images = []
128
+
129
+ # # resize input image to fit the mesh image height
130
+ # img_height = img_res
131
+ # img_width = int(img_height * img.shape[1] / img.shape[0])
132
+ # img = cv2.resize(img, (img_width, img_height))
133
+ # mesh_images.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
134
+
135
+ # for sideview_angle in [0, 90, 180, 270]:
136
+ # out_mesh = mesh.copy()
137
+ # rot = trimesh.transformations.rotation_matrix(
138
+ # np.radians(sideview_angle), [0, 1, 0])
139
+ # out_mesh.apply_transform(rot)
140
+ # out_mesh = pyrender.Mesh.from_trimesh(
141
+ # out_mesh,
142
+ # material=material)
143
+ # mesh_pose = np.eye(4)
144
+ # scene.add(out_mesh, pose=mesh_pose, name='mesh')
145
+ # output_img = render_image(scene, img_res)
146
+ # output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
147
+ # output_img = np.asarray(output_img)[:, :, :3]
148
+ # mesh_images.append(output_img)
149
+ # # delete the previous mesh
150
+ # prev_mesh = scene.get_nodes(name='mesh').pop()
151
+ # scene.remove_node(prev_mesh)
152
+
153
+ # # show upside down view
154
+ # for topview_angle in [90, 270]:
155
+ # out_mesh = mesh.copy()
156
+ # rot = trimesh.transformations.rotation_matrix(
157
+ # np.radians(topview_angle), [1, 0, 0])
158
+ # out_mesh.apply_transform(rot)
159
+ # out_mesh = pyrender.Mesh.from_trimesh(
160
+ # out_mesh,
161
+ # material=material)
162
+ # mesh_pose = np.eye(4)
163
+ # scene.add(out_mesh, pose=mesh_pose, name='mesh')
164
+ # output_img = render_image(scene, img_res)
165
+ # output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
166
+ # output_img = np.asarray(output_img)[:, :, :3]
167
+ # mesh_images.append(output_img)
168
+ # # delete the previous mesh
169
+ # prev_mesh = scene.get_nodes(name='mesh').pop()
170
+ # scene.remove_node(prev_mesh)
171
+
172
+ # # stack images
173
+ # IMG = np.hstack(mesh_images)
174
+ # IMG = pil_img.fromarray(IMG)
175
+ # IMG.thumbnail((3000, 3000))
176
+ # return IMG
177
 
178
  def main(pil_img, out_dir='demo_out', model_path='checkpoint/deco_best.pth', mesh_colour=[130, 130, 130, 255], annot_colour=[0, 255, 0, 255]):
179
  deco_model = initiate_model(model_path)
 
208
  body_model_smpl.visual.vertex_colors[vert] = mesh_colour
209
  body_model_smpl.visual.vertex_colors[cont_smpl] = annot_colour
210
 
211
+ # rend = create_scene(body_model_smpl, img)
212
+ # os.makedirs(os.path.join(out_dir, 'Renders'), exist_ok=True)
213
+ # rend.save(os.path.join(out_dir, 'Renders', 'pred.png'))
214
+ rend = img
215
 
216
  mesh_out_dir = os.path.join(out_dir, 'Preds')
217
  os.makedirs(mesh_out_dir, exist_ok=True)
packages.txt CHANGED
@@ -1,2 +1 @@
1
- freeglut3-dev
2
- libegl1
 
1
+ freeglut3-dev