基于Keras的多类别CNN图像分类模型改造：Bounding Box预测实现问询

Hey there! Let's walk through exactly what you need to modify to add bounding box prediction to your existing CNN, plus the visualization part you mentioned. Your current model is built for pure classification—we'll expand it to a multi-task learning model that handles both category prediction and bounding box regression (predicting x,y,w,h coordinates).

1. Fix the Existing Model's Minor Bug First #

Quick note: Your original Dense(128, input_dim=11) line has an incorrect input_dim parameter. Sequential models automatically infer input shapes from the previous layer, so remove this to avoid shape mismatches:

# Replace this incorrect line
# classifier.add(Dense(128, input_dim = 11, activation = 'relu'))
# With this:
classifier.add(Dense(128, activation = 'relu'))

2. Rewrite the Model for Multi-Task Output #

Your current model only outputs 4 class probabilities. We need to add a regression branch to predict the 4 bounding box coordinates. Sequential isn't ideal for multi-task setups (since we need different activation functions for classification vs regression), so switch to Keras' Functional API for clarity:

from keras.models import Model
from keras.layers import Input, Convolution2D, MaxPooling2D, Flatten, Dense

# Input layer (matches your original input shape)
input_layer = Input(shape=(64, 64, 3))

# Reuse your existing convolutional/pooling layers
x = Convolution2D(filters=32, kernel_size=(3,3), data_format="channels_last", activation="relu")(input_layer)
x = MaxPooling2D(pool_size=(2,2))(x)
x = Convolution2D(filters=32, kernel_size=(3,3), data_format="channels_last", activation="relu")(x)
x = MaxPooling2D(pool_size=(2,2))(x)
x = Flatten()(x)
x = Dense(128, activation='relu')(x)

# Branch 1: Classification output (4 classes, softmax for single-class prediction)
class_output = Dense(4, activation='softmax', name='class_output')(x)

# Branch 2: Bounding box regression output (4 continuous values, linear activation)
bbox_output = Dense(4, activation='linear', name='bbox_output')(x)

# Build the multi-task model
classifier = Model(inputs=input_layer, outputs=[class_output, bbox_output])

3. Adjust the Loss Function & Compile Step #

We need separate loss functions for each task:

• Classification: categorical_crossentropy (matches your original setup)
• Bounding box regression: mse (mean squared error, ideal for continuous value prediction)

You can also weight the tasks if one is more important than the other:

classifier.compile(
    optimizer='adam',
    loss={
        'class_output': 'categorical_crossentropy',
        'bbox_output': 'mse'
    },
    loss_weights={
        'class_output': 1.0,  # Adjust weights based on your priority
        'bbox_output': 0.5
    },
    metrics={'class_output': 'accuracy'}
)

4. Prepare Training Data with Bounding Box Labels #

Your current flow_from_directory only loads class labels—we need to load the corresponding (x,y,w,h) coordinates from your CSV. Use a custom data generator (or flow_from_dataframe if your CSV includes image paths):

Assumptions about your CSV: #

Let's say your CSV has columns: image_path, class_name, x, y, w, h

import pandas as pd
from keras.preprocessing.image import ImageDataGenerator

# Load CSV data
train_df = pd.read_csv('train_bbox.csv')
test_df = pd.read_csv('test_bbox.csv')

# Map class names to numeric labels (reuse your original class mapping)
class_mapping = {v: k for k, v in training_set.class_indices.items()}
train_df['class'] = train_df['class_name'].map(class_mapping)
test_df['class'] = test_df['class_name'].map(class_mapping)

# Data augmentation (same as your original setup)
train_datagen = ImageDataGenerator(rescale=1./255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True)
test_datagen = ImageDataGenerator(rescale=1./255)

# Custom generator to return images + [class labels, bbox coordinates]
def multi_task_generator(df, datagen, target_size=(64,64), batch_size=32):
    base_gen = datagen.flow_from_dataframe(
        dataframe=df,
        x_col='image_path',
        y_col='class',
        target_size=target_size,
        batch_size=batch_size,
        class_mode='raw',
        shuffle=True
    )
    while True:
        imgs, class_labels = next(base_gen)
        # Get current batch's bounding box coordinates
        batch_start = base_gen.batch_index * batch_size
        batch_end = batch_start + batch_size
        bboxes = df.iloc[batch_start:batch_end][['x','y','w','h']].values
        
        # Critical: Rescale bbox coordinates to match model input size (64x64)
        # Example: If your original images are 256x256, scale by 64/256 = 0.25
        # Replace with your actual original image dimensions
        orig_w, orig_h = 256, 256
        scale_x = target_size[0] / orig_w
        scale_y = target_size[1] / orig_h
        bboxes = bboxes * [scale_x, scale_y, scale_x, scale_y]
        
        yield imgs, [class_labels, bboxes]

# Initialize generators
train_generator = multi_task_generator(train_df, train_datagen)
test_generator = multi_task_generator(test_df, test_datagen)

5. Train the Multi-Task Model #

Update your fit_generator call to use the new multi-task generator:

classifier.fit_generator(
    train_generator,
    steps_per_epoch=len(train_df)//32,
    epochs=25,
    validation_data=test_generator,
    validation_steps=len(test_df)//32
)

6. Test & Visualize Predictions #

Once trained, load test images, predict both class and bounding box, then draw the box and label on the image (using OpenCV or PIL):

import cv2
import numpy as np
from keras.preprocessing import image

# Load a test image
test_img_path = 'dataset/Test/class2/sample_img.jpg'
img = image.load_img(test_img_path, target_size=(64,64))
img_array = image.img_to_array(img) / 255.0
img_array = np.expand_dims(img_array, axis=0)

# Predict class and bounding box
class_pred, bbox_pred = classifier.predict(img_array)

# Get predicted class name
pred_class_idx = np.argmax(class_pred)
pred_class = [k for k, v in class_mapping.items() if v == pred_class_idx][0]

# Rescale bbox back to original image dimensions
orig_img = cv2.imread(test_img_path)
orig_h, orig_w = orig_img.shape[:2]
scale_x = orig_w / 64
scale_y = orig_h / 64
x, y, w, h = bbox_pred[0]
x = int(x * scale_x)
y = int(y * scale_y)
w = int(w * scale_x)
h = int(h * scale_y)

# Draw bounding box and label
cv2.rectangle(orig_img, (x, y), (x+w, y+h), (0, 255, 0), 2)
cv2.putText(orig_img, pred_class, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2)

# Save or display the result
cv2.imwrite('predicted_result.jpg', orig_img)
cv2.imshow('Prediction', orig_img)
cv2.waitKey(0)

• Normalize Bounding Boxes: For training, normalize coordinates to the [0,1] range (divide x/w by image width, y/h by image height) to stabilize training.
• Use Pre-Trained Models: Swap your simple CNN with a pre-trained model like VGG16 or ResNet as a feature extractor—this will drastically improve bounding box prediction accuracy.
• Tune Loss Weights: If bounding box prediction is more critical, increase the bbox_output loss weight (e.g., set it to 1.0 and class_output to 0.5).

内容的提问来源于stack exchange，提问作者vignesh charan