def input_fn(mode):
import os
resource_path = os.path.join(os.path.split(__file__)[0], "../resources")
if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
image_folder = os.path.join(resource_path, "cat_dog")
image_set = ImageSet.read(image_folder, with_label=True, sc=self.sc,
one_based_label=False)
transformer = ChainedPreprocessing([ImageResize(256, 256),
ImageRandomCrop(224, 224, True),
ImageMatToTensor(format="NHWC"),
ImageSetToSample(input_keys=["imageTensor"],
target_keys=["label"])])
image_set = image_set.transform(transformer)
dataset = TFDataset.from_image_set(image_set,
image=(tf.float32, [224, 224, 3]),
label=(tf.int32, [1]),
batch_size=8)
else:
image_folder = os.path.join(resource_path, "cat_dog/*/*")
image_set = ImageSet.read(image_folder, with_label=False, sc=self.sc,
one_based_label=False)
transformer = ChainedPreprocessing([ImageResize(256, 256),
ImageRandomCrop(224, 224, True),
ImageMatToTensor(format="NHWC"),
ImageSetToSample(
input_keys=["imageTensor"])])
image_set = image_set.transform(transformer)
dataset = TFDataset.from_image_set(image_set,
image=(tf.float32, [224, 224, 3]),
batch_per_thread=8)
return dataset
def test_predict_for_imageset(self):
model = self.create_image_model()
image_set = self.create_image_set(with_label=False)
predict_dataset = TFDataset.from_image_set(image_set,
image=(tf.float32, [224, 224, 3]),
batch_per_thread=1)
results = model.predict(predict_dataset).get_predict().collect()
assert all(r[1] is not None for r in results)
def test_training_for_imageset(self):
model = self.create_image_model()
image_set = self.create_image_set(with_label=True)
training_dataset = TFDataset.from_image_set(image_set,
image=(tf.float32, [224, 224, 3]),
label=(tf.int32, [1]),
batch_size=4)
model.fit(training_dataset)
def test_evaluation_for_imageset(self):
model = self.create_image_model()
image_set = self.create_image_set(with_label=True)
eval_dataset = TFDataset.from_image_set(image_set,
image=(tf.float32, [224, 224, 3]),
label=(tf.int32, [1]),
batch_per_thread=1)
model.evaluate(eval_dataset)
