As stated by title, I am building a playing card detection model that will identify cards and keep score. I am currently running into a problem right out of the gate. I’m unsure of why the shape is incorrect. I have pulled one image, used img_to_array and got something like (250,176,3).

import os
import numpy as np
import cv2
import math
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from PIL import Image
from ultralytics import YOLO
from keras._tf_keras.keras.utils import img_to_array, array_to_img, load_img
from keras._tf_keras.keras.preprocessing.image import ImageDataGenerator

# setup model
model = YOLO('../model/yolov8s_playing_cards.pt')
names = model.names

# setup paths
TRAIN_PATH = "../data/train/images"
TEST_PATH = "../data/test/images"

# seed and batch size
SEED = 1324
BATCH_SIZE = 32

# generator to flip, rotate, and shift height and width
train_gen = ImageDataGenerator(horizontal_flip=True,
                               rotation_range=20,
                               width_shift_range=[-2,2],
                               height_shift_range=[-2,2])

def plot_iter(obj, n=3):
    """ plot generator images
    Args: 
        obj (iter): image iter
        n (int, optional): plot matrix size . default to 3
    """
    fig, ax = plt.subplots(n, n)
    for i in range(n):
        for j in range(n):
            batch = next(obj)
            image = batch[0].astype('uint8')
            ax[i][j].set_axis_off()
            ax[i][j].imshow(image)
    plt.show()
    
it = train_gen.flow_from_directory(directory=TRAIN_PATH, batch_size=BATCH_SIZE)
plot_iter(it)

The error provided:

TypeError                                 Traceback (most recent call last)
Cell In[23], line 47
     44     plt.show()
     46 it = train_gen.flow_from_directory(directory=TRAIN_PATH, batch_size=BATCH_SIZE)
---> 47 plot_iter(it)

Cell In[23], line 43
     41         image = batch[0].astype('uint8')
     42         ax[i][j].set_axis_off()
---> 43         ax[i][j].imshow(image)
     44 plt.show()

File ~/miniconda3/envs/tenfl/lib/python3.10/site-packages/matplotlib/__init__.py:1473, in _preprocess_data.<locals>.inner(ax, data, *args, **kwargs)
   1470 @functools.wraps(func)
   1471 def inner(ax, *args, data=None, **kwargs):
   1472     if data is None:
-> 1473         return func(
   1474             ax,
   1475             *map(sanitize_sequence, args),
   1476             **{k: sanitize_sequence(v) for k, v in kwargs.items()})
   1478     bound = new_sig.bind(ax, *args, **kwargs)
   1479     auto_label = (bound.arguments.get(label_namer)
   1480                   or bound.kwargs.get(label_namer))

    701     # - otherwise casting wraps extreme values, hiding outliers and
    702     # making reliable interpretation impossible.
    703     high = 255 if np.issubdtype(A.dtype, np.integer) else 1

TypeError: Invalid shape (0, 256, 256, 3) for image data

I’d imagine this is happening because I need to use np.expand_dims() and add one dimension but I am unsure how to go about doing this to the whole dataset. Any ideas?