ObjectDatasetTools生成linemod数据集预处理Linemod_preprocessed过程记录

在之前用ObjectDatasetTools制作真实数据后，还需要进一步处理才能用于算法训练。

主要包括以下预处理步骤：
在LINEMOD/Gold⽂件夹下依次创建如下python脚本并运行

import os
import numpy as np
import matplotlib.image


f = open("gt.txt", "w")

count = len(os.listdir("./transforms"))
for k in range(count):
    print('正在读取第'+str(k)+"张\n")
    data_load = np.load("transforms" + "/" + str(k) + ".npy")
    cam_r = []
    for i in range(3):
        for j in range(3):
            cam_r.append(data_load[i][j])      #cam_R_m2c

    cam_t = [data_load[0][3] * 1000,data_load[1][3] * 1000,data_load[2][3] * 1000]   #cam_t_m2c

    im = matplotlib.image.imread('mask/' + str(k) +'.png')
    r = []
    c = []
    ls1 = [0]
    ls2 = [0]

    for i in range(480):
        for j in range(1, 0):
            if im[i][j - 1] == 0 and im[i][j] == 1:
                r.append(i)
                c.append(j)
                break
    for i in range(480):
        for j in range(1, 0):
            if im[i][j - 1] == 1 and im[i][j] == 0:
                ls1[0] = i
                ls2[0] = j
    r.append(ls1[0])
    c.append(ls2[0])
    rmin = min(r)
    rmax = max(r)
    cmin = min(c)
    cmax = max(c)
    r.clear()
    c.clear()
    bb=[]
    bb.append(rmin)
    bb.append(rmax)
    bb.append(cmin)
    bb.append(cmax)
    print(cam_t)
    print(cam_r)
    print(bb)
    f.write("{}:\n".format(k))
    f.write("- cam_R_m2c: [{}, {}, {}, {}, {}, {}, {}, {}, {}]\n".format(cam_r[0],cam_r[1],cam_r[2],cam_r[3],cam_r[4],cam_r[5],cam_r[6],cam_r[7],cam_r[8]))
    f.write("  cam_t_m2c: [{}, {}, {}]\n".format(cam_t[0],cam_t[1],cam_t[2]))
    f.write("  obj_bb: [{}, {}, {}, {}]\n".format(bb[0],bb[1],bb[2],bb[3]))
    f.write("  obj_id: 1\n")
    cam_r.clear()
    bb.clear()
    cam_t.clear()
f.close()

os.rename("./gt.txt","./gt.yml")

运行后得到gt.yml文件

2. rename.py

import os
import cv2


root_rgb = "./JPEGImages/"
root_mask = "./mask/"
root_depth = "./depth/"

ls_rgb = os.listdir(root_rgb)
ls_mask = os.listdir(root_mask)
ls_depth = os.listdir(root_depth)

os.mkdir("rgb")

for file in ls_rgb:
    os.rename(root_rgb + file,"./rgb/" +"0" * int(4 - len(file[:-4])) + file[:-4] + ".jpg")
for file in ls_mask:
    os.rename(root_mask + file,root_mask +"0" * int(4 - len(file[:-4])) + file[:-4] + ".png")
for file in ls_depth:
    os.rename(root_depth + file,root_depth +"0" * int(4 - len(file[:-4])) + file[:-4] + ".png")

os.rmdir(root_rgb)

rgb = "./rgb"
ls__rgb = os.listdir(rgb)
i = 0
for file in ls__rgb:
    print("正在进行图片" + str(i) + "的转码")
    img = cv2.imread("./rgb/" + file)
    cv2.imwrite("./rgb/" + file[:-3] + "png",img)
    os.remove("./rgb/"+ file[:-3] + "jpg")
    i += 1

运行后会将jpg换成png格式，图片名字也会改变

3. info.py
   记得将fx,0.,cx,0.,fy,cy等参数按照intrinsics.json换成自己的相机参数

import yaml
import os

