[SubImage]

- handle different image shapes

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):

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))
+	img_rbb = np.reshape(np.linspace(1, n*c*w*h, num=n*c*w*h, dtype=np.float64), (n, c, w, h))
-#	print(img)
-
 	sub = SubImage(4, 4, 1, 1)
+	sub(img_rbb, 3)
 
-	images = sub(img)
+	img_gray = np.reshape(np.linspace(1, n*w*h, num=n*w*h, dtype=np.float64), (n, w, h))
-	print(images)
+	sub(img_gray, 1)
 
 
 	print("Test: [passed]")