Address remaining PR review comments across workflows and tensor/model handling

Author: tomas-gajarsky

.github/workflows/auto-release.yml

@@ -23,12 +23,7 @@ jobs:
       - name: Read version
         id: version
         run: |
-          VERSION=$(python3 -c "
-          import re
-          with open('pyproject.toml') as f:
-              match = re.search(r'^version\s*=\s*\"(.+?)\"', f.read(), re.MULTILINE)
-          print(match.group(1))
-          ")
+          VERSION=$(python3 -c "import re; print(re.search(r'^version\s*=\s*\"(.+?)\"', open('pyproject.toml').read(), re.MULTILINE).group(1))")
           echo "VERSION=$VERSION" >> $GITHUB_OUTPUT
           echo "TAG=v$VERSION" >> $GITHUB_OUTPUT
           echo "VERSION=$VERSION" >> $GITHUB_ENV
@@ -106,10 +101,7 @@ jobs:
 
       - name: Set VERSION variable
         run: |
-          VERSION=$(python3 -c 'import re
-          with open("pyproject.toml") as f:
-              match = re.search(r"^version\s*=\s*\"(.+?)\"", f.read(), re.MULTILINE)
-          print(match.group(1))')
+          VERSION=$(python3 -c "import re; print(re.search(r'^version\s*=\s*\"(.+?)\"', open('pyproject.toml').read(), re.MULTILINE).group(1))")
           echo "VERSION=$VERSION" >> $GITHUB_ENV
 
       - name: Login to Docker Hub

environment.yml

@@ -10,7 +10,7 @@ dependencies:
   - numpy>=1.18.2
   - pip>=20.0.2
   - python-json-logger>=2.0.0
-  - pytorch-cpu>=2.0.0
+  - pytorch-cpu>=2.3.0
   - requests>=2.20.0
   - timm>=0.6.0
   - torchvision>=0.15.0

facetorch/analyzer/core.py

@@ -278,19 +278,19 @@ def _read_input(
             image_tensor = numpy_to_chw_tensor(image_source)
             return self.reader.process_tensor(image_tensor, fix_img_size=fix_img_size)
 
-        opened_here = False
         if isinstance(image_source, bytes):
-            image_source = Image.open(io.BytesIO(image_source))
-            opened_here = True
+            with Image.open(io.BytesIO(image_source)) as img:
+                pil_image = img.convert("RGB") if img.mode != "RGB" else img.copy()
+            image_tensor = torchvision.transforms.functional.pil_to_tensor(pil_image)
+            return self.reader.process_tensor(image_tensor, fix_img_size=fix_img_size)
 
         if isinstance(image_source, Image.Image):
-            if image_source.mode != "RGB":
-                image_source = image_source.convert("RGB")
-            image_tensor = torchvision.transforms.functional.pil_to_tensor(
-                image_source
+            pil_image = (
+                image_source.convert("RGB")
+                if image_source.mode != "RGB"
+                else image_source
             )
-            if opened_here:
-                image_source.close()
+            image_tensor = torchvision.transforms.functional.pil_to_tensor(pil_image)
             return self.reader.process_tensor(image_tensor, fix_img_size=fix_img_size)
 
         raise TypeError(

facetorch/analyzer/detector/core.py

@@ -46,13 +46,24 @@ def run(self, data: ImageData) -> ImageData:
         Returns:
             ImageData: Image data object with Detection tensors and detected Face objects.
         """
-        img_h, img_w = data.tensor.shape[-2], data.tensor.shape[-1]
+        orig_tensor = data.tensor
+        img_h, img_w = orig_tensor.shape[-2], orig_tensor.shape[-1]
         data = self.preprocessor.run(data)
         logits = self.inference(data.tensor)
         data = self.postprocessor.run(data, logits)
 
         if data.tensor.shape[-2] != img_h or data.tensor.shape[-1] != img_w:
-            data.tensor = data.tensor[:, :, :img_h, :img_w]
+            data.tensor = orig_tensor
             data.set_dims()
 
+            if hasattr(data.det, "dets") and data.det.dets.numel() > 0:
+                data.det.dets[:, 0].clamp_(0, img_w)
+                data.det.dets[:, 2].clamp_(0, img_w)
+                data.det.dets[:, 1].clamp_(0, img_h)
+                data.det.dets[:, 3].clamp_(0, img_h)
+
+            data.faces = []
+            if hasattr(self.postprocessor, "_extract_faces"):
+                data = self.postprocessor._extract_faces(data)
+
         return data

facetorch/base.py

@@ -104,12 +104,15 @@ def run(self, path: str) -> ImageData:
         pass
 
     def process_tensor(self, tensor: torch.Tensor, fix_img_size: bool) -> ImageData:
-        """Read a tensor and return a data object containing a tensor of the image with
-        shape (batch, channels, height, width).
+        """Read an input tensor and normalize it to shape (B, C, H, W).
 
         Args:
-            tensor (torch.Tensor): Image tensor with values between 0-255. Accepted shapes: (H, W), (C, H, W), or (B, C, H, W) where C in {1, 3, 4}. Single-channel inputs are expanded to 3 channels, RGBA inputs have the alpha channel dropped.
-            fix_img_size (bool): Whether to resize the image to a fixed size. If False, the size_portrait and size_landscape are ignored. Default is False.
+            tensor (torch.Tensor): Image tensor with values between 0-255. Accepted
+                shapes are (H, W), (C, H, W), (H, W, C), or (B, C, H, W), where
+                C is in {1, 3, 4}. Unambiguous HWC tensors are converted to CHW.
+                Batched tensors currently support only B=1.
+            fix_img_size (bool): Whether to resize the image to a fixed size. If
+                False, size_portrait and size_landscape are ignored.
         """
 
         data = ImageData(path_input=None)
@@ -119,11 +122,43 @@ def process_tensor(self, tensor: torch.Tensor, fix_img_size: bool) -> ImageData:
             data.tensor = data.tensor.unsqueeze(0)
 
         if data.tensor.dim() == 3:
+            c0 = data.tensor.shape[0]
+            c2 = data.tensor.shape[2]
+            chw_like = c0 in (1, 3, 4)
+            hwc_like = c2 in (1, 3, 4)
+            if hwc_like and not chw_like:
+                data.tensor = data.tensor.permute(2, 0, 1)
+            elif not chw_like and not hwc_like:
+                raise ValueError(
+                    "Invalid 3D tensor shape. Expected CHW with C in {1,3,4} or "
+                    "HWC with channels in the last dimension."
+                )
+            elif chw_like and hwc_like:
+                raise ValueError(
+                    "Ambiguous 3D tensor layout: both first and last dimensions "
+                    "look like channel dimensions. Please pass CHW explicitly."
+                )
             data.tensor = data.tensor.unsqueeze(0)
 
-        if data.tensor.shape[1] == 1:
+        if data.tensor.dim() != 4:
+            raise ValueError(
+                f"Unsupported tensor rank {data.tensor.dim()}. Expected 2D, 3D, or 4D input."
+            )
+
+        if data.tensor.shape[0] != 1:
+            raise ValueError(
+                f"Batched tensor input is not supported yet. Expected B=1, got B={data.tensor.shape[0]}."
+            )
+
+        channels = data.tensor.shape[1]
+        if channels not in (1, 3, 4):
+            raise ValueError(
+                f"Unsupported channel count: {channels}. Expected channels in {{1,3,4}}."
+            )
+
+        if channels == 1:
             data.tensor = data.tensor.repeat(1, 3, 1, 1)
-        elif data.tensor.shape[1] == 4:
+        elif channels == 4:
             data.tensor = data.tensor[:, :3, :, :]
 
         data.tensor = data.tensor.to(self.device)
@@ -265,13 +300,14 @@ def _load_exported_model(self) -> torch.nn.Module:
             if any(k in err_msg for k in ("schema version", "serialized version", "example_inputs")):
                 raise RuntimeError(
                     f"Cannot load {self.path_local}: the .pt2 model was exported with a "
-                    f"different PyTorch version. The bundled models require torch >=2.3.0,<2.4.0. "
+                    f"different PyTorch version. The bundled models require torch >=2.3.0,<2.5.0. "
                     f"Current version: {torch.__version__}. Install a compatible version or "
                     f"re-export the model with your current PyTorch."
                 ) from e
             raise
         model = ep.module()
         model.to(self.device)
+        model.eval()
         return model
 
     def _load_native_model(self) -> torch.nn.Module:

gpu.environment.yml

@@ -11,7 +11,7 @@ dependencies:
   - numpy>=1.18.2
   - pip>=20.0.2
   - python-json-logger>=2.0.0
-  - pytorch-gpu>=2.0.0
+  - pytorch-gpu>=2.3.0
   - requests>=2.20.0
   - timm>=0.6.0
   - torchvision>=0.15.0

pyproject.toml

@@ -63,7 +63,7 @@ Documentation = "https://tomas-gajarsky.github.io/facetorch/facetorch/index.html
 
 [tool.uv]
 extra-index-url = ["https://download.pytorch.org/whl/cpu"]
-constraint-dependencies = ["torch<2.4.0"]
+constraint-dependencies = ["torch<2.5.0"]
 
 [tool.setuptools.packages.find]
 exclude = ["tests*", "model_defs*", "scripts*"]

tests/test_base_model.py

@@ -210,3 +210,27 @@ def fake_download():
 
         m = ConcreteModel(downloader=dl, device=torch.device("cpu"))
         assert m.model is not None
+
+
+    def test_exported_model_is_on_device_and_eval(self, tmp_path):
+        bad_pt2 = str(tmp_path / "model.pt2")
+        with open(bad_pt2, "wb") as f:
+            f.write(b"placeholder")
+        dl = _make_dummy_downloader(bad_pt2)
+
+        class _FakeExported(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(2, 2)
+
+        fake_model = _FakeExported()
+
+        class _EP:
+            def module(self):
+                return fake_model
+
+        with patch("torch.export.load", return_value=_EP()):
+            m = ConcreteModel(downloader=dl, device=torch.device("cpu"))
+
+        assert m.model.training is False
+        assert next(m.model.parameters()).device.type == "cpu"

tests/test_reader.py

@@ -245,3 +245,27 @@ def test_process_tensor_rgba(analyzer):
     tensor_input = torch.randn(4, 224, 224)
     result = analyzer.reader.process_tensor(tensor_input, fix_img_size=False)
     assert result.tensor.size() == torch.Size([1, 3, 224, 224])
+
+
+@pytest.mark.unit
+@pytest.mark.reader
+def test_process_tensor_hwc_rgb(analyzer):
+    tensor_input = torch.randn(224, 224, 3)
+    result = analyzer.reader.process_tensor(tensor_input, fix_img_size=False)
+    assert result.tensor.size() == torch.Size([1, 3, 224, 224])
+
+
+@pytest.mark.unit
+@pytest.mark.reader
+def test_process_tensor_batched_not_supported(analyzer):
+    tensor_input = torch.randn(2, 3, 224, 224)
+    with pytest.raises(ValueError, match="B=1"):
+        analyzer.reader.process_tensor(tensor_input, fix_img_size=False)
+
+
+@pytest.mark.unit
+@pytest.mark.reader
+def test_process_tensor_ambiguous_3d_raises(analyzer):
+    tensor_input = torch.randn(3, 224, 3)
+    with pytest.raises(ValueError, match="Ambiguous 3D tensor layout"):
+        analyzer.reader.process_tensor(tensor_input, fix_img_size=False)