count = len(os.listdir("./rgb"))
for i in range(0,600):
    list = [fx,0.,cx,0.,fy,cy,0.,0.,1.]
    d={
        i:{
            "cam_K": list,
            "depth_scale": 0.001,
        }
    }
    f=open("info.yml","a",encoding="utf-8")
    yaml.dump(d,f)
    f.close()

4. calc_model_info.py

import yaml
import numpy as np
import struct
import math
import trimesh
 
def load_ply(path):
  """Loads a 3D mesh model from a PLY file.
  :param path: Path to a PLY file.
  :return: The loaded model given by a dictionary with items:
   - 'pts' (nx3 ndarray)
   - 'normals' (nx3 ndarray), optional
   - 'colors' (nx3 ndarray), optional
   - 'faces' (mx3 ndarray), optional
   - 'texture_uv' (nx2 ndarray), optional
   - 'texture_uv_face' (mx6 ndarray), optional
   - 'texture_file' (string), optional
  """
  f = open(path, 'rb')
 
  # Only triangular faces are supported.
  face_n_corners = 3
 
  n_pts = 0
  n_faces = 0
  pt_props = []
  face_props = []
  is_binary = False
  header_vertex_section = False
  header_face_section = False
  texture_file = None
 
  # Read the header.
  while True:
 
    # Strip the newline character(s).
    line = f.readline().decode('utf8').rstrip('\n').rstrip('\r')
 
    if line.startswith('comment TextureFile'):
      texture_file = line.split()[-1]
    elif line.startswith('element vertex'):
      n_pts = int(line.split()[-1])
      header_vertex_section = True
      header_face_section = False
    elif line.startswith('element face'):
      n_faces = int(line.split()[-1])
      header_vertex_section = False
      header_face_section = True
    elif line.startswith('element'):  # Some other element.
      header_vertex_section = False
      header_face_section = False
    elif line.startswith('property') and header_vertex_section:
      # (name of the property, data type)
      pt_props.append((line.split()[-1], line.split()[-2]))
    elif line.startswith('property list') and header_face_section:
      elems = line.split()
      if elems[-1] == 'vertex_indices' or elems[-1] == 'vertex_index':
        # (name of the property, data type)
        face_props.append(('n_corners', elems[2]))
        for i in range(face_n_corners):
          face_props.append(('ind_' + str(i), elems[3]))
      elif elems[-1] == 'texcoord':
        # (name of the property, data type)
        face_props.append(('texcoord', elems[2]))
        for i in range(face_n_corners * 2):
          face_props.append(('texcoord_ind_' + str(i), elems[3]))
      else:
        print('Warning: Not supported face property: ' + elems[-1])
    elif line.startswith('format'):
      if 'binary' in line:
        is_binary = True
    elif line.startswith('end_header'):
      break
 
  # Prepare data structures.
  model = {}
  if texture_file is not None:
    model['texture_file'] = texture_file
  model['pts'] = np.zeros((n_pts, 3), np.float)
  if n_faces > 0:
    model['faces'] = np.zeros((n_faces, face_n_corners), np.float)
 
  pt_props_names = [p[0] for p in pt_props]
  face_props_names = [p[0] for p in face_props]
 
  is_normal = False
  if {'nx', 'ny', 'nz'}.issubset(set(pt_props_names)):
    is_normal = True
    model['normals'] = np.zeros((n_pts, 3), np.float)
 
  is_color = False
  if {'red', 'green', 'blue'}.issubset(set(pt_props_names)):
    is_color = True
    model['colors'] = np.zeros((n_pts, 3), np.float)
 
  is_texture_pt = False
  if {'texture_u', 'texture_v'}.issubset(set(pt_props_names)):
    is_texture_pt = True
    model['texture_uv'] = np.zeros((n_pts, 2), np.float)
 
  is_texture_face = False
  if {'texcoord'}.issubset(set(face_props_names)):
    is_texture_face = True
    model['texture_uv_face'] = np.zeros((n_faces, 6), np.float)
 
  # Formats for the binary case.
  formats = {
    'float': ('f', 4),
    'double': ('d', 8),
    'int': ('i', 4),
    'uchar': ('B', 1)
  }
 
  # Load vertices.
  for pt_id in range(n_pts):
    prop_vals = {}
    load_props = ['x', 'y', 'z', 'nx', 'ny', 'nz',
                  'red', 'green', 'blue', 'texture_u', 'texture_v']
    if is_binary:
      for prop in pt_props:
        format = formats[prop[1]]
        read_data = f.read(format[1])
        val = struct.unpack(format[0], read_data)[0]
        if prop[0] in load_props:
          prop_vals[prop[0]] = val
    else:
      elems = f.readline().decode('utf8').rstrip('\n').rstrip('\r').split()
      for prop_id, prop in enumerate(pt_props):
        if prop[0] in load_props:
          prop_vals[prop[0]] = elems[prop_id]
 
    model['pts'][pt_id, 0] = float(prop_vals['x'])
    model['pts'][pt_id, 1] = float(prop_vals['y'])
    model['pts'][pt_id, 2] = float(prop_vals['z'])
 
    if is_normal:
      model['normals'][pt_id, 0] = float(prop_vals['nx'])
      model['normals'][pt_id, 1] = float(prop_vals['ny'])
      model['normals'][pt_id, 2] = float(prop_vals['nz'])
 
    if is_color:
      model['colors'][pt_id, 0] = float(prop_vals['red'])
      model['colors'][pt_id, 1] = float(prop_vals['green'])
      model['colors'][pt_id, 2] = float(prop_vals['blue'])
 
    if is_texture_pt:
      model['texture_uv'][pt_id, 0] = float(prop_vals['texture_u'])
      model['texture_uv'][pt_id, 1] = float(prop_vals['texture_v'])
 
  # Load faces.
  for face_id in range(n_faces):
    prop_vals = {}
    if is_binary:
      for prop in face_props:
        format = formats[prop[1]]
        val = struct.unpack(format[0], f.read(format[1]))[0]
        if prop[0] == 'n_corners':
          if val != face_n_corners:
            raise ValueError('Only triangular faces are supported.')
        elif prop[0] == 'texcoord':
          if val != face_n_corners * 2:
            raise ValueError('Wrong number of UV face coordinates.')
        else:
          prop_vals[prop[0]] = val
    else:
      elems = f.readline().decode('utf8').rstrip('\n').rstrip('\r').split()
      for prop_id, prop in enumerate(face_props):
        if prop[0] == 'n_corners':
          if int(elems[prop_id]) != face_n_corners:
            raise ValueError('Only triangular faces are supported.')
        elif prop[0] == 'texcoord':
          if int(elems[prop_id]) != face_n_corners * 2:
            raise ValueError('Wrong number of UV face coordinates.')
        else:
          prop_vals[prop[0]] = elems[prop_id]
 
    model['faces'][face_id, 0] = int(prop_vals['ind_0'])
    model['faces'][face_id, 1] = int(prop_vals['ind_1'])
    model['faces'][face_id, 2] = int(prop_vals['ind_2'])
 
    if is_texture_face:
      for i in range(6):
        model['texture_uv_face'][face_id, i] = float(
          prop_vals['texcoord_ind_{}'.format(i)])
 
  f.close()
  return model
 
