refactor: extract frame helper utilities

Author: ALexFIlik

app/workers/frame_helpers.py

@@ -0,0 +1,112 @@
+"""Pure frame preview and geometry helpers used by worker/viewer paths."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from app.utils.write_output import is_normalized_float_range
+
+
+def coerce_preview_frame(frame) -> np.ndarray | None:
+    arr = np.asarray(frame)
+    if arr.ndim == 2:
+        gray = arr.astype(np.float32)
+        min_val = float(np.nanmin(gray)) if gray.size else 0.0
+        max_val = float(np.nanmax(gray)) if gray.size else 0.0
+        if gray.dtype != np.uint8:
+            if is_normalized_float_range(min_val, max_val):
+                gray = np.clip(gray, 0.0, 1.0)
+                gray = np.power(gray, 1.0 / 2.2) * 255.0
+            gray = np.clip(gray, 0.0, 255.0).astype(np.uint8)
+        else:
+            gray = arr
+        return np.stack([gray] * 3, axis=-1)
+    if arr.ndim == 3 and arr.shape[2] == 1:
+        return coerce_preview_frame(arr[:, :, 0])
+    if arr.ndim == 3 and arr.shape[2] >= 3:
+        if arr.dtype != np.uint8 and arr.shape[2] >= 4:
+            rgb_premul = arr[:, :, :3].astype(np.float32)
+            alpha_ch = np.clip(arr[:, :, 3:4].astype(np.float32), 0.0, 1.0)
+            min_val = float(np.nanmin(rgb_premul)) if rgb_premul.size else 0.0
+            max_val = float(np.nanmax(rgb_premul)) if rgb_premul.size else 0.0
+            if is_normalized_float_range(min_val, max_val):
+                bg_lin = 0.214
+                comp_lin = np.clip(rgb_premul + bg_lin * (1.0 - alpha_ch), 0.0, 1.0)
+                comp_srgb = np.where(
+                    comp_lin <= 0.0031308,
+                    comp_lin * 12.92,
+                    1.055 * np.power(np.clip(comp_lin, 1e-9, 1.0), 1.0 / 2.4) - 0.055,
+                )
+                return np.clip(comp_srgb * 255.0, 0.0, 255.0).astype(np.uint8)
+            return np.clip(rgb_premul, 0.0, 255.0).astype(np.uint8)
+
+        rgb = arr[:, :, :3].astype(np.float32)
+        if arr.dtype != np.uint8:
+            min_val = float(np.nanmin(rgb)) if rgb.size else 0.0
+            max_val = float(np.nanmax(rgb)) if rgb.size else 0.0
+            if is_normalized_float_range(min_val, max_val):
+                rgb = np.clip(rgb, 0.0, 1.0)
+                rgb = np.power(rgb, 1.0 / 2.2) * 255.0
+            rgb = np.clip(rgb, 0.0, 255.0).astype(np.uint8)
+        else:
+            rgb = arr[:, :, :3]
+        return rgb
+    return None
+
+
+def frame_bbox(frame) -> tuple[int, int, int, int]:
+    arr = np.asarray(frame)
+    if arr.ndim == 2:
+        coverage = np.asarray(arr, dtype=np.float32) > 1e-6
+    elif arr.ndim == 3 and arr.shape[2] >= 4:
+        coverage = np.asarray(arr[:, :, 3], dtype=np.float32) > 1e-6
+    elif arr.ndim == 3:
+        coverage = np.any(np.asarray(arr[:, :, :3], dtype=np.float32) > 1e-6, axis=2)
+    else:
+        return (0, 0, 0, 0)
+
+    ys, xs = np.where(coverage)
+    if ys.size == 0 or xs.size == 0:
+        return (0, 0, 0, 0)
+    x0 = int(np.min(xs))
+    y0 = int(np.min(ys))
+    x1 = int(np.max(xs)) + 1
+    y1 = int(np.max(ys)) + 1
+    return (x0, y0, x1, y1)
+
+
+def bbox_union(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> tuple[int, int, int, int]:
+    if a[2] <= a[0] or a[3] <= a[1]:
+        return b
+    if b[2] <= b[0] or b[3] <= b[1]:
+        return a
+    return (min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3]))
+
+
+def bbox_intersection(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> tuple[int, int, int, int]:
+    x0 = max(a[0], b[0])
+    y0 = max(a[1], b[1])
+    x1 = min(a[2], b[2])
+    y1 = min(a[3], b[3])
+    if x1 <= x0 or y1 <= y0:
+        return (0, 0, 0, 0)
+    return (x0, y0, x1, y1)
+
+
+def clip_frame_to_bbox(frame: np.ndarray, bbox: tuple[int, int, int, int]) -> np.ndarray:
+    out = np.asarray(frame).copy()
+    h, w = out.shape[:2]
+    x0 = max(0, min(w, int(bbox[0])))
+    y0 = max(0, min(h, int(bbox[1])))
+    x1 = max(0, min(w, int(bbox[2])))
+    y1 = max(0, min(h, int(bbox[3])))
+    if x1 <= x0 or y1 <= y0:
+        return np.zeros_like(out)
+
+    mask = np.zeros((h, w), dtype=bool)
+    mask[y0:y1, x0:x1] = True
+    if out.ndim == 2:
+        out[~mask] = 0
+    else:
+        out[~mask, ...] = 0
+    return out

app/workers/inference_worker.py

@@ -42,6 +42,13 @@
     initialize_graph_write_plan,
     write_graph_plan_frame,
 )
+from app.workers.frame_helpers import (
+    bbox_intersection,
+    bbox_union,
+    clip_frame_to_bbox,
+    coerce_preview_frame,
+    frame_bbox,
+)
 from app.utils.media import (
     is_supported_image_file,
     is_numbered_image_sequence,
@@ -795,115 +802,23 @@ def _prepare_graph_write_targets(
 
     @staticmethod
     def _coerce_preview_frame(frame) -> np.ndarray | None:
-        arr = np.asarray(frame)
-        if arr.ndim == 2:
-            gray = arr.astype(np.float32)
-            min_val = float(np.nanmin(gray)) if gray.size else 0.0
-            max_val = float(np.nanmax(gray)) if gray.size else 0.0
-            if gray.dtype != np.uint8:
-                if InferenceWorker._is_normalized_float_range(min_val, max_val):
-                    gray = np.clip(gray, 0.0, 1.0)
-                    gray = np.power(gray, 1.0 / 2.2) * 255.0
-                gray = np.clip(gray, 0.0, 255.0).astype(np.uint8)
-            else:
-                gray = arr
-            return np.stack([gray] * 3, axis=-1)
-        if arr.ndim == 3 and arr.shape[2] == 1:
-            return InferenceWorker._coerce_preview_frame(arr[:, :, 0])
-        if arr.ndim == 3 and arr.shape[2] >= 3:
-            if arr.dtype != np.uint8 and arr.shape[2] >= 4:
-                # 4-channel float data = linear premultiplied RGBA (e.g. CorridorKey
-                # 'processed' output).  Stripping alpha and gamma-correcting the
-                # premultiplied RGB produces false blue/dark on transparent areas
-                # (e.g. glass: premul_b > premul_r after gamma).
-                # Correct approach: composite over a neutral mid-grey so that
-                # transparent areas show as grey, not as dark premultiplied colour.
-                rgb_premul = arr[:, :, :3].astype(np.float32)
-                alpha_ch = np.clip(arr[:, :, 3:4].astype(np.float32), 0.0, 1.0)
-                min_val = float(np.nanmin(rgb_premul)) if rgb_premul.size else 0.0
-                max_val = float(np.nanmax(rgb_premul)) if rgb_premul.size else 0.0
-                if InferenceWorker._is_normalized_float_range(min_val, max_val):
-                    # linear premul + neutral grey bg (linear 0.214 ≈ sRGB 0.5)
-                    bg_lin = 0.214
-                    comp_lin = np.clip(rgb_premul + bg_lin * (1.0 - alpha_ch), 0.0, 1.0)
-                    # linear → sRGB
-                    comp_srgb = np.where(comp_lin <= 0.0031308,
-                                         comp_lin * 12.92,
-                                         1.055 * np.power(np.clip(comp_lin, 1e-9, 1.0), 1.0 / 2.4) - 0.055)
-                    return np.clip(comp_srgb * 255.0, 0.0, 255.0).astype(np.uint8)
-                else:
-                    return np.clip(rgb_premul, 0.0, 255.0).astype(np.uint8)
-
-            rgb = arr[:, :, :3].astype(np.float32)
-            if arr.dtype != np.uint8:
-                min_val = float(np.nanmin(rgb)) if rgb.size else 0.0
-                max_val = float(np.nanmax(rgb)) if rgb.size else 0.0
-                if InferenceWorker._is_normalized_float_range(min_val, max_val):
-                    rgb = np.clip(rgb, 0.0, 1.0)
-                    rgb = np.power(rgb, 1.0 / 2.2) * 255.0
-                rgb = np.clip(rgb, 0.0, 255.0).astype(np.uint8)
-            else:
-                rgb = arr[:, :, :3]
-            return rgb
-        return None
+        return coerce_preview_frame(frame)
 
     @staticmethod
     def _frame_bbox(frame) -> tuple[int, int, int, int]:
-        arr = np.asarray(frame)
-        if arr.ndim == 2:
-            coverage = np.asarray(arr, dtype=np.float32) > 1e-6
-        elif arr.ndim == 3 and arr.shape[2] >= 4:
-            coverage = np.asarray(arr[:, :, 3], dtype=np.float32) > 1e-6
-        elif arr.ndim == 3:
-            coverage = np.any(np.asarray(arr[:, :, :3], dtype=np.float32) > 1e-6, axis=2)
-        else:
-            return (0, 0, 0, 0)
-
-        ys, xs = np.where(coverage)
-        if ys.size == 0 or xs.size == 0:
-            return (0, 0, 0, 0)
-        x0 = int(np.min(xs))
-        y0 = int(np.min(ys))
-        x1 = int(np.max(xs)) + 1
-        y1 = int(np.max(ys)) + 1
-        return (x0, y0, x1, y1)
+        return frame_bbox(frame)
 
     @staticmethod
     def _bbox_union(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> tuple[int, int, int, int]:
-        if a[2] <= a[0] or a[3] <= a[1]:
-            return b
-        if b[2] <= b[0] or b[3] <= b[1]:
-            return a
-        return (min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3]))
+        return bbox_union(a, b)
 
     @staticmethod
     def _bbox_intersection(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> tuple[int, int, int, int]:
-        x0 = max(a[0], b[0])
-        y0 = max(a[1], b[1])
-        x1 = min(a[2], b[2])
-        y1 = min(a[3], b[3])
-        if x1 <= x0 or y1 <= y0:
-            return (0, 0, 0, 0)
-        return (x0, y0, x1, y1)
+        return bbox_intersection(a, b)
 
     @staticmethod
     def _clip_frame_to_bbox(frame: np.ndarray, bbox: tuple[int, int, int, int]) -> np.ndarray:
-        out = np.asarray(frame).copy()
-        h, w = out.shape[:2]
-        x0 = max(0, min(w, int(bbox[0])))
-        y0 = max(0, min(h, int(bbox[1])))
-        x1 = max(0, min(w, int(bbox[2])))
-        y1 = max(0, min(h, int(bbox[3])))
-        if x1 <= x0 or y1 <= y0:
-            return np.zeros_like(out)
-
-        mask = np.zeros((h, w), dtype=bool)
-        mask[y0:y1, x0:x1] = True
-        if out.ndim == 2:
-            out[~mask] = 0
-        else:
-            out[~mask, ...] = 0
-        return out
+        return clip_frame_to_bbox(frame, bbox)
 
     @staticmethod
     def _to_u8_frame(src: np.ndarray) -> np.ndarray:

tests/test_inference_worker_failures.py

@@ -18,6 +18,31 @@
 
 
 class InferenceWorkerFailurePathTests(unittest.TestCase):
+    def test_preview_frame_wrapper_handles_linear_premultiplied_rgba(self):
+        frame = np.zeros((1, 2, 4), dtype=np.float32)
+        frame[0, 0] = [0.0, 0.0, 0.0, 0.0]
+        frame[0, 1] = [0.25, 0.25, 0.25, 1.0]
+
+        preview = InferenceWorker._coerce_preview_frame(frame)
+
+        self.assertIsNotNone(preview)
+        self.assertEqual(preview.shape, (1, 2, 3))
+        self.assertEqual(preview.dtype, np.uint8)
+        self.assertGreater(int(preview[0, 0, 0]), 100)
+
+    def test_bbox_wrappers_preserve_geometry_helpers(self):
+        frame = np.zeros((4, 5, 4), dtype=np.float32)
+        frame[1:3, 2:5, 3] = 1.0
+
+        bbox = InferenceWorker._frame_bbox(frame)
+        clipped = InferenceWorker._clip_frame_to_bbox(frame, (2, 1, 4, 3))
+
+        self.assertEqual(bbox, (2, 1, 5, 3))
+        self.assertEqual(InferenceWorker._bbox_intersection(bbox, (3, 0, 5, 2)), (3, 1, 5, 2))
+        self.assertEqual(InferenceWorker._bbox_union((0, 0, 0, 0), bbox), bbox)
+        self.assertEqual(float(clipped[1, 2, 3]), 1.0)
+        self.assertEqual(float(clipped[1, 4, 3]), 0.0)
+
     def test_birefnet_runtime_notice_is_exposed_for_mps_fallback(self):
         service = BiRefNetService()
         service.device = "mps"