r

alignment3d/tfrecord2mx.py

@@ -1,47 +0,0 @@
-
-import tensorflow as tf
-import cv2
-import sys
-import mxnet as mx
-import os
-import numpy as np
-
-input_dir = sys.argv[1]
-output_dir = './data'
-writer = mx.recordio.MXIndexedRecordIO(os.path.join(output_dir, 'train.idx'), os.path.join(output_dir, 'train.rec'), 'w')
-
-idx = 1
-for _file in os.listdir(input_dir):
-  if not _file.endswith('tfrecords'):
-    continue
-  data_file = os.path.join(input_dir, _file)
-  for serialized_example in tf.python_io.tf_record_iterator(data_file):
-    example = tf.train.Example()
-    example.ParseFromString(serialized_example)
-    features = example.features.feature
-    image = features['image'].bytes_list.value[0]
-    width = features['width'].int64_list.value[0]
-    height = features['height'].int64_list.value[0]
-    image = np.fromstring(image, dtype=np.uint8)
-    image = cv2.imdecode(image, cv2.CV_LOAD_IMAGE_COLOR)
-    #print(image.shape)
-    n_landmarks = features['n_landmarks'].int64_list.value[0]
-    mask_index = features['mask_index'].bytes_list.value[0]
-    status = features['status'].int64_list.value[0]
-    gt_mask = features['gt_mask'].bytes_list.value[0]
-    gt_mask = np.fromstring(gt_mask, dtype=np.uint8)
-    #print(gt_mask.shape)
-    gt_pts = features['gt_pts'].bytes_list.value[0]
-    gt_pts = np.fromstring(gt_pts, dtype=np.float32)
-    #print(gt_pts.shape, n_landmarks)
-    #print(gt_pts)
-    #for k in features:
-    #  print(k)
-
-    #print(len(image),width, height, n_landmarks, status, gt_mask, gt_pts)
-    nlabel = list(gt_pts)
-    nheader = mx.recordio.IRHeader(0, nlabel, idx, 0)
-    s = mx.recordio.pack_img(nheader, image, quality=95, img_fmt='.jpg')
-    writer.write_idx(idx, s)
-    idx+=1
-