def calc_pts_diameter(pts):
  """Calculates the diameter of a set of 3D points (i.e. the maximum distance
  between any two points in the set).
  :param pts: nx3 ndarray with 3D points.
  :return: The calculated diameter.
  """
  diameter = -1.0
  for pt_id in range(pts.shape[0]):
    pt_dup = np.tile(np.array([pts[pt_id, :]]), [pts.shape[0] - pt_id, 1])
    pts_diff = pt_dup - pts[pt_id:, :]
    max_dist = math.sqrt((pts_diff * pts_diff).sum(axis=1).max())
    if max_dist > diameter:
      diameter = max_dist
  return diameter
 
 
def distance(point_one, point_two):
    return ((point_one[0] - point_two[0]) ** 2 +
            (point_one[1] - point_two[1]) ** 2 + (point_one[2] - point_two[2]) ** 2) ** 0.5
 
def max_distance(points):
    return max(distance(p1, p2) for p1, p2 in zip(points, points[1:]))
 
 
def calc_pts_diameter1(path):
    mesh = trimesh.load(path)
    vertices = mesh.vertices
    maxD = max_distance(vertices.tolist())
    return maxD
	
 
if __name__=="__main__":
    obj_ids =[1]
    model_lpath = r"registeredScene.ply"
    models_info_path=r"models_info.yml"
 
    models_info = {}
    for obj_id in obj_ids:
        print('Processing model of object {}...'.format(obj_id))
 
        model = load_ply(model_lpath)
 
        ref_pt = model['pts'].min(axis=0).flatten()
        size = model['pts'].max(axis=0) - ref_pt
        print(ref_pt)
        print(size)
 
    # Calculated diameter.
        diameter = calc_pts_diameter(model['pts'])
        #diameter = calc_pts_diameter1(model_lpath)
        print(diameter)
 
        models_info[obj_id] = {
            'min_x': float(ref_pt[0]), 'min_y': float(ref_pt[1]), 'min_z': float(ref_pt[2]),
            'size_x': float(size[0]), 'size_y': float(size[1]), 'size_z': float(size[2]),
            'diameter': float(diameter)
        }
    print(yaml.dump(models_info))
    # Save the calculated info about the object models.
    #inout.save_json(dp_model['models_info_path'], models_info)
    with open(models_info_path, 'w') as f:
        yaml.dump(models_info, f)

