如何使用 ImageDataGenerator 和 OpenCV 训练 CNN 的裁剪图像

Question

我有一个由 ['images', 'xmin', 'ymin', 'xmax', 'ymax', 'label'] 组成的头盔数据框。 images 列是图像的绝对路径，xmin、ymin、xmax 和 ymax 是边界框值和标签（with_helmet 和 without_helmet）。

我想使用 Keras ImageDataGenerator 为 target_size(150, 150) 的图像创建训练数据集。

Answer 1

这是代码

类 CustomBoundingBoxGenerator(ImageDataGenerator): def init(self, *args, **kwargs): super().init(*args, **kwargs)

def flow_from_dataframe(self, dataframe, x_col, y_col, bounding_box_cols, batch_size=32, target_size=(150, 150), class_mode='binary', subset='training'):
    if subset == 'training':
        dataframe, _ = train_test_split(dataframe, test_size=0.2, random_state=42)
    elif subset == 'validation':
        _, dataframe = train_test_split(dataframe, test_size=0.2, random_state=42)
    else:
        raise ValueError("Invalid subset. Use 'training' or 'validation'.")

    generator = super().flow_from_dataframe(dataframe, x_col=x_col, y_col=y_col, batch_size=batch_size, target_size=target_size, class_mode=class_mode)

    while True:
        batch_x, batch_y = next(generator)

        processed_images = []
        processed_labels = []

        for i in range(len(batch_x)):
            img_path = dataframe[x_col].iloc[i]
            label = batch_y[i]

            img = cv.imread(img_path)

            xmin, ymin, xmax, ymax = dataframe.loc[dataframe.index[i], bounding_box_cols]

            b_box = (xmin, ymin, xmax, ymax)
            x, y, w, h = map(int, b_box)

            cropped_image = img[y:y+(h-y), x:x+(w-x)]


            cropped_image = cv.resize(cropped_image, target_size)

            processed_images.append(cropped_image)
            processed_labels.append(label)


        processed_images = np.array(processed_images)
        processed_labels = np.array(processed_labels)

        yield processed_images, processed_labels