如何在TensorFlow的CNN测试阶段按类别保存预测图像？

Got it, let's get your test images sorted into category folders based on the model's predictions. First, I noticed a quick mismatch in your code: you mentioned your dataset has 5 classes, but FLAGS.num_classes is set to 3. We'll fix that first, then add the image-saving logic step by step.

Step 1: Prep Output Folders #

First, we'll define a root directory to store categorized images and create subfolders for each class automatically—no manual folder setup needed.

Step 2: Modify the Testing Loop #

We need to:

• Pull the predicted class for every image in each batch
• Grab the original image paths from your BatchPreprocessor (assuming it stores paths in an image_paths attribute; adjust if your setup uses a different name)
• Save each image to its predicted class folder

Here's the modified complete code with comments highlighting changes:

import os
import datetime
import numpy as np
import tensorflow as tf
from PIL import Image  # Use cv2 instead if you prefer OpenCV

# --- Updated: Set num_classes to match your 5-category dataset ---
tf.app.flags.DEFINE_float('learning_rate', 0.0001, 'Learning rate for adam optimizer')
tf.app.flags.DEFINE_integer('num_classes', 5, 'Number of classes (updated to 5)')
tf.app.flags.DEFINE_integer('batch_size', 128, 'Batch size')
tf.app.flags.DEFINE_float('keep_prob', 0.8, 'Dropout keep probability')
tf.app.flags.DEFINE_integer('num_channel',3 , 'Image channel, RGB=3, Grayscale=1')
tf.app.flags.DEFINE_integer('img_size', 80, 'Size of images')
tf.app.flags.DEFINE_string('test_file', 'data/test.txt', 'Test dataset file')
FLAGS = tf.app.flags.FLAGS
checkpoint_dir = '/home/xyrio/Desktop/classier/training/checkpoints/model_epoch.ckpt89'

# --- Added: Define and create output directories ---
OUTPUT_ROOT = './test_predicted_classes'
os.makedirs(OUTPUT_ROOT, exist_ok=True)
for class_idx in range(FLAGS.num_classes):
    class_folder = os.path.join(OUTPUT_ROOT, f'class_{class_idx}')
    os.makedirs(class_folder, exist_ok=True)

