diff --git a/atr_data_generator/extract/base.py b/atr_data_generator/extract/base.py
index 7295e1312b2545299057a20ff184b0f1d58bf14f..c899df16d3083d93e4ed0d57f0b9fd19795e899d 100644
--- a/atr_data_generator/extract/base.py
+++ b/atr_data_generator/extract/base.py
@@ -123,7 +123,8 @@ class DataGenerator:
         )
 
         # Resize to required height if needed
-        if self.image.fixed_height and bbox.height != self.image.fixed_height:
+        # image.shape = (height, width)
+        if self.image.fixed_height and image.shape[0] != self.image.fixed_height:
             image = np.asarray(
                 resize_image_height(
                     img=Image.fromarray(image),