├── .gitignore ├── README.md ├── __pycache__ └── depth_renderer.cpython-36.pyc ├── depth_renderer.py ├── example_images ├── fuse0.png ├── fuse1.png ├── rnd0.png └── rnd1.png ├── fuse.py ├── gen_render.sh ├── rastertriangle.h ├── rastertriangle_so.cpp ├── rastertriangle_so.sh ├── rastertriangle_so.so ├── rgbd_renderer.py └── sampled_poses └── .gitkeep /.gitignore: -------------------------------------------------------------------------------- 1 | Linemod_preprocessed 2 | SUN2012pascalformat 3 | sampled_poses/*.pkl 4 | *.so 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Raster Triangle 2 | This is a simple renderer with z-buffer for synthesis data generating. With an input mesh and its pose, the rendered RGB and depth map are generated. It has been applied to [PVN3D](https://github.com/ethnhe/PVN3D) and [FFB6D](https://github.com/ethnhe/FFB6D). 3 | 4 | ## Installation 5 | Compile the source code by 6 | ```shell 7 | chmod +x ./rastertriangle_so.sh 8 | ./rastertriangle_so.sh 9 | ``` 10 | 11 | ## Datasets: 12 | - Download the LineMOD dataset from [here](https://drive.google.com/drive/folders/19ivHpaKm9dOrr12fzC8IDFczWRPFxho7). Unzip it and link the unzipped folder to this folder: 13 | ```shell 14 | ln -s path_to_Linemod_preprocessed ./Linemod_preprocessed 15 | ``` 16 | - Download SUN2012pascalformat dataset from [here](http://groups.csail.mit.edu/vision/SUN/releases/SUN2012pascalformat.tar.gz)(works with FireFox or right click and open link in new tab in Chrome) or [here](https://github.com/ShapeNet/RenderForCNN/blob/master/datasets/get_sun2012pascalformat.sh). Unzip it and link the unzipped folder to this folder: 17 | ```shell 18 | ln -s path_to_SUN2012pascalformat ./ 19 | ``` 20 | - Download the sample poses from [here](https://hkustconnect-my.sharepoint.com/:f:/g/personal/yhebk_connect_ust_hk/End-Ha7PuQFNktD_ZqBIuQgBwR0wNVDPi-Bneulo7Dy-JA?e=WQba57) and move the pickle files to folder ```sampled_poses/```. 21 | 22 | ## Generate synthesis data 23 | - Generate rendered data. Sampled poses for each object are provided in ``sampled_poses/``, which are generated by scripts [here](https://github.com/zju3dv/pvnet-rendering). Run the following command to generate rendered data: 24 | ```shell 25 | python3 rgbd_renderer.py --help 26 | python3 rgbd_renderer.py --cls ape --render_num 70000 27 | ``` 28 | The generated depth map is in unit meter(m). 29 | Example rendered images are as follows: 30 | ![render0](./example_images/rnd0.png) 31 | ![render1](./example_images/rnd1.png) 32 | 33 | 34 | - Run the following command to generate fusing data: 35 | ```shell 36 | python3 fuse.py --help 37 | python3 fuse.py --cls ape --fuse_num 10000 38 | ``` 39 | The generated depth map is in unit meter(m). 40 | Example fused images are as follows: 41 | ![fuse0](./example_images/fuse0.png) 42 | ![fuse1](./example_images/fuse1.png) 43 | -------------------------------------------------------------------------------- /__pycache__/depth_renderer.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethnhe/raster_triangle/b31501a0f90b1e7f1270862a8d7decb435ea6f6a/__pycache__/depth_renderer.cpython-36.pyc -------------------------------------------------------------------------------- /depth_renderer.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import os 3 | import cv2 4 | import trimesh 5 | import numpy as np 6 | import ctypes as ct 7 | from glob import glob 8 | import matplotlib.pyplot as plt 9 | 10 | 11 | mdl_ptn = '/data/6D_pose_data/BOP/lm/models/*.ply' 12 | mdl_pth_lst = glob(mdl_ptn) 13 | cur_dir = os.path.dirname(os.path.realpath(__file__)) 14 | so_p = os.path.join(cur_dir, 'rastertriangle_so.so') 15 | print("raster triangle so path:", so_p) 16 | dll = np.ctypeslib.load_library(so_p, '.') 17 | 18 | 19 | def load_trimesh(pth, scale2m=1.): 20 | mesh = trimesh.load(pth, force='mesh') 21 | mesh.vertices = mesh.vertices / scale2m 22 | return mesh 23 | 24 | 25 | def depth_renderer(mesh, K, RT=np.eye(4), h=480, w=640, scale2m=1.): 26 | vtxs = np.array(mesh.vertices) / scale2m 27 | faces = np.array(mesh.faces) 28 | n_face = faces.shape[0] 29 | face = faces.flatten().copy() 30 | 31 | R, T = RT[:3, :3], RT[:3, 3] 32 | 33 | new_xyz = vtxs.copy() 34 | new_xyz = np.dot(new_xyz, R.T) + T[None, :] 35 | p2ds = np.dot(new_xyz.copy(), K.T) 36 | p2ds = p2ds[:, :2] / p2ds[:, 2:] 37 | p2ds = np.require(p2ds.flatten(), 'float32', 'C') 38 | 39 | face = np.require(face, 'int32', 'C') 40 | new_xyz = np.require(new_xyz, 'float32', 'C') 41 | zs = np.require(new_xyz[:, 2].copy(), 'float32', 'C') 42 | zbuf = np.require(np.zeros(h*w), 'float32', 'C') 43 | 44 | dll.zbuffer( 45 | ct.c_int(h), 46 | ct.c_int(w), 47 | p2ds.ctypes.data_as(ct.c_void_p), 48 | # new_xyz.ctypes.data_as(ct.c_void_p), 49 | zs.ctypes.data_as(ct.c_void_p), 50 | ct.c_int(n_face), 51 | face.ctypes.data_as(ct.c_void_p), 52 | zbuf.ctypes.data_as(ct.c_void_p), 53 | ) 54 | 55 | zbuf.resize((h, w)) 56 | msk = (zbuf > 1e-8).astype('uint8') 57 | zbuf *= msk.astype(zbuf.dtype) # * 1000.0 58 | return zbuf, msk 59 | 60 | 61 | def dpt2cld(dpt, cam_scale, K): 62 | """ 63 | dpt: h x w depth image 64 | cam_scale: scale depth to (m). e.g: dpt in mm, cam_scale should be 1000 65 | K: camera intrinsic 66 | """ 67 | h, w = dpt.shape[0], dpt.shape[1] 68 | xmap = np.array([[j for i in range(w)] for j in range(h)]) 69 | ymap = np.array([[i for i in range(w)] for j in range(h)]) 70 | 71 | if len(dpt.shape) > 2: 72 | dpt = dpt[:, :, 0] 73 | msk_dp = dpt > 1e-6 74 | choose = msk_dp.flatten() 75 | if choose.sum() < 1: 76 | return None, None 77 | 78 | dpt_mskd = dpt.flatten()[choose][:, np.newaxis].astype(np.float32) 79 | xmap_mskd = xmap.flatten()[choose][:, np.newaxis].astype(np.float32) 80 | ymap_mskd = ymap.flatten()[choose][:, np.newaxis].astype(np.float32) 81 | 82 | pt2 = dpt_mskd / cam_scale 83 | cam_cx, cam_cy = K[0][2], K[1][2] 84 | cam_fx, cam_fy = K[0][0], K[1][1] 85 | pt0 = (ymap_mskd - cam_cx) * pt2 / cam_fx 86 | pt1 = (xmap_mskd - cam_cy) * pt2 / cam_fy 87 | cld = np.concatenate((pt0, pt1, pt2), axis=1) 88 | return cld 89 | 90 | 91 | def dpt2heat(dpt): 92 | vd = dpt[dpt > 0] 93 | dpt[dpt > 0] = (dpt[dpt > 0] - vd.min() + 0.4) / (vd.max() - vd.min()) 94 | dpt = (dpt * 255).astype(np.uint8) 95 | colormap = plt.get_cmap('inferno') 96 | heatmap = (colormap(dpt.copy()) * 2**16).astype(np.uint16)[:, :, :3] 97 | heatmap = cv2.cvtColor(heatmap, cv2.COLOR_RGB2BGR) 98 | heatmap[dpt == 0, :] = np.array([255, 255, 255]) 99 | return heatmap 100 | 101 | 102 | def main(): 103 | for mdl_pth in mdl_pth_lst: 104 | mesh = load_trimesh(mdl_pth, scale2m=1000.) 105 | # mesh = load_trimesh(mdl_pth, scale2m=1.) 106 | K = np.array([[700., 0., 320.], [0., 700., 240.], [0., 0., 1.]]) 107 | RT = np.eye(4) 108 | RT[:3, 3] = np.array([0., 0., 1.]) 109 | depth, msk = depth_renderer(mesh, K, RT=RT) 110 | 111 | try: 112 | from neupeak.utils.webcv2 import imshow, waitKey 113 | except ImportError: 114 | from cv2 import imshow, waitKey 115 | show_dpt = dpt2heat(depth.copy()) 116 | imshow("render_depth:", show_dpt) 117 | waitKey(0) 118 | 119 | 120 | if __name__ == "__main__": 121 | main() 122 | # vim: ts=4 sw=4 sts=4 expandtab 123 | -------------------------------------------------------------------------------- /example_images/fuse0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethnhe/raster_triangle/b31501a0f90b1e7f1270862a8d7decb435ea6f6a/example_images/fuse0.png -------------------------------------------------------------------------------- /example_images/fuse1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethnhe/raster_triangle/b31501a0f90b1e7f1270862a8d7decb435ea6f6a/example_images/fuse1.png -------------------------------------------------------------------------------- /example_images/rnd0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethnhe/raster_triangle/b31501a0f90b1e7f1270862a8d7decb435ea6f6a/example_images/rnd0.png -------------------------------------------------------------------------------- /example_images/rnd1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethnhe/raster_triangle/b31501a0f90b1e7f1270862a8d7decb435ea6f6a/example_images/rnd1.png -------------------------------------------------------------------------------- /fuse.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import cv2 4 | import pickle 5 | import yaml 6 | import numpy as np 7 | from glob import glob 8 | from tqdm import tqdm 9 | from PIL import ImageFile, Image 10 | from plyfile import PlyData 11 | from concurrent.futures import ProcessPoolExecutor 12 | from argparse import ArgumentParser 13 | try: 14 | from neupeak.utils.webcv2 import imshow, waitKey 15 | except: 16 | from cv2 import imshow, waitKey 17 | 18 | 19 | parser = ArgumentParser() 20 | parser.add_argument( 21 | "--cls", type=str, default="ape", 22 | help="Target object from {ape, benchvise, cam, can, cat, driller, duck, \ 23 | eggbox, glue, holepuncher, iron, lamp, phone} (default ape)" 24 | ) 25 | parser.add_argument( 26 | '--fuse_num', type=int, default=10000, 27 | help="Number of images you want to generate." 28 | ) 29 | parser.add_argument( 30 | '--DEBUG', action="store_true", 31 | help="To show the generated images or not." 32 | ) 33 | args = parser.parse_args() 34 | DEBUG = args.DEBUG 35 | 36 | Intrinsic_matrix = { 37 | 'linemod': np.array([[572.4114, 0., 325.2611], 38 | [0., 573.57043, 242.04899], 39 | [0., 0., 1.]]), 40 | 'blender': np.array([[700., 0., 320.], 41 | [0., 700., 240.], 42 | [0., 0., 1.]]) 43 | } 44 | lm_obj_dict={ 45 | 'ape':1, 46 | 'benchvise':2, 47 | 'cam':4, 48 | 'can':5, 49 | 'cat':6, 50 | 'driller':8, 51 | 'duck':9, 52 | 'eggbox':10, 53 | 'glue':11, 54 | 'holepuncher':12, 55 | 'iron':13, 56 | 'lamp':14, 57 | 'phone':15, 58 | } 59 | root = './Linemod_preprocessed' 60 | cls_root_ptn = os.path.join(root, "data/%02d/") 61 | 62 | 63 | def ensure_dir(pth): 64 | if not os.path.exists(pth): 65 | os.system("mkdir -p {}".format(pth)) 66 | 67 | 68 | def read_lines(pth): 69 | with open(pth, 'r') as f: 70 | return [ 71 | line.strip() for line in f.readlines() 72 | ] 73 | 74 | 75 | def read_pickle(pkl_path): 76 | with open(pkl_path, 'rb') as f: 77 | return pickle.load(f) 78 | 79 | 80 | def save_pickle(data, pkl_path): 81 | with open(pkl_path, 'wb') as f: 82 | pickle.dump(data, f) 83 | 84 | 85 | def collect_train_info(cls_name): 86 | cls_id = lm_obj_dict[cls_name] 87 | cls_root = cls_root_ptn % cls_id 88 | tr_pth = os.path.join( 89 | cls_root, "train.txt" 90 | ) 91 | train_fns = read_lines(tr_pth) 92 | 93 | return train_fns 94 | 95 | 96 | def collect_linemod_set_info( 97 | linemod_dir, linemod_cls_name, cache_dir='./data/LINEMOD/cache' 98 | ): 99 | database = [] 100 | if not os.path.exists(cache_dir): 101 | os.mkdir(cache_dir) 102 | if os.path.exists( 103 | os.path.join(cache_dir,'{}_info.pkl').format(linemod_cls_name) 104 | ): 105 | return read_pickle( 106 | os.path.join(cache_dir,'{}_info.pkl').format(linemod_cls_name) 107 | ) 108 | 109 | train_fns = collect_train_info(linemod_cls_name) 110 | cls_id = lm_obj_dict[linemod_cls_name] 111 | cls_root = cls_root_ptn % cls_id 112 | meta_file = open(os.path.join(cls_root, 'gt.yml'), "r") 113 | meta_lst = yaml.load(meta_file) 114 | 115 | print('begin generate database {}'.format(linemod_cls_name)) 116 | rgb_ptn = os.path.join(cls_root, "rgb/{}.png") 117 | msk_ptn = os.path.join(cls_root, "mask/{}.png") 118 | dpt_ptn = os.path.join(cls_root, "depth/{}.png") 119 | for item in train_fns: 120 | data={} 121 | data['rgb_pth'] = rgb_ptn.format(item) 122 | data['dpt_pth'] = dpt_ptn.format(item) 123 | data['msk_pth'] = msk_ptn.format(item) 124 | 125 | meta = meta_lst[int(item)] 126 | if cls_id == 2: 127 | for i in range(0, len(meta)): 128 | if meta[i]['obj_id'] == 2: 129 | meta = meta[i] 130 | break 131 | else: 132 | meta = meta[0] 133 | R = np.resize(np.array(meta['cam_R_m2c']), (3, 3)) 134 | T = np.array(meta['cam_t_m2c']) / 1000.0 135 | RT = np.concatenate((R, T[:, None]), axis=1) 136 | data['RT'] = RT 137 | database.append(data) 138 | 139 | print( 140 | 'successfully generate database {} len {}'.format( 141 | linemod_cls_name, len(database) 142 | ) 143 | ) 144 | save_pickle( 145 | database, os.path.join(cache_dir,'{}_info.pkl').format(linemod_cls_name) 146 | ) 147 | return database 148 | 149 | 150 | def randomly_read_background(background_dir,cache_dir): 151 | if os.path.exists(os.path.join(cache_dir,'background_info.pkl')): 152 | fns = read_pickle(os.path.join(cache_dir,'background_info.pkl')) 153 | else: 154 | fns = glob(os.path.join(background_dir,'*.jpg')) + \ 155 | glob(os.path.join(background_dir,'*.png')) 156 | save_pickle(fns, os.path.join(cache_dir,'background_info.pkl')) 157 | 158 | return cv2.imread(fns[np.random.randint(0,len(fns))])[:, :, ::-1] 159 | 160 | 161 | def fuse_regions(rgbs, masks, depths, begins, cls_ids, background, th, tw, cls): 162 | fuse_order = np.arange(len(rgbs)) 163 | np.random.shuffle(fuse_order) 164 | fuse_img = background 165 | fuse_img = cv2.resize(fuse_img,(tw,th),interpolation=cv2.INTER_LINEAR) 166 | fuse_mask = np.zeros([fuse_img.shape[0],fuse_img.shape[1]],np.int32) 167 | INF = pow(2,15) 168 | fuse_depth = np.ones([fuse_img.shape[0], fuse_img.shape[1]], np.uint16) * INF 169 | t_cls_id = lm_obj_dict[cls] 170 | if len(background.shape) < 3: 171 | return None, None, None, None 172 | for idx in fuse_order: 173 | if len(rgbs[idx].shape) < 3: 174 | continue 175 | cls_id = cls_ids[idx] 176 | rh,rw = masks[idx].shape 177 | if cls_id == t_cls_id: 178 | bh, bw = begins[idx][0], begins[idx][1] 179 | else: 180 | bh = np.random.randint(0,fuse_img.shape[0]-rh) 181 | bw = np.random.randint(0,fuse_img.shape[1]-rw) 182 | 183 | silhouette = masks[idx]>0 184 | out_silhouette = np.logical_not(silhouette) 185 | fuse_depth_patch = fuse_depth[bh:bh+rh, bw:bw+rw].copy() 186 | cover = (depths[idx] < fuse_depth_patch) * silhouette 187 | not_cover = np.logical_not(cover) 188 | 189 | fuse_mask[bh:bh+rh,bw:bw+rw] *= not_cover.astype(fuse_mask.dtype) 190 | cover_msk = masks[idx] * cover.astype(masks[idx].dtype) 191 | fuse_mask[bh:bh+rh,bw:bw+rw] += cover_msk 192 | 193 | fuse_img[bh:bh+rh,bw:bw+rw] *= not_cover.astype(fuse_img.dtype)[:,:,None] 194 | cover_rgb = rgbs[idx] * cover.astype(rgbs[idx].dtype)[:,:,None] 195 | fuse_img[bh:bh+rh,bw:bw+rw] += cover_rgb 196 | 197 | fuse_depth[bh:bh+rh, bw:bw+rw] *= not_cover.astype(fuse_depth.dtype) 198 | cover_dpt = depths[idx] * cover.astype(depths[idx].dtype) 199 | fuse_depth[bh:bh+rh, bw:bw+rw] += cover_dpt.astype(fuse_depth.dtype) 200 | 201 | begins[idx][0] = -begins[idx][0]+bh 202 | begins[idx][1] = -begins[idx][1]+bw 203 | 204 | dp_bg = (fuse_depth == INF) 205 | dp_bg_filter = np.logical_not(dp_bg) 206 | fuse_depth *= dp_bg_filter.astype(fuse_depth.dtype) 207 | 208 | return fuse_img, fuse_mask, fuse_depth, begins 209 | 210 | 211 | def randomly_sample_foreground(image_db, linemod_dir): 212 | idx = np.random.randint(0,len(image_db)) 213 | rgb_pth = image_db[idx]['rgb_pth'] 214 | dpt_pth = image_db[idx]['dpt_pth'] 215 | msk_pth = image_db[idx]['msk_pth'] 216 | with Image.open(dpt_pth) as di: 217 | depth = np.array(di).astype(np.int16) 218 | with Image.open(msk_pth) as li: 219 | mask = np.array(li).astype(np.int16) 220 | with Image.open(rgb_pth) as ri: 221 | rgb = np.array(ri)[:, :, :3].astype(np.uint8) 222 | 223 | mask = np.sum(mask,2)>0 224 | mask = np.asarray(mask,np.int32) 225 | 226 | hs, ws = np.nonzero(mask) 227 | hmin, hmax = np.min(hs),np.max(hs) 228 | wmin, wmax = np.min(ws),np.max(ws) 229 | 230 | mask = mask[hmin:hmax,wmin:wmax] 231 | rgb = rgb[hmin:hmax,wmin:wmax] 232 | depth = depth[hmin:hmax, wmin:wmax] 233 | 234 | rgb *= mask.astype(np.uint8)[:,:,None] 235 | depth *= mask.astype(np.uint16)[:,:] 236 | begin = [hmin,wmin] 237 | pose = image_db[idx]['RT'] 238 | 239 | return rgb, mask, depth, begin, pose 240 | 241 | 242 | def save_fuse_data( 243 | output_dir, idx, fuse_img, fuse_mask, fuse_depth, fuse_begins, t_pose, cls 244 | ): 245 | cls_id = lm_obj_dict[cls] 246 | if (fuse_mask == cls_id).sum() < 20: 247 | return None 248 | os.makedirs(output_dir, exist_ok=True) 249 | fuse_mask = fuse_mask.astype(np.uint8) 250 | data = {} 251 | data['rgb'] = fuse_img 252 | data['mask'] = fuse_mask 253 | data['depth'] = fuse_depth.astype(np.float32) / 1000.0 254 | data['K'] = Intrinsic_matrix['linemod'] 255 | data['RT'] = t_pose 256 | data['cls_typ'] = cls 257 | data['rnd_typ'] = 'fuse' 258 | data['begins'] = fuse_begins 259 | if DEBUG: 260 | imshow("rgb", fuse_img[:, :, ::-1]) 261 | imshow("depth", (fuse_depth / fuse_depth.max() * 255).astype('uint8')) 262 | imshow("label", (fuse_mask / fuse_mask.max() * 255).astype("uint8")) 263 | waitKey(0) 264 | sv_pth = os.path.join(output_dir, "{}.pkl".format(idx)) 265 | pickle.dump(data, open(sv_pth, 'wb')) 266 | sv_pth = os.path.abspath(sv_pth) 267 | return sv_pth 268 | 269 | 270 | def prepare_dataset_single( 271 | output_dir, idx, linemod_dir, background_dir, cache_dir, seed, cls 272 | ): 273 | time_begin = time.time() 274 | np.random.seed(seed) 275 | rgbs, masks, depths, begins, poses, cls_ids = [], [], [], [], [], [] 276 | image_dbs={} 277 | for cls_name in lm_obj_dict.keys(): 278 | cls_id = lm_obj_dict[cls_name] 279 | image_dbs[cls_id] = collect_linemod_set_info( 280 | linemod_dir, cls_name, cache_dir 281 | ) 282 | 283 | for cls_name in lm_obj_dict.keys(): 284 | cls_id = lm_obj_dict[cls_name] 285 | rgb, mask, depth, begin, pose = randomly_sample_foreground( 286 | image_dbs[cls_id], linemod_dir 287 | ) 288 | if cls_name == cls: 289 | t_pose = pose 290 | mask *= cls_id 291 | rgbs.append(rgb) 292 | masks.append(mask) 293 | depths.append(depth) 294 | begins.append(begin) 295 | poses.append(pose) 296 | cls_ids.append(cls_id) 297 | 298 | background = randomly_read_background(background_dir, cache_dir) 299 | 300 | fuse_img, fuse_mask, fuse_depth, fuse_begins= fuse_regions( 301 | rgbs, masks, depths, begins, cls_ids, background, 480, 640, cls 302 | ) 303 | 304 | if fuse_img is not None: 305 | sv_pth = save_fuse_data( 306 | output_dir, idx, fuse_img, fuse_mask, fuse_depth, fuse_begins, 307 | t_pose, cls 308 | ) 309 | return sv_pth 310 | 311 | 312 | def prepare_dataset_parallel( 313 | output_dir, linemod_dir, fuse_num, background_dir, cache_dir, 314 | worker_num=8, cls="ape" 315 | ): 316 | exector = ProcessPoolExecutor(max_workers=worker_num) 317 | futures = [] 318 | 319 | for cls_name in lm_obj_dict.keys(): 320 | collect_linemod_set_info(linemod_dir, cls_name, cache_dir) 321 | randomly_read_background(background_dir, cache_dir) 322 | 323 | for idx in np.arange(fuse_num): 324 | seed = np.random.randint(500000) 325 | futures.append(exector.submit( 326 | prepare_dataset_single, output_dir, idx, linemod_dir, 327 | background_dir, cache_dir, seed, cls 328 | )) 329 | 330 | pth_lst = [] 331 | for f in tqdm(futures): 332 | res = f.result() 333 | if res is not None: 334 | pth_lst.append(res) 335 | f_lst_pth = os.path.join(output_dir, "file_list.txt") 336 | with open(f_lst_pth, "w") as f: 337 | for item in pth_lst: 338 | print(item, file=f) 339 | 340 | 341 | if __name__=="__main__": 342 | cls = args.cls 343 | linemod_dir = './Linemod_preprocessed' 344 | output_dir = os.path.join(linemod_dir, "fuse", cls) 345 | ensure_dir(output_dir) 346 | background_dir = './SUN2012pascalformat/JPEGImages' 347 | cache_dir = './' 348 | fuse_num = args.fuse_num 349 | worker_num = 20 350 | prepare_dataset_parallel( 351 | output_dir, linemod_dir, fuse_num, background_dir, cache_dir, 352 | worker_num, cls=cls 353 | ) 354 | -------------------------------------------------------------------------------- /gen_render.sh: -------------------------------------------------------------------------------- 1 | cls_lst=('ape' 'benchvise' 'cam' 'can' 'cat' 'driller' 'duck' 'eggbox' 'glue' 'holepuncher' 'iron' 'lamp' 'phone') 2 | cls=${cls_lst[12]} 3 | python3 rgbd_renderer.py --cls_type ${cls} 4 | -------------------------------------------------------------------------------- /rastertriangle.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #ifndef RASTERTRIANGLE_H 4 | #define RASTERTRIANGLE_H 5 | 6 | /* 7 | 8 | Enumerate pixels of a triangle within a rectangular region 9 | Pixels on the boundary are considered inclusively -> connecting triangles may share pixels 10 | 11 | Iterator usage: 12 | 13 | PixelEnumerator pixels( 14 | x0,y0, //point 0, barycentric coordinate (0,0) 15 | x1,y1, //point 1, barycentric coordinate (1,0) 16 | x2,y2, //point 2, barycentric coordinate (0,1) 17 | x_lower_bound, //for xx_upper_bound (inclusive bounds), the pixel is clipped 18 | x_upper_bound, // 19 | y_lower_bound, //bounds for y 20 | y_upper_bound // 21 | ); 22 | int x,y; 23 | float u,v; 24 | while (pixels.getNext(&x,&y,&u,&v)){ 25 | //(u,v) is the barycentric coordinate: 26 | //(x,y) = (x0,y0)+ (x1-x0,y1-y0)*u + (x2-x0,y2-y0)*v 27 | 28 | //blablabla 29 | } 30 | 31 | There is also a function allPixels that packs all generated pixels into a std::vector 32 | 33 | */ 34 | 35 | #include 36 | #include 37 | #include 38 | #include 39 | 40 | namespace RasterTriangle{ 41 | 42 | template 43 | struct PixelEnumerator{ 44 | int lx,hx,ly,hy; 45 | Float x0,y0,x1,y1,x2,y2,u0,v0,u1,v1,u2,v2; 46 | int x,y,yhb; 47 | Float ubeg,vbeg,udelta,vdelta; 48 | Float ybeg; 49 | void sortPoints(){ 50 | if (x0>x1){ 51 | std::swap(x0,x1); 52 | std::swap(y0,y1); 53 | std::swap(u0,u1); 54 | std::swap(v0,v1); 55 | } 56 | if (x0>x2){ 57 | std::swap(x0,x2); 58 | std::swap(y0,y2); 59 | std::swap(u0,u2); 60 | std::swap(v0,v2); 61 | } 62 | if (x1>x2){ 63 | std::swap(x1,x2); 64 | std::swap(y1,y2); 65 | std::swap(u1,u2); 66 | std::swap(v1,v2); 67 | } 68 | } 69 | PixelEnumerator(Float x0,Float y0,Float x1,Float y1,Float x2,Float y2,int lx,int hx,int ly,int hy){ 70 | this->x0=x0; 71 | this->y0=y0; 72 | this->u0=0; 73 | this->v0=0; 74 | this->x1=x1; 75 | this->y1=y1; 76 | this->u1=1; 77 | this->v1=0; 78 | this->x2=x2; 79 | this->y2=y2; 80 | this->u2=0; 81 | this->v2=1; 82 | this->lx=lx; 83 | this->hx=hx; 84 | this->ly=ly; 85 | this->hy=hy; 86 | sortPoints(); 87 | x=int(std::ceil(std::max(Float(lx),std::min(Float(hx),this->x0)))); 88 | //printf("%f %f %f\n",this->x0,this->x1,this->x2); 89 | //printf("%f %f %f\n",this->y0,this->y1,this->y2); 90 | if (x>=this->x0 && x<=this->x2){ 91 | setUpY(); 92 | }else{ 93 | x=hx+1; 94 | yhb=hy; 95 | y=yhb+1; 96 | ybeg=0; 97 | ubeg=0; 98 | udelta=0; 99 | vbeg=0; 100 | vdelta=0; 101 | } 102 | } 103 | void setUpY(){ 104 | Float yb,yc,ub,uc,vb,vc; 105 | if (x0==x2){ 106 | if (y0yc){ 127 | yc=y2; 128 | uc=0; 129 | vc=1; 130 | } 131 | }else{ 132 | ub=0; 133 | vb=(Float(x)-x0)/(x2-x0); 134 | yb=y0+(y2-y0)*vb; 135 | if (x<=x1){ 136 | if (x1!=x0){ 137 | uc=(Float(x)-x0)/(x1-x0); 138 | vc=0; 139 | yc=y0+(y1-y0)*uc; 140 | }else{ 141 | uc=1; 142 | vc=0; 143 | yc=y1; 144 | } 145 | }else{ 146 | Float f=(Float(x)-x1)/(x2-x1); 147 | yc=y1+f*(y2-y1); 148 | uc=1-f; 149 | vc=f; 150 | } 151 | if (yb>yc){ 152 | std::swap(yb,yc); 153 | std::swap(ub,uc); 154 | std::swap(vb,vc); 155 | } 156 | } 157 | y=int(std::ceil(std::max(Float(ly),std::min(Float(hy+1),yb)))); 158 | yhb=int(std::floor(std::max(Float(ly-1),std::min(Float(hy),yc)))); 159 | ybeg=yb; 160 | ubeg=ub; 161 | vbeg=vb; 162 | if (yb!=yc){ 163 | udelta=(uc-ub)/(yc-yb); 164 | vdelta=(vc-vb)/(yc-yb); 165 | }else{ 166 | udelta=0; 167 | vdelta=0; 168 | } 169 | //printf("setUpY x=%d y=%d %d u %f %f v %f %f ybeg %f\n",x,y,yhb,ubeg,udelta,vbeg,vdelta,ybeg); 170 | } 171 | bool getNext(int *x_out,int *y_out,Float *u,Float *v){ 172 | 173 | while (y>yhb){ 174 | x++; 175 | if (x>hx || x>std::floor(x2)){ 176 | return false; 177 | } 178 | setUpY(); 179 | } 180 | x_out[0]=x; 181 | y_out[0]=y; 182 | if (u || v){ 183 | Float u_=ubeg+(Float(y)-ybeg)*udelta; 184 | Float v_=vbeg+(Float(y)-ybeg)*vdelta; 185 | if (u) 186 | u[0]=u0+(u1-u0)*u_+(u2-u0)*v_; 187 | if (v) 188 | v[0]=v0+(v1-v0)*u_+(v2-v0)*v_; 189 | } 190 | y++; 191 | return true; 192 | } 193 | }; 194 | 195 | template 196 | std::vector,std::pair > > allPixels(Float x0,Float y0,Float x1,Float y1,Float x2,Float y2,int lx,int hx,int ly,int hy){ 197 | PixelEnumerator pixels(x0,y0,x1,y1,x2,y2,lx,hx,ly,hy); 198 | int x,y; 199 | Float u,v; 200 | std::vector,std::pair > > ret; 201 | while (pixels.getNext(&x,&y,&u,&v)){ 202 | ret.push_back(std::make_pair( 203 | std::make_pair(x,y),std::make_pair(u,v))); 204 | } 205 | return ret; 206 | } 207 | 208 | }//namespace 209 | 210 | #endif 211 | -------------------------------------------------------------------------------- /rastertriangle_so.cpp: -------------------------------------------------------------------------------- 1 | #include "rastertriangle.h" 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | using namespace cv; 10 | 11 | 12 | 13 | extern "C"{ 14 | 15 | //n * 3 * 2 floats in xys 16 | void rasterTrianglesRGB(int h,int w,unsigned char *img,int n,float * xys,unsigned char * colors){ 17 | for (int i=0;i pixels( 19 | xys[i*6+0], 20 | xys[i*6+1], 21 | xys[i*6+2], 22 | xys[i*6+3], 23 | xys[i*6+4], 24 | xys[i*6+5], 25 | 0,h-1,0,w-1); 26 | //h/6,h-h/6-1,w/6,w-w/6-1); 27 | int x,y; 28 | float u,v; 29 | while (pixels.getNext(&x,&y,&u,&v)){ 30 | img[(x*w+y)*3+0]=colors[i*9+0*3+0]*(1-u-v)+colors[i*9+1*3+0]*u+colors[i*9+2*3+0]*v; 31 | img[(x*w+y)*3+1]=colors[i*9+0*3+1]*(1-u-v)+colors[i*9+1*3+1]*u+colors[i*9+2*3+1]*v; 32 | img[(x*w+y)*3+2]=colors[i*9+0*3+2]*(1-u-v)+colors[i*9+1*3+2]*u+colors[i*9+2*3+2]*v; 33 | } 34 | } 35 | } 36 | void getAllPixels_phase1(int lx,int hx,int ly,int hy,float * xy3,void ** handle,int *n){ 37 | std::vector,std::pair > > *all_pixels=new std::vector,std::pair > >(RasterTriangle::allPixels( 38 | xy3[0],xy3[1],xy3[2],xy3[3],xy3[4],xy3[5],lx,hx,ly,hy)); 39 | handle[0]=all_pixels; 40 | n[0]=all_pixels->size(); 41 | } 42 | void getAllPixels_phase2(void ** handle,int * xys,float * uvs){ 43 | std::vector,std::pair > > *all_pixels=(std::vector,std::pair > >*)handle[0]; 44 | int n=all_pixels->size(); 45 | for (int i=0;i,std::pair > > *all_pixels=(std::vector,std::pair > >*)handle[0]; 57 | int n=all_pixels->size(); 58 | for (int i=0;ipoints_2d[j*2]) 125 | swap(i,j); 126 | if (points_2d[j*2]>points_2d[k*2]) 127 | swap(j,k); 128 | if (points_2d[i*2]>points_2d[j*2]) 129 | swap(i,j); 130 | 131 | if ((ii>points_2d[j*2]) && (iimax(cof[i],max(cof[j],cof[k])) ) 148 | cof_ret = max(cof[i],max(cof[j],cof[k])); 149 | 150 | return cof_ret; 151 | 152 | } 153 | 154 | 155 | bool isline(float* xy3){ 156 | float ax = xy3[0]-xy3[4]; 157 | float ay = xy3[1]-xy3[5]; 158 | float bx = xy3[2]-xy3[4]; 159 | float by = xy3[3]-xy3[5]; 160 | float cross = ax*by-bx*ay; 161 | if (abs(cross)<1e-5) 162 | return 1; 163 | return 0; 164 | 165 | } 166 | 167 | 168 | void cover_rgbd(float* p , float* pz, float* pr, float* pg, float* pb, int m1,int m2,int m3, int idx, int* triangle ,float* zbuf, int* rbuf, int* gbuf, int* bbuf, int h,int w, int * xys, int debug){ 169 | float xy3[6]; 170 | xy3[0] = p[m1*2+0]; 171 | xy3[1] = p[m1*2+1]; 172 | xy3[2] = p[m2*2+0]; 173 | xy3[3] = p[m2*2+1]; 174 | xy3[4] = p[m3*2+0]; 175 | xy3[5] = p[m3*2+1]; 176 | 177 | if (debug == 1){ 178 | printf("%d, %d, %d\n", m1, m2, m3); 179 | } 180 | 181 | if (isline(xy3)) 182 | return; 183 | 184 | int num = pixelsInTriangle(h,w,xy3,xys); 185 | for (int i =0; i=tw)|(y2>=th)){ 272 | rgbzbuf[y*w*4+x*4+0] = trgb[y1*tw*3+x1*3+0]; 273 | rgbzbuf[y*w*4+x*4+1] = trgb[y1*tw*3+x1*3+1]; 274 | rgbzbuf[y*w*4+x*4+2] = trgb[y1*tw*3+x1*3+2]; 275 | }else{ 276 | rgbzbuf[y*w*4+x*4+0] = bilinear_inter( 277 | x1, y1, x2, y2, tx, ty, trgb[y1*tw*3+x1*3+0], 278 | trgb[y1*tw*3+x2*3+0], trgb[y2*tw*3+x1*3+0], 279 | trgb[y2*tw*3+x2*3+0] 280 | ); 281 | rgbzbuf[y*w*4+x*4+1] = bilinear_inter( 282 | x1, y1, x2, y2, tx, ty, trgb[y1*tw*3+x1*3+1], 283 | trgb[y1*tw*3+x2*3+1], trgb[y2*tw*3+x1*3+1], 284 | trgb[y2*tw*3+x2*3+1] 285 | ); 286 | rgbzbuf[y*w*4+x*4+2] = bilinear_inter( 287 | x1, y1, x2, y2, tx, ty, trgb[y1*tw*3+x1*3+2], 288 | trgb[y1*tw*3+x2*3+2], trgb[y2*tw*3+x1*3+2], 289 | trgb[y2*tw*3+x2*3+2] 290 | ); 291 | } 292 | rgbzbuf[y*w*4+x*4+3] = zz; 293 | triangle[y*w+x] = idx; 294 | } 295 | } 296 | } 297 | 298 | void uv_rgbzbuffer( 299 | int h,int w, float* points_onface, float* points_onface_ori, float* points_z, float* points_u, float* points_v, int len_mesh, int* mesh, 300 | int th, int tw, float* tmap, float* rgbzbuf 301 | ){ 302 | 303 | int * triangle = (int *) malloc(sizeof(int)*h*w); 304 | 305 | for (int i=0; i 1e-8).astype('uint8') 162 | if len(np.where(msk.flatten() > 0)[0]) < 500: 163 | continue 164 | zbuf *= msk.astype(zbuf.dtype) # * 1000.0 165 | 166 | bbuf.resize((h, w)), rbuf.resize((h, w)), gbuf.resize((h, w)) 167 | bgr = np.concatenate((bbuf[:, :, None], gbuf[:, :, None], rbuf[:, :, None]), axis=2) 168 | bgr = bgr.astype('uint8') 169 | 170 | bg = None 171 | len_bg_lst = len(self.bg_img_pth_lst) 172 | while bg is None or len(bg.shape) < 3: 173 | bg_pth = self.bg_img_pth_lst[randint(0, len_bg_lst-1)] 174 | bg = cv2.imread(bg_pth) 175 | if len(bg.shape) < 3: 176 | bg = None 177 | continue 178 | bg_h, bg_w, _ = bg.shape 179 | if bg_h < h: 180 | new_w = int(float(h) / bg_h * bg_w) 181 | bg = cv2.resize(bg, (new_w, h)) 182 | bg_h, bg_w, _ = bg.shape 183 | if bg_w < w: 184 | new_h = int(float(w) / bg_w * bg_h) 185 | bg = cv2.resize(bg, (w, new_h)) 186 | bg_h, bg_w, _ = bg.shape 187 | if bg_h > h: 188 | sh = randint(0, bg_h-h) 189 | bg = bg[sh:sh+h, :, :] 190 | bg_h, bg_w, _ = bg.shape 191 | if bg_w > w: 192 | sw = randint(0, bg_w-w) 193 | bg = bg[:, sw:sw+w, :] 194 | 195 | msk_3c = np.repeat(msk[:, :, None], 3, axis=2) 196 | bgr = bg * (msk_3c <= 0).astype(bg.dtype) + bgr * (msk_3c).astype(bg.dtype) 197 | rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) 198 | 199 | if args.vis: 200 | try: 201 | from neupeak.utils.webcv2 import imshow, waitKey 202 | except ImportError: 203 | from cv2 import imshow, waitKey 204 | imshow("bgr", bgr.astype("uint8")) 205 | show_zbuf = zbuf.copy() 206 | min_d, max_d = show_zbuf[show_zbuf > 0].min(), show_zbuf.max() 207 | show_zbuf[show_zbuf > 0] = (show_zbuf[show_zbuf > 0] - min_d) / (max_d - min_d) * 255 208 | show_zbuf = show_zbuf.astype(np.uint8) 209 | imshow("dpt", show_zbuf) 210 | show_msk = (msk / msk.max() * 255).astype("uint8") 211 | imshow("msk", show_msk) 212 | waitKey(0) 213 | 214 | data = {} 215 | data['depth'] = zbuf 216 | data['rgb'] = rgb 217 | data['mask'] = msk 218 | data['K'] = self.K 219 | data['RT'] = RT 220 | data['cls_typ'] = self.cls_type 221 | data['rnd_typ'] = 'render' 222 | sv_pth = os.path.join(self.render_dir, "{}.pkl".format(idx)) 223 | if DEBUG: 224 | imshow("rgb", rgb[:, :, ::-1].astype("uint8")) 225 | imshow("depth", (zbuf / zbuf.max() * 255).astype("uint8")) 226 | imshow("mask", (msk / msk.max() * 255).astype("uint8")) 227 | waitKey(0) 228 | pkl.dump(data, open(sv_pth, "wb")) 229 | pth_lst.append(os.path.abspath(sv_pth)) 230 | 231 | plst_pth = os.path.join(self.render_dir, "file_list.txt") 232 | with open(plst_pth, 'w') as of: 233 | for pth in pth_lst: 234 | print(pth, file=of) 235 | 236 | 237 | def main(): 238 | cls_type = 'cam' 239 | render_num = 70000 240 | if len(args.cls) > 0: 241 | cls_type = args.cls 242 | if args.render_num > 0: 243 | render_num = args.render_num 244 | print("cls: ", cls_type) 245 | gen = LineModRenderDB(cls_type, render_num, True) 246 | gen.gen_pack_zbuf_render() 247 | 248 | # for cls_type in OBJ_ID_DICT.keys(): 249 | # pth_lst = [] 250 | # render_num = 70000 251 | # gen = LineModRenderDB(cls_type, render_num, True) 252 | # gen.gen_pack_zbuf_render() 253 | 254 | 255 | 256 | if __name__ == "__main__": 257 | main() 258 | # vim: ts=4 sw=4 sts=4 expandtab 259 | -------------------------------------------------------------------------------- /sampled_poses/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ethnhe/raster_triangle/b31501a0f90b1e7f1270862a8d7decb435ea6f6a/sampled_poses/.gitkeep --------------------------------------------------------------------------------