nh-rep/code/pre_processing/split_and_gen_fea.py

import os
import numpy as np
from load_obj import *
import yaml
import platform
sys_info = platform.system()

in_path = 'abc_data'
out_path = 'raw_input'

def split_single_sample(mesh_file):
  overwrite_files = True
  def normalize_mesh_padding(mesh, padding=0.05):
    coord_max = mesh.vertices.max(axis=0)
    coord_min = mesh.vertices.min(axis=0)
    center = (coord_max + coord_min) / 2.0
    bbox_length = (coord_max - coord_min).max()
    mesh.vertices = (mesh.vertices - center) * (1-padding) / bbox_length
    return mesh
  
  #sample points on each curve and save curve informations
  mesh_fileall = mesh_file
  mesh = load_obj_simple(mesh_file)
  vert2comp = color_connected_mesh(mesh)
  face2comp = [0] * len(mesh.faces)
  print('filename {} comp: {}'.format(mesh_file, max(vert2comp) + 1))
  nface = len(mesh.faces)
  for i in range(nface):
    face2comp[i] = vert2comp[mesh.faces[i][0]]
  ncomp = max(vert2comp) + 1
  comp2faces = [[] for i in range(ncomp)]
  for i in range(nface):
    comp2faces[face2comp[i]].append(i)
  perpatch_faceo2n = [{} for i in range(ncomp)]

  # hangning vert??
  for i in range(ncomp):
    if not len(comp2faces[i]) > 0:
      print('!!!! component with no face exists !!!!')
    assert(len(comp2faces[i]) > 0)
    for j in range(len(comp2faces[i])):
      perpatch_faceo2n[i][comp2faces[i][j]] = j



  yaml_file = mesh_file.replace(".obj", ".yml")
  assert(os.path.exists(yaml_file))
  with open(yaml_file, "r") as f: feature_dict_all = yaml.safe_load(f)
  
  curves_all = []
  patches_all = []
  if 'curves' in feature_dict_all:
    curves_all = feature_dict_all['curves']
  patches_all = feature_dict_all['surfaces']

  #classifiy patches and curves
  comp2curves = [[] for i in range(ncomp)]
  comp2patches = [[] for i in range(ncomp)]
  for i in range(len(curves_all)):
    cur_comp = vert2comp[curves_all[i]['vert_indices'][0]]
    comp2curves[cur_comp].append(curves_all[i])

  for i in range(len(patches_all)):
    cur_comp = face2comp[patches_all[i]['face_indices'][0]]
    comp2patches[cur_comp].append(patches_all[i])

  #update face id of comp2patches, remove vert id
  for i in range(ncomp):
    for j in range(len(comp2patches[i])):
      comp2patches[i][j].pop('vert_indices', None)
      for k in range(len(comp2patches[i][j]['face_indices'])):
        comp2patches[i][j]['face_indices'][k] = perpatch_faceo2n[i][comp2patches[i][j]['face_indices'][k]]

  for compid in range(ncomp):
    print('compid: ', compid)
    mesh_file = mesh_fileall.replace('.obj', '_{}.obj'.format(compid))
    curves = comp2curves[compid]
    patches = comp2patches[compid]
    feature_dict = {'curves': curves, 'surfaces': patches}
    #update vert id
    vert_valid_id = np.unique(mesh.faces[comp2faces[compid]].reshape(-1))
    # print('unique vert size:', vert_valid_id.shape[0])
    vert_o2n = -1 * np.ones(mesh.vertices.shape[0], dtype = np.int64)
    vert_o2n[vert_valid_id] = np.arange(vert_valid_id.shape[0])

    cur_faces = vert_o2n[mesh.faces[comp2faces[compid]].reshape(-1)].reshape(-1,3)
    cur_verts = mesh.vertices[vert_valid_id]
    write_obj_simple(os.path.join(out_path, os.path.basename(mesh_file)), cur_verts, cur_faces)
    output_yaml = os.path.join(out_path, os.path.basename(mesh_file.replace(".obj", ".yml")))

    #update face id of patches
    flag_cylinder_cube = False
    with open(output_yaml, 'w') as yaml_file:
      yaml.dump(feature_dict, yaml_file, default_flow_style=True)

    if len(feature_dict['surfaces']) == 3 or len(feature_dict['surfaces']) == 6:
      count_cylinder = 0
      count_plane = 0
      for i in range(len(feature_dict['surfaces'])):
        if feature_dict['surfaces'][i]['type'] == 'Cylinder':
          count_cylinder += 1
        if feature_dict['surfaces'][i]['type'] == 'Plane':
          count_plane += 1

      if count_cylinder == 1 and count_plane == 2:
        flag_cylinder_cube = True

      if count_cylinder == 0  and count_plane == 6:
        flag_cylinder_cube = True

    if flag_cylinder_cube:
      tmpfile = os.path.join(out_path, os.path.basename(mesh_file.replace(".obj", ".cylindercube")))

      f = open(tmpfile, 'w')
      f.close()

