From 92afc16137d9a9ff6dcc51927f944f9e3c61be71 Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Sun, 11 Jan 2026 11:21:42 +0100
Subject: [PATCH] [SubImage] - handle different image shapes

---
 src/rbm/image.py            | 10 +++++++---
 src/tests/test_sub_image.py |  9 ++++-----
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/rbm/image.py b/src/rbm/image.py
index 6884603..83d8d44 100644
--- a/src/rbm/image.py
+++ b/src/rbm/image.py
@@ -13,8 +13,6 @@ class SubImage:
 		self.sy = sy
 		self.px = px
 		self.py = py
-		self.x = 0
-		self.y = 0
 
 	def check(self, image: Mat):
 		nx = (image.shape[0] - self.sx)
@@ -28,7 +26,13 @@ class SubImage:
 		ny = int(ny / self.py) + 1
 		return nx, ny
 
-	def __call__(self, images: Mat):
+	def __call__(self, images: Mat, n_ch: int = 3):
+		if len(images.shape) == 3:
+			return self.process_image(np.reshape(images, (images.shape[0], n_ch) + images.shape[1:3]))
+
+		if len(images.shape) == 2:
+			return self.process_image(np.reshape(images, (1,n_ch,) + images.shape)), images.shape
+
 		return self.process_image(images)
 
 	def process_image(self, image: Mat):
diff --git a/src/tests/test_sub_image.py b/src/tests/test_sub_image.py
index 0fc0885..338c9dc 100644
--- a/src/tests/test_sub_image.py
+++ b/src/tests/test_sub_image.py
@@ -7,13 +7,12 @@ if __name__ == "__main__":
 	c = 3
 	w = 8
 	h = 8
-	img = np.reshape(np.linspace(1, n*c*w*h, num=n*c*w*h, dtype=np.float64), (n, c, w, h))
-#	print(img)
-
+	img_rbb = np.reshape(np.linspace(1, n*c*w*h, num=n*c*w*h, dtype=np.float64), (n, c, w, h))
 	sub = SubImage(4, 4, 1, 1)
+	sub(img_rbb, 3)
 
-	images = sub(img)
-	print(images)
+	img_gray = np.reshape(np.linspace(1, n*w*h, num=n*w*h, dtype=np.float64), (n, w, h))
+	sub(img_gray, 1)
 
 
 	print("Test: [passed]")