运行后得到models_info.yml文件

5. re-format.py
   这一步主要是生成训练需要的文件目录结构，第21行os.rename(“./registeredScene.ply”, “./obj_01.ply”) ##记得改成自己那个ply的文件名

import os
import shutil

os.mkdir("./data")
os.mkdir("./data/01")
os.mkdir("./models")
os.mkdir("./segnet_results")
os.mkdir("./segnet_results/01_label")

msk = os.listdir("./mask")
for file in msk:
    shutil.copy("./mask/" + file,"./segnet_results/01_label/")


shutil.move("./rgb", "./data/01/")
shutil.move("./mask", "./data/01/")
shutil.move("./depth", "./data/01/")
shutil.move("./gt.yml", "./data/01/")
shutil.move("./info.yml", "./data/01/")

os.rename("./registeredScene.ply", "./obj_01.ply")   ##记得改成自己那个ply的文件名
shutil.move("./obj_01.ply", "./models")
shutil.move("./models_info.yml", "./models")
shutil.rmtree("./labels")
shutil.rmtree("./transforms")
'''os.remove("intrinsics.json")
os.remove("registeredScene.ply")
os.remove("transforms.npy")

root = os.listdir("./")
for file in root:
    if file[-2:] == "py":
        if file == "compute_model_info.py":
            continue
        os.remove(file)'''

os.mkdir("./Linemod_preprocessed")
shutil.move("data", "./Linemod_preprocessed/")
shutil.move("models", "./Linemod_preprocessed/")
shutil.move("segnet_results", "./Linemod_preprocessed/")

6. train_test_txt.py
   这一步主要是划分测试集和训练集

import os

files = len(os.listdir("./Linemod_preprocessed/data/01/depth/"))

_train = open("./Linemod_preprocessed/data/01/train.txt","w")
_test = open("./Linemod_preprocessed/data/01/test.txt","w")
for i in range(files):
    num = (4-len(str(i))) * '0' + str(i)
    if i % 5 == 4:
        _test.write(num + "\n")
    else:
        _train.write(num + "\n")
_train.close()
_test.close()

.
├── data
│   └── 01
├── models
│   └── obj_01.ply
│   └── models_info.yml
└── segnet_results
    └── 01_label