from learning.utils.trainer import Trainer
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
#%%
data = Dataset([
'cylinder-cube-1',
'cylinder-1',
'cylinder-2',
'cube-1',
])
train_set, validation_set = data.load_data(
force_images=False,
suffixes=[['ed-v']],
# suffixes=['ed-side_b-0.400'],
)
#%%
def define_model(number_actions):
image = tk.Input(shape=(None, None, 1), name='image')
bayes_conv_block = bayes_conv_block_gen()
x = bayes_conv_block(image, 32, kernel_size=(5, 5), strides=(2, 2))
x = bayes_conv_block(x, 32)
x = bayes_conv_block(x, 48)
x = bayes_conv_block(x, 48)
#%%
data = Dataset([
'small-cubes-2',
'cylinder-cube-mc-1',
'cylinder-cube-1',
# 'cylinder-1',
'cylinder-2',
'cube-1',
'cube-3',
],
indexer=GraspIndexer(gripper_classes=Config.gripper_classes)
)
train_set, validation_set = data.load_data(
force_images=False,
suffixes=[
['ed-v'],
],
)
#%%
def define_model(number_actions: float, number_types: float):
image = tk.Input(shape=(None, None, 1), name='image')
conv_block = conv_block_gen(l2_reg=0.02, dropout_rate=0.4)
x = conv_block(image, 32, kernel_size=(5, 5))
x = conv_block(x, 32)
x = conv_block(x, 48)
x = conv_block(x, 48)
x = conv_block(x, 48, strides=(2, 2))
'cylinder-2',
'cylinder-cube-1',
'cylinder-cube-2',
'cylinder-cube-mc-1',
'shifting',
'small-cubes-2',
],
indexer=GraspShiftIndexer(gripper_classes=[0.05, 0.07, 0.086],
shift_distance=0.03),
)
train_set, validation_set = data.load_data(
# max_number=5000,
force_images=False,
scale_around_zero=True,
size_cropped=(256, 256),
size_output=(64, 64),
suffixes=[
['ed-after', 'ed-v'],
],
)
#%%
image_generator_type = 'Bicycle'
if image_generator_type == 'Pix2Pix':
image_generator = Pix2Pix(
training_generator=DataGenerator(train_set, shuffle=True),
validation_generator=DataGenerator(validation_set, shuffle=True),
result_path=data.result_path,
generator_path=data.model_path / 'predict-gen-6.h5',
class TFRecord(object):
def __init__(self):
self.data_path = path_params['data_path']
self.tfrecord_dir = path_params['tfrecord_dir']
self.train_tfrecord_name = path_params['train_tfrecord_name']
self.input_width = model_params['input_width']
self.input_height = model_params['input_height']
self.channels = model_params['channels']
self.grid_height = model_params['grid_height']
self.grid_width = model_params['grid_width']
self.class_num = len(model_params['classes'])
self.batch_size = solver_params['batch_size']
self.dataset = Dataset()
# 数值形式的数据,首先转换为string,再转换为int形式进行保存
def _int64_feature(self, value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
def _float_feature(self, value):
return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
# 数组形式的数据,首先转换为string,再转换为二进制形式进行保存
def _bytes_feature(self, value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def create_tfrecord(self):
# 获取作为训练验证集的图片序列
trainval_path = os.path.join(self.data_path, 'trainval.txt')
tf_file = os.path.join(self.tfrecord_dir, self.train_tfrecord_name)
if os.path.exists(tf_file):
os.remove(tf_file)
writer = tf.python_io.TFRecordWriter(tf_file)
with open(trainval_path, 'r') as read:
lines = read.readlines()
for line in lines:
filename = line[0:-1]
image_raw, bbox_raw, image_shape = self.dataset.load_data(
filename)
example = tf.train.Example(features=tf.train.Features(
feature={
'image':
tf.train.Feature(bytes_list=tf.train.BytesList(
value=[image_raw])),
'bbox':
tf.train.Feature(bytes_list=tf.train.BytesList(
value=[bbox_raw])),
'height':
tf.train.Feature(int64_list=tf.train.Int64List(
value=[image_shape[0]])),
'width':
tf.train.Feature(int64_list=tf.train.Int64List(
value=[image_shape[1]])),
}))
writer.write(example.SerializeToString())
writer.close()
print('Finish trainval.tfrecord Done')
def parse_single_example(self, serialized_example):
"""
:param serialized_example:待解析的tfrecord文件的名称
:return: 从文件中解析出的单个样本的相关特征，image, label
"""
# 解析单个样本文件
features = tf.parse_single_example(serialized_example,
features={
'image':
tf.FixedLenFeature([],
tf.string),
'bbox':
tf.FixedLenFeature([],
tf.string),
'height':
tf.FixedLenFeature([],
tf.int64),
'width':
tf.FixedLenFeature([], tf.int64)
})
# 进行解码
tf_image = tf.decode_raw(features['image'], tf.uint8)
tf_bbox = tf.decode_raw(features['bbox'], tf.float32)
tf_height = features['height']
tf_width = features['width']
# 转换为网络输入所要求的形状
tf_image = tf.reshape(tf_image, [tf_height, tf_width, 3])
tf_label = tf.reshape(tf_bbox, [150, 5])
tf_image, y_true = tf.py_func(self.dataset.preprocess_true_data,
inp=[tf_image, tf_label],
Tout=[tf.float32, tf.float32])
y_true = tf.reshape(y_true, [self.grid_height, self.grid_width, 5, 6])
return tf_image, y_true
def create_dataset(self,
filenames,
batch_num,
batch_size=1,
is_shuffle=False):
"""
:param filenames: record file names
:param batch_size: batch size
:param is_shuffle: whether shuffle
:param n_repeats: number of repeats
:return:
"""
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(self.parse_single_example, num_parallel_calls=4)
if is_shuffle:
dataset = dataset.shuffle(batch_num)
dataset = dataset.batch(batch_size)
dataset = dataset.repeat()
dataset = dataset.prefetch(batch_size)
return dataset
def test_dataset(self):
ds = Dataset('all-1')
ds.load_data(
force_images=True,
suffixes=[['ed-v']],
)