fix checking logic

weecology · Dec 5, 2023 · 3a01d21 · 3a01d21
1 parent 52bc353
commit 3a01d21
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepforest/preprocess.py b/deepforest/preprocess.py
@@ -187,8 +187,8 @@ def split_raster(annotations_file,
             raise (IOError("If passing an numpy_image, please also specify a image_name"
                            " to match the column in the annotation.csv file"))
 
-    # Confirm that raster is H x W x C, if not, convert
-    if numpy_image.shape[0] < numpy_image.shape[1]:
+    # Confirm that raster is H x W x C, if not, convert, assuming image is wider/taller than channels
+    if numpy_image.shape[0] < numpy_image.shape[-1]:
         warnings.warn(
             "Input rasterio had shape {}, assuming channels first. Converting to channels last".format(
                 numpy_image.shape), UserWarning)