def sample_points_on_single_curve(points, sample_num):
  #calculate curve length
  curve_length = 0
  cum_line_length = []
  for i in range(points.shape[0]-1):
    length = np.sqrt((np.square(points[i+1] - points[i])).sum())
    curve_length += length
    cum_line_length.append(length)
  
  cum_line_length = np.cumsum(np.array(cum_line_length))
  segment_curve_length = curve_length / (sample_num - 1)
  sampled_points = np.zeros([sample_num, 3])
  sampled_points[0] = points[0]
  
  cur_line_id = 0
  for i in range(1, sample_num-1):
    target_length = i*segment_curve_length
    while(target_length > cum_line_length[cur_line_id]): cur_line_id+=1
    start_points = points[cur_line_id]
    end_points = points[cur_line_id+1]
    cur_line_length = np.sqrt((np.square(end_points - start_points)).sum())
    proportion = (cum_line_length[cur_line_id] - target_length) / cur_line_length
    if(proportion < -1e-5 or proportion > 1.00001):
      print("Assertion Violation: proportion out of range: {}".format(proportion))
    if(proportion < 0):
      proportion = 0
    elif(proportion > 1):
      proportion = 1
    sampled_points[i] = proportion*start_points + (1-proportion)*end_points
  sampled_points[sample_num-1] = points[-1]
  return sampled_points


def sample_points_on_each_curve(mesh, curves, sample_num=100):  
  curve_points = []
  for curve_id in range(len(curves)):
    curve_points.append(sample_points_on_single_curve(mesh.vertices[curves[curve_id]['vert_indices']], sample_num))
  if len(curve_points) > 0:
    return np.concatenate(curve_points, axis = 0)
  else:
    return np.array([])