def main(_):
    # --- Adjusted placeholder to handle variable batch sizes (better for leftover images) ---
    x = tf.placeholder(tf.float32, shape=[None, FLAGS.img_size, FLAGS.img_size, FLAGS.num_channel], name='x')
    y_true = tf.placeholder(tf.float32, shape=[None, FLAGS.num_classes], name='y_true')
    y_true_cls = tf.argmax(y_true, axis=1)
    filter_size_conv1 = 3
    num_filters_conv1 = 32
    filter_size_conv2 = 3
    num_filters_conv2 = 32
    filter_size_conv3 = 3
    num_filters_conv3 = 64
    filter_size_conv4 = 3
    num_filters_conv4 = 128
    filter_size_conv5 = 3
    num_filters_conv5 = 256
    fc_layer_size = 512
    fc_layer_size2 = 128

    def create_weights(shape):
        return tf.Variable(tf.truncated_normal(shape, mean=0, stddev=0.01))

    def create_biases(size):
        return tf.Variable(tf.constant(0.01, shape=[size]))

    def create_convolutional_layer(input, num_input_channels, conv_filter_size, num_filters, useBatchNorm=False, usePooling=True):
        weights = create_weights(shape=[conv_filter_size, conv_filter_size, num_input_channels, num_filters])
        biases = create_biases(num_filters)
        layer = tf.nn.conv2d(input=input, filter=weights, strides=[1, 1, 1, 1], padding='SAME')
        layer += biases
        layer = tf.nn.relu(layer)
        if useBatchNorm == True:
            layer = tf.layers.batch_normalization(layer)
        if usePooling:
            layer = tf.nn.max_pool(value=layer, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        return layer

    def create_flatten_layer(layer):
        layer_shape = layer.get_shape()
        num_features = layer_shape[1:4].num_elements()
        layer = tf.reshape(layer, [-1, num_features])
        return layer

    def create_fc_layer(input, num_inputs, num_outputs, useRelu=True, useDropout=False):
        weights = create_weights(shape=[num_inputs, num_outputs])
        biases = create_biases(num_outputs)
        layer = tf.matmul(input, weights) + biases
        if useRelu:
            layer = tf.nn.relu(layer)
        if useDropout == True:
            layer = tf.nn.dropout(layer, keep_prob=FLAGS.keep_prob)
        return layer

    layer_conv1 = create_convolutional_layer(x, FLAGS.num_channel, filter_size_conv1, num_filters_conv1, useBatchNorm=True, usePooling=True)
    layer_conv2 = create_convolutional_layer(layer_conv1, num_filters_conv1, filter_size_conv2, num_filters_conv2, useBatchNorm=True, usePooling=True)
    layer_conv3 = create_convolutional_layer(layer_conv2, num_filters_conv2, filter_size_conv3, num_filters_conv3, useBatchNorm=True, usePooling=True)
    layer_conv4 = create_convolutional_layer(layer_conv3, num_filters_conv3, filter_size_conv4, num_filters_conv4, useBatchNorm=True, usePooling=True)
    layer_conv5 = create_convolutional_layer(layer_conv4, num_filters_conv4, filter_size_conv5, num_filters_conv5, useBatchNorm=True, usePooling=True)
    layer_flat = create_flatten_layer(layer_conv5)
    layer_fc1 = create_fc_layer(layer_flat, layer_flat.get_shape()[1:4].num_elements(), fc_layer_size, useRelu=True, useDropout=False)
    layer_fc2 = create_fc_layer(layer_fc1, fc_layer_size, fc_layer_size2, useRelu=True, useDropout=True)
    layer_fc3 = create_fc_layer(layer_fc2, fc_layer_size2, FLAGS.num_classes, useRelu=False)
    y_pred = tf.nn.softmax(layer_fc3, name='y_pred', axis=1)
    y_pred_cls = tf.argmax(y_pred, axis=1)
    correct_prediction = tf.equal(y_pred_cls, y_true_cls)
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    saver = tf.train.Saver()

    test_preprocessor = BatchPreprocessor(dataset_file_path=FLAGS.test_file, num_classes=FLAGS.num_classes, output_size=[FLAGS.img_size, FLAGS.img_size])
    test_batches_per_epoch = np.floor(len(test_preprocessor.labels) / FLAGS.batch_size).astype(np.int16)
    conf_mat = tf.confusion_matrix(y_true_cls,y_pred_cls,FLAGS.num_classes)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver.restore(sess, checkpoint_dir)
        # Start Testing
        test_acc = 0.
        test_count = 0
        cm_total = None
        current_idx = 0

        for _ in range(test_batches_per_epoch):
            batch_tx, batch_ty = test_preprocessor.next_batch(FLAGS.batch_size)
            # --- Added: Fetch predicted classes for the batch ---
            acc, conf_m, pred_cls = sess.run([accuracy, conf_mat, y_pred_cls], feed_dict={x: batch_tx, y_true: batch_ty})
            
            # --- Added: Save each image to its predicted class folder ---
            batch_image_paths = test_preprocessor.image_paths[current_idx:current_idx+FLAGS.batch_size]
            for img_path, pred in zip(batch_image_paths, pred_cls):
                # Load original image (use batch_tx if you want to save preprocessed version instead)
                img = Image.open(img_path)
                img_filename = os.path.basename(img_path)
                output_path = os.path.join(OUTPUT_ROOT, f'class_{pred}', img_filename)
                img.save(output_path)
            
            current_idx += FLAGS.batch_size

            if cm_total is None:
                cm_total = conf_m
            else:
                cm_total += conf_m
            test_acc += acc
            test_count += 1

        # --- Added: Handle leftover images that don't fill a full batch ---
        remaining_images = len(test_preprocessor.labels) - current_idx
        if remaining_images > 0:
            batch_tx, batch_ty = test_preprocessor.next_batch(remaining_images)
            acc, conf_m, pred_cls = sess.run([accuracy, conf_mat, y_pred_cls], feed_dict={x: batch_tx, y_true: batch_ty})
            batch_image_paths = test_preprocessor.image_paths[current_idx:]
            for img_path, pred in zip(batch_image_paths, pred_cls):
                img = Image.open(img_path)
                img_filename = os.path.basename(img_path)
                output_path = os.path.join(OUTPUT_ROOT, f'class_{pred}', img_filename)
                img.save(output_path)
            cm_total += conf_m
            test_acc += acc
            test_count += 1

        test_acc /= test_count
        print("{} Testing Accuracy: {:.4f}".format(datetime.now(), test_acc))
        print("Confusion Matrix:")
        print(cm_total)

if __name__ == '__main__':
    tf.app.run()

Key Details to Note: #

• Variable Batch Size: I adjusted the x placeholder to accept variable batch sizes, which makes handling leftover test images much cleaner.
• Image Saving: We save the original image (not the preprocessed version) so you can easily verify the model's predictions with the raw input. If you want to save the preprocessed image instead, convert the numpy array from batch_tx to an image using Image.fromarray().
• Folder Safety: os.makedirs with exist_ok=True ensures you won't get errors if the output folders already exist.

内容的提问来源于stack exchange，提问作者A.M