def get_single_picture_prediction(model, picture_dir):
image_tensor = load_and_preprocess_image(
tf.io.read_file(filename=picture_dir), data_augmentation=False)
image = tf.expand_dims(image_tensor, axis=0)
prediction = model(image, training=False)
pred_class = tf.math.argmax(prediction, axis=-1)
return pred_class
def main():
# 测试数据集，
image_raw = tf.io.read_file(filename=test_image_dir)
image_tensor = load_and_preprocess_image(image_raw)
image_tensor = tf.expand_dims(image_tensor, axis=0)
# Load frozen graph using TensorFlow 1.x functions
with tf.io.gfile.GFile("./frozen_models2/model.pb", "rb") as f:
graph_def = tf.compat.v1.GraphDef()
loaded = graph_def.ParseFromString(f.read())
# Wrap frozen graph to ConcreteFunctions
frozen_func = wrap_frozen_graph(graph_def=graph_def,
inputs=["Input:0"],
outputs=["Identity:0"],
print_graph=True)
print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
# Get predictions for test images
predictions = frozen_func(Input=image_tensor)[0]
# Print the prediction for the first image
print("-" * 50)
print("Example prediction reference:")
print(predictions[0].numpy())
def get_list_picture_prediction(model, picture_dir):
class_name_list = sorted(
[s for s in os.listdir(picture_dir) if 'README' not in s])
print('class name:')
print(class_name_list)
image_tensor_list = []
for class_name_item in class_name_list:
picture_path = os.path.join(picture_dir, class_name_item)
picture_path_list = os.listdir(picture_path)
for picture_path_item in picture_path_list:
image_path = os.path.join(picture_path, picture_path_item)
image_tensor = load_and_preprocess_image(
tf.io.read_file(filename=image_path), data_augmentation=False)
image_tensor_list.append(image_tensor)
image_tensor_list_length = len(image_tensor_list)
images = tf.stack(image_tensor_list, axis=0)
prediction = model(images, training=False)
pred_class_num_list = tf.math.argmax(prediction, axis=-1)
pred_class_name_list = []
true_num = 0
for index, pred_class_num_item in enumerate(pred_class_num_list):
if class_name_list[pred_class_num_item] == class_name_list[int(index //
5)]:
true_num = true_num + 1
pred_class_name_item = class_name_list[pred_class_num_item]
pred_class_name_list.append(pred_class_name_item)
return pred_class_name_list, image_tensor_list_length, true_num
def get_single_picture_prediction(model, picture_dir):
image_tensor = load_and_preprocess_image(
tf.io.read_file(filename=picture_dir), data_augmentation=False)
image = tf.expand_dims(image_tensor, axis=0)
prediction = model(image, training=False)
#print(np.shape(prediction.numpy()))
return prediction.numpy()
def process_features(features, data_augmentation):
image_raw = features['image_raw'].numpy()
image_tensor_list = []
for image in image_raw:
image_tensor = load_and_preprocess_image(image, data_augmentation=data_augmentation)
image_tensor_list.append(image_tensor)
images = tf.stack(image_tensor_list, axis=0)
labels = features['label'].numpy()
return images, labels
def get_single_picture_prediction(model, picture_dir):
class_name_list = [
'dirty_defect', 'edge_defect', 'linear_defect', 'no_defect',
'yarn_defect'
]
picture_list = os.listdir(picture_dir)
picture_path = os.path.join(picture_dir, picture_list[0])
image_tensor = load_and_preprocess_image(
tf.io.read_file(filename=picture_path), data_augmentation=False)
image = tf.expand_dims(image_tensor, axis=0)
prediction = model(image, training=False)
pred_class_num = tf.math.argmax(prediction, axis=-1)
pred_class_name = class_name_list[int(pred_class_num)]
return pred_class_name
def get_list_picture_prediction(model, picture_dir):
class_name_list = ['defect', 'no_defect']
if 'no_defect' in picture_dir:
picture_path_list = os.listdir(picture_dir)
elif 'defect' in picture_dir:
picture_path_list = os.listdir(picture_dir)
image_tensor_list = []
for picture_path_item in picture_path_list:
picture_path = os.path.join(picture_dir, picture_path_item)
image_tensor = load_and_preprocess_image(tf.io.read_file(filename=picture_path), data_augmentation=False)
image_tensor_list.append(image_tensor)
images = tf.stack(image_tensor_list, axis=0)
prediction = model(images, training=False)
pred_class_num_list = tf.math.argmax(prediction, axis=-1)
pred_class_name_list = []
for pred_class_num_item in pred_class_num_list:
pred_class_name_item = class_name_list[pred_class_num_item]
pred_class_name_list.append(pred_class_name_item)
return pred_class_name_list
id_cls[i] = item
return id_cls
if __name__ == '__main__':
# GPU settings
gpus = tf.config.list_physical_devices('GPU')
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
model = tf.keras.models.load_model('model.h5',
custom_objects={
'KerasLayer': hub.KerasLayer,
'Dense': tf.keras.layers.Dense
},
compile=False)
model.summary()
image_raw = tf.io.read_file(filename=test_image_dir)
image_tensor = load_and_preprocess_image(image_raw)
image_tensor = tf.expand_dims(image_tensor, axis=0)
pred = model(image_tensor, training=False)
print(pred)
idx = tf.math.argmax(pred, axis=-1).numpy()[0]
print(idx)
# id_cls = get_class_id("./original_dataset")
#
# print("The predicted category of this picture is: {}".format(id_cls[idx]))