def gen_fea_file(mesh_file):
  overwrite_files = True
  def normalize_mesh_padding(mesh, padding=0.05):
    coord_max = mesh.vertices.max(axis=0)
    coord_min = mesh.vertices.min(axis=0)
    center = (coord_max + coord_min) / 2.0
    bbox_length = (coord_max - coord_min).max()
    mesh.vertices = (mesh.vertices - center) * (1-padding) / bbox_length
    return mesh
  
  #sample points on each curve and save curve informations
  mesh_fileall = mesh_file
  mesh = load_obj_simple(mesh_file)
  vert2comp = color_connected_mesh(mesh)
  face2comp = [0] * len(mesh.faces)
  print('filename {} comp: {}'.format(mesh_file, max(vert2comp) + 1))
  nface = len(mesh.faces)
  for i in range(nface):
    face2comp[i] = vert2comp[mesh.faces[i][0]]
  ncomp = max(vert2comp) + 1
  comp2faces = [[] for i in range(ncomp)]
  for i in range(nface):
    comp2faces[face2comp[i]].append(i)
  perpatch_faceo2n = [{} for i in range(ncomp)]

  # hanging vert
  for i in range(ncomp):
    if not len(comp2faces[i]) > 0:
      print('!!!! component with no face exists !!!!')
    assert(len(comp2faces[i]) > 0)
    for j in range(len(comp2faces[i])):
      perpatch_faceo2n[i][comp2faces[i][j]] = j


  yaml_file = mesh_file.replace(".obj", ".yml")
  assert(os.path.exists(yaml_file))
  with open(yaml_file, "r") as f: feature_dict_all = yaml.safe_load(f)
  
  curves_all = []
  patches_all = []
  if 'curves' in feature_dict_all:
    curves_all = feature_dict_all['curves']
  patches_all = feature_dict_all['surfaces']

  #classifiy patches and curves
  comp2curves = [[] for i in range(ncomp)]
  comp2patches = [[] for i in range(ncomp)]
  for i in range(len(curves_all)):
    cur_comp = vert2comp[curves_all[i]['vert_indices'][0]]
    comp2curves[cur_comp].append(curves_all[i])

  for i in range(len(patches_all)):
    cur_comp = face2comp[patches_all[i]['face_indices'][0]]
    comp2patches[cur_comp].append(patches_all[i])

  #update face id of comp2patches, remove vert id
  for i in range(ncomp):
    for j in range(len(comp2patches[i])):
      comp2patches[i][j].pop('vert_indices', None)
      for k in range(len(comp2patches[i][j]['face_indices'])):
        comp2patches[i][j]['face_indices'][k] = perpatch_faceo2n[i][comp2patches[i][j]['face_indices'][k]]
  
  for compid in range(ncomp):
    mesh_file = mesh_fileall.replace('.obj', '_{}.obj'.format(compid))
    curves = comp2curves[compid]
    patches = comp2patches[compid]
    feature_dict = {'curves': curves, 'surfaces': patches}
    
    curve_sample_points_file = os.path.join(out_path, os.path.basename(mesh_file.replace(".obj", "_curve.sample.xyz")))
    if True:
      curve_sample_points = np.reshape(sample_points_on_each_curve(mesh, curves), [-1,3])
      np.savetxt(curve_sample_points_file, curve_sample_points)
    
    #update vert id
    vert_valid_id = np.unique(mesh.faces[comp2faces[compid]].reshape(-1))
    vert_o2n = -1 * np.ones(mesh.vertices.shape[0], dtype = np.int64)
    vert_o2n[vert_valid_id] = np.arange(vert_valid_id.shape[0])
    cur_faces = vert_o2n[mesh.faces[comp2faces[compid]].reshape(-1)].reshape(-1,3)
    cur_verts = mesh.vertices[vert_valid_id]
    write_obj_simple(os.path.join(out_path, os.path.basename(mesh_file)), cur_verts, cur_faces)
    
    #save feature file
    feature_list = []
    flag_only_sharp = False
    for curve_idx in range(len(curves)):
      if('closed' in curves[curve_idx]):
        assert(curves[curve_idx]['type'] == 'BSpline')
        closed_curve = curves[curve_idx]['closed']
      else:
        assert(curves[curve_idx]['type'] != 'BSpline')
        curve_vert_idx = curves[curve_idx]['vert_indices']
        if(curve_vert_idx[0] == curve_vert_idx[-1]):
          closed_curve = True
        else:
          closed_curve = False
      if flag_only_sharp and not curves[curve_idx]['sharp']:
        continue

      if True:
        for vid in range(len(curves[curve_idx]['vert_indices']) - 1):
          feature_list.append(vert_o2n[curves[curve_idx]['vert_indices'][vid]])
          feature_list.append(vert_o2n[curves[curve_idx]['vert_indices'][vid + 1]])
        if closed_curve:
          feature_list.append(vert_o2n[curves[curve_idx]['vert_indices'][-1]])
          feature_list.append(vert_o2n[curves[curve_idx]['vert_indices'][0]])

    feature_file = os.path.join(out_path, os.path.basename(mesh_file.replace(".obj", ".fea")))

    f = open(feature_file, 'w')
    f.write('{}\n'.format(len(feature_list)//2))
    for i in range(len(feature_list)//2):
      f.write('{} {}\n'.format(feature_list[2 * i], feature_list[2 * i + 1]))
    f.close()


def split_all_models():

  if not os.path.exists(out_path):
      os.mkdir(out_path)

  task_list = []
  #source folder set below
  file_list = os.listdir(in_path)
  for file in file_list:
    if file.endswith(".obj"):
      task_list.append(os.path.join(in_path, file))

  for t in task_list:
      split_single_sample(t)
      gen_fea_file(t)

def main():
  split_all_models()

if __name__ == '__main__':
  print('Program running...')
  main()