1 # -*- coding: utf-8 -*-
3 from __future__ import absolute_import
4 from __future__ import print_function
5 from __future__ import division
8 import tensorflow as tf
10 from data.io import image_preprocess
11 from libs.configs import cfgs
13 def read_single_example_and_decode(filename_queue):
15 # tfrecord_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
17 # reader = tf.TFRecordReader(options=tfrecord_options)
18 reader = tf.TFRecordReader()
19 _, serialized_example = reader.read(filename_queue)
21 features = tf.parse_single_example(
22 serialized=serialized_example,
24 'img_name': tf.FixedLenFeature([], tf.string),
25 'img_height': tf.FixedLenFeature([], tf.int64),
26 'img_width': tf.FixedLenFeature([], tf.int64),
27 'img': tf.FixedLenFeature([], tf.string),
28 'gtboxes_and_label': tf.FixedLenFeature([], tf.string),
29 'num_objects': tf.FixedLenFeature([], tf.int64)
32 img_name = features['img_name']
33 img_height = tf.cast(features['img_height'], tf.int32)
34 img_width = tf.cast(features['img_width'], tf.int32)
35 img = tf.decode_raw(features['img'], tf.uint8)
37 img = tf.reshape(img, shape=[img_height, img_width, 3])
39 gtboxes_and_label = tf.decode_raw(features['gtboxes_and_label'], tf.int32)
40 gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 5])
42 num_objects = tf.cast(features['num_objects'], tf.int32)
43 return img_name, img, gtboxes_and_label, num_objects
46 def read_and_prepocess_single_img(filename_queue, shortside_len, is_training):
48 img_name, img, gtboxes_and_label, num_objects = read_single_example_and_decode(filename_queue)
50 img = tf.cast(img, tf.float32)
53 img, gtboxes_and_label = image_preprocess.short_side_resize(img_tensor=img, gtboxes_and_label=gtboxes_and_label,
54 target_shortside_len=shortside_len,
55 length_limitation=cfgs.IMG_MAX_LENGTH)
56 img, gtboxes_and_label = image_preprocess.random_flip_left_right(img_tensor=img,
57 gtboxes_and_label=gtboxes_and_label)
60 img, gtboxes_and_label = image_preprocess.short_side_resize(img_tensor=img, gtboxes_and_label=gtboxes_and_label,
61 target_shortside_len=shortside_len,
62 length_limitation=cfgs.IMG_MAX_LENGTH)
63 img = img - tf.constant([[cfgs.PIXEL_MEAN]]) # sub pixel mean at last
64 return img_name, img, gtboxes_and_label, num_objects
67 def next_batch(dataset_name, batch_size, shortside_len, is_training):
70 img_name_batch: shape(1, 1)
71 img_batch: shape:(1, new_imgH, new_imgW, C)
72 gtboxes_and_label_batch: shape(1, Num_Of_objects, 5] .each row is [x1, y1, x2, y2, label]
74 assert batch_size == 1, "we only support batch_size is 1.We may support large batch_size in the future"
76 if dataset_name not in ['ship', 'spacenet', 'pascal', 'coco','pcb']:
77 raise ValueError('dataSet name must be in pascal, coco spacenet and ship')
80 pattern = os.path.join('../data/tfrecord', dataset_name + '_train*')
82 pattern = os.path.join('../data/tfrecord', dataset_name + '_test*')
84 print('tfrecord path is -->', os.path.abspath(pattern))
86 filename_tensorlist = tf.train.match_filenames_once(pattern)
88 filename_queue = tf.train.string_input_producer(filename_tensorlist)
90 img_name, img, gtboxes_and_label, num_obs = read_and_prepocess_single_img(filename_queue, shortside_len,
91 is_training=is_training)
92 img_name_batch, img_batch, gtboxes_and_label_batch, num_obs_batch = \
94 [img_name, img, gtboxes_and_label, num_obs],
95 batch_size=batch_size,
99 return img_name_batch, img_batch, gtboxes_and_label_batch, num_obs_batch