Fix: Image batch dimension not expanded.

This commit is contained in:
h-alice 2025-01-03 16:22:55 +08:00
parent bb820a7ef4
commit e1822851a5
No known key found for this signature in database
GPG Key ID: 5708F34144A70909

View File

@ -43,6 +43,6 @@ class Demography(ABC):
# Check input dimension
if len(image_batch.shape) == 3:
# Single image - add batch dimension
imgs = np.expand_dims(image_batch, axis=0)
image_batch = np.expand_dims(image_batch, axis=0)
return image_batch