Add test coverage for preprocessing modes and implement resize_mode support

- Added new tests in `tests/test_preprocessing.py` to cover stretch, crop, and pad preprocessing modes.
- Refactored `serve.py` to include a reusable `preprocess_image` function that handles multiple resize modes (`pad`, `crop`, and `stretch`).
- Updated `read_image_from_conn` and related methods to support dynamic resize modes based on an optional `resize_mode` parameter.
- Expanded README with detailed explanations of resize modes and updated request examples.

Author: artryazanov

README.md

@@ -91,8 +91,17 @@ Steps:
     {
         "type": "predict",
         "len": 12345
+    {
+        "type": "predict",
+        "len": 12345,
+        "resize_mode": "pad"
     }
     ```
+    
+    **Resize Modes (`resize_mode`):**
+    *   `pad` (Default): Pads the image with black borders to preserve aspect ratio (adds bars), then resizes to 256x256.
+    *   `crop`: Center-crops a square from the image, then resizes to 256x256.
+    *   `stretch`: Stretches the image to fit 256x256 (may distort aspect ratio).
 3.  **Send Image**:
     *   **Option A (Recommended):** Send a standard image file (PNG, BMP, JPG). The server uses `cv2.imdecode` to parse it automatically.
     *   **Option B (Fallback):** Send **196,608 bytes** of raw RGB pixel data (256x256). If `len` matches exactly, it is treated as raw buffer.

scripts/serve.py

@@ -13,6 +13,60 @@
 # Lock to ensure thread safety for the stateful InferenceSession
 session_lock = threading.Lock()
 
+def preprocess_image(img, mode="pad"):
+    """
+    Resizes image to 256x256 based on the mode:
+    - stretch: simple resize (default old behavior)
+    - crop: center crop to square, then resize
+    - pad: pad with black to square, then resize (default new behavior)
+    """
+    target_size = (256, 256)
+    # Check if image is valid
+    if img is None:
+        return None
+        
+    h, w = img.shape[:2]
+
+    if mode == "stretch":
+        if (w, h) != target_size:
+            return cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
+        return img
+
+    elif mode == "crop":
+        min_dim = min(h, w)
+        if h != w:
+            center_h, center_w = h // 2, w // 2
+            half_dim = min_dim // 2
+            start_h = max(0, center_h - half_dim)
+            start_w = max(0, center_w - half_dim)
+            end_h = start_h + min_dim
+            end_w = start_w + min_dim
+            img = img[start_h:end_h, start_w:end_w]
+        
+        if img.shape[:2] != (256, 256):
+             return cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
+        return img
+
+    elif mode == "pad":
+        max_dim = max(h, w)
+        if h != w:
+            top = (max_dim - h) // 2
+            bottom = max_dim - h - top
+            left = (max_dim - w) // 2
+            right = max_dim - w - left
+            img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])
+        
+        if img.shape[:2] != (256, 256):
+             return cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
+        return img
+    
+    # Fallback to pad if unknown mode
+    if h != w:
+        return preprocess_image(img, "pad")
+    if img.shape[:2] != target_size:
+        return cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
+    return img
+
 def handle_request(session, request, raw_image=None):
     """Universal request handler for ZeroMQ+Pickle and TCP+JSON+RawBytes protocols."""
     with session_lock:
@@ -29,7 +83,7 @@ def handle_request(session, request, raw_image=None):
         return {"status": "error", "message": "Unknown type"}
 
 
-def read_image_from_conn(conn, expected_size=None):
+def read_image_from_conn(conn, expected_size=None, resize_mode='pad'):
     """
     Reads an image from the connection.
     If expected_size is provided, reads exactly that many bytes.
@@ -53,8 +107,7 @@ def read_image_from_conn(conn, expected_size=None):
                      # OpenCV loads as BGR. 
                      # We need RGB.
                      img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-                     if img.shape[0] != 256 or img.shape[1] != 256:
-                          img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
+                     img = preprocess_image(img, resize_mode)
                      return img
              except Exception:
                  pass
@@ -143,7 +196,7 @@ def read_image_from_conn(conn, expected_size=None):
                 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
                 if actual_width != 256 or actual_height != 256:
-                    img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
+                    img = preprocess_image(img, resize_mode)
 
                 return img
         except struct.error:
@@ -233,7 +286,8 @@ def run_tcp_server(session, port):
                 img = None
                 if req.get("type") == "predict":
                     expected_len = req.get("len")
-                    img = read_image_from_conn(conn, expected_size=expected_len)
+                    resize_mode = req.get("resize_mode", "pad")
+                    img = read_image_from_conn(conn, expected_size=expected_len, resize_mode=resize_mode)
                     if img is None:
                         print("Incomplete or invalid image data received")
                         break

tests/test_preprocessing.py

@@ -0,0 +1,118 @@
+import sys
+import os
+import pytest
+from unittest.mock import MagicMock, call
+
+# Add scripts to path
+sys.path.append(os.path.join(os.path.dirname(__file__), "../scripts"))
+
+# Import serve (mocks are already applied by conftest.py)
+try:
+    import serve
+except ImportError:
+    serve = None
+
+import cv2
+import numpy as np
+
+@pytest.mark.skipif(serve is None, reason="Dependencies missing")
+class TestPreprocessing:
+    
+    def setup_method(self):
+        # Reset mocks before each test
+        cv2.reset_mock()
+        # Setup common mock behavior
+        # cv2.resize returns a new mock
+        cv2.resize.return_value = MagicMock(shape=(256, 256, 3))
+        cv2.copyMakeBorder.return_value = MagicMock(shape=(256, 256, 3))
+
+    def test_preprocess_stretch(self):
+        # Setup Mock Image
+        img = MagicMock()
+        img.shape = (200, 100, 3) # Height 200, Width 100
+        
+        # Call
+        res = serve.preprocess_image(img, "stretch")
+        
+        # Verify
+        cv2.resize.assert_called_once()
+        args, kwargs = cv2.resize.call_args
+        assert args[0] == img # First arg is image
+        assert args[1] == (256, 256) # Target size
+        assert kwargs.get('interpolation') == cv2.INTER_AREA
+
+    def test_preprocess_stretch_no_op(self):
+        # If already 256x256
+        img = MagicMock()
+        img.shape = (256, 256, 3)
+        
+        res = serve.preprocess_image(img, "stretch")
+        
+        # Should return original image without resize
+        cv2.resize.assert_not_called()
+        assert res == img
+
+    def test_preprocess_crop(self):
+        # 100x200 (Height 100, Width 200)
+        img = MagicMock()
+        img.shape = (100, 200, 3)
+        # Slicing returns a new mock
+        sliced_img = MagicMock()
+        sliced_img.shape = (100, 100, 3) # After crop it should be square
+        img.__getitem__.return_value = sliced_img
+        
+        res = serve.preprocess_image(img, "crop")
+        
+        # Verify Slicing (Center Crop)
+        # min_dim = 100. Center w=100. start_w = 50. end_w = 150.
+        # Img should be sliced [0:100, 50:150]
+        # Since we can't easily check slice args on __getitem__ with simple mocks without complex setup,
+        # we focus on the fact that it was sliced and then resized.
+        
+        img.__getitem__.assert_called()
+        
+        # And then resized
+        cv2.resize.assert_called_once()
+        args, kwargs = cv2.resize.call_args
+        assert args[0] == sliced_img # Should resize the sliced result
+        assert args[1] == (256, 256)
+
+    def test_preprocess_pad(self):
+        # 200x100 (Height 200, Width 100)
+        img = MagicMock()
+        img.shape = (200, 100, 3)
+        
+        padded_img = MagicMock()
+        padded_img.shape = (200, 200, 3) # Square after padding
+        cv2.copyMakeBorder.return_value = padded_img
+        
+        res = serve.preprocess_image(img, "pad")
+        
+        # Verify Padding
+        # max_dim = 200. Padding on left/right.
+        # Height 200. Width 100.
+        # top=0, bottom=0.
+        # left=(200-100)//2=50. right=50.
+        
+        cv2.copyMakeBorder.assert_called_once_with(
+            img, 0, 0, 50, 50, cv2.BORDER_CONSTANT, value=[0, 0, 0]
+        )
+        
+        # And then resized
+        cv2.resize.assert_called_once()
+        args, kwargs = cv2.resize.call_args
+        assert args[0] == padded_img
+        assert args[1] == (256, 256)
+
+    def test_preprocess_default(self):
+        # Default -> Pad
+        img = MagicMock()
+        img.shape = (200, 100, 3)
+        
+        padded_img = MagicMock()
+        padded_img.shape = (200, 200, 3)
+        cv2.copyMakeBorder.return_value = padded_img
+        
+        res = serve.preprocess_image(img) # No mode
+        
+        cv2.copyMakeBorder.assert_called_once() # Should use pad logic