generate_tfrecord.py

def class_text_to_int(klassen_label):
  if klassen_label == 'person':
    return 1
  else:

def split(df, gruppiereBei):
  data = namedtuple('data', ['filename', 'object'])
  gb = df.groupby(gruppiereBei)
  return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]

def bmp_to_jpg(bmp_data):
  bmp_bild = tf.image.decode_bmp(bmp_data, channels = 3)
  return tf.Session().run(tf.image.encode_jpeg(bmp_bild, format='rgb', quality = 100, channels = 3))

def is_bmp(bild_datei):
  return ".bmp" in bild_datei

def create_tf_example(gruppe, path):
  with tf.gfile.GFile(os.path.join(path, '{}'.format(gruppe.filename)), 'rb') as fid:
  raw_bild = fid.read()
  if(is_bmp(gruppe.filename)):
    raw_bild = bmp_to_jpg(raw_bild)
    decoded_bild = tf.image.decode_jpeg(raw_bild, channels = 3)
    breite = tf.Session().run(tf.shape(decoded_bild)[0])
    hoehe = tf.Session().run(tf.shape(decoded_bild)[1])
    depth = tf.Session().run(tf.shape(decoded_bild)[2])

    dateiname = gruppe.filename.encode('utf8')
    bild_format = b'JPEG'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    klassen_text = []
    klassen = []

    for index, zeile in gruppe.object.iterrows():
      xmins.append(zeile['xmin'] / breite)
      xmaxs.append(zeile['xmax'] / breite)
      ymins.append(zeile['ymin'] / hoehe)
      ymaxs.append(zeile['ymax'] / hoehe)

    klassen_text.append(zeile['class'].encode('utf8'))

    klassen.append(class_text_to_int(zeile['class']))

    tf_example = tf.train.Example(features=tf.train.Features(feature={
    #'image/height': dataset_util.int64_feature(hoehe),
    #'image/width': dataset_util.int64_feature(breite),
    #'image/depth': dataset_util.int64_feature(depth),
    #'image/filename': dataset_util.bytes_feature(dateiname),
    #'image/source_id': dataset_util.bytes_feature(dateiname),
    #'image/encoded': dataset_util.bytes_feature(raw_bild),
    #'image/format': dataset_util.bytes_feature(bild_format),
    'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
    'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
    'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
    'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
    'image/class/text': dataset_util.bytes_list_feature(klassen_text),
    'image/class/label': dataset_util.int64_list_feature(klassen),
        }))
    return tf_example

  def main(_):
    current_pfad = os.getcwd()
    input_pfad = os.path.join(current_pfad, 'data/')
    output_pfad = input_pfad
    for csv_datei in os.listdir(input_pfad):
      if csv_datei.endswith(".csv"):
        output_file = csv_datei.replace("_labels.csv",".record")
        output_file = os.path.join(output_pfad, output_file)
        writer = tf.python_io.TFRecordWriter(output_file)
        examples = pd.read_csv(os.path.join(input_pfad, csv_datei))
        gruppiert = split(examples, 'filename')
          for gruppe in gruppiert:
            tf_example = create_tf_example(gruppe,
            os.path.join(current_pfad+"/images/",
            csv_datei.replace("_labels.csv","")))
            writer.write(tf_example.SerializeToString())
          writer.close()