Upload Get_Correct_Batch_Img.py

Get_Correct_Batch_Img.py

@@ -0,0 +1,171 @@
+import torch
+
+
+class Get_Correct_Batch_Img:
+    """
+    Given a batch of RGBA images, selects:
+      - the sprite with the widest visible span along a given Y row (max_img)
+      - the sprite with the thinnest visible span along that same row (min_img)
+      - the sprite whose width is closest to the midpoint between min/max widths (avg_img)
+
+    Visibility is determined from the alpha channel (A > 0).
+    Only images within [start_index, end_index] (inclusive) are considered.
+    """
+
+    # Where this node appears in the right-click menu:
+    CATEGORY = "image/batch"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                # RGBA image batch: torch.Tensor [B, H, W, 4]
+                "images": ("IMAGE",),
+
+                # Sub-batch start index (inclusive, 0-based)
+                "start_index": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 2_147_483_647,
+                        "step": 1,
+                    },
+                ),
+
+                # Sub-batch end index (inclusive, 0-based)
+                "end_index": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 2_147_483_647,
+                        "step": 1,
+                    },
+                ),
+
+                # Y coordinate (row) used for the horizontal scan
+                "y_coord": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 2_147_483_647,
+                        "step": 1,
+                    },
+                ),
+            }
+        }
+
+    # Three RGBA images out now
+    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
+    RETURN_NAMES = ("max_img", "min_img", "avg_img")
+    FUNCTION = "select"
+
+    def select(self, images, start_index, end_index, y_coord):
+        # Basic sanity checks
+        if not isinstance(images, torch.Tensor):
+            raise TypeError(f"Expected IMAGE tensor, got {type(images)}")
+
+        if images.ndim != 4:
+            raise ValueError(
+                f"Expected IMAGE of shape [B,H,W,C], got {tuple(images.shape)}"
+            )
+
+        batch_size, height, width, channels = images.shape
+
+        if channels != 4:
+            raise ValueError(
+                f"Expected RGBA image with 4 channels, got {channels}. "
+                "Make sure your input batch is RGBA (not RGB)."
+            )
+
+        if batch_size == 0:
+            raise ValueError("Empty image batch passed to Get_Correct_Batch_Img.")
+
+        # Clamp and normalize indices
+        start = max(0, min(int(start_index), batch_size - 1))
+        end = max(0, min(int(end_index), batch_size - 1))
+        if start > end:
+            start, end = end, start  # swap so start <= end
+
+        # Clamp Y coordinate into image bounds
+        y = max(0, min(int(y_coord), height - 1))
+
+        # Track widest and thinnest sprite
+        max_width = None
+        min_width = None
+        max_idx = start
+        min_idx = start
+
+        # For AVG: store (index, width_px) for all valid sprites
+        widths = []
+
+        # Small alpha threshold; alpha > 0 is "visible"
+        alpha_threshold = 0.0
+        any_visible = False
+
+        # Loop over the requested sub-batch only
+        for i in range(start, end + 1):
+            # row_alpha shape: [W]
+            row_alpha = images[i, y, :, 3]
+            visible = row_alpha > alpha_threshold
+
+            if not torch.any(visible):
+                # No visible pixels on this row for this image; skip it
+                continue
+
+            any_visible = True
+
+            # Indices of visible pixels along X
+            visible_indices = torch.nonzero(visible, as_tuple=False).squeeze(1)
+            left_x = int(visible_indices[0])
+            right_x = int(visible_indices[-1])
+            width_px = right_x - left_x + 1  # inclusive distance
+
+            widths.append((i, width_px))
+
+            # Update max width (widest sprite)
+            if max_width is None or width_px > max_width:
+                max_width = width_px
+                max_idx = i
+
+            # Update min width (thinnest sprite)
+            if min_width is None or width_px < min_width:
+                min_width = width_px
+                min_idx = i
+
+        # If nothing had visible pixels on that Y, just return the first image
+        # in the sub-batch as all three outputs (so the node never crashes).
+        if not any_visible:
+            base_img = images[start].unsqueeze(0)
+            return (base_img, base_img, base_img)
+
+        # Compute midpoint between MIN and MAX widths
+        center_width = (min_width + max_width) / 2.0
+
+        # Find sprite whose width is closest to this center_width
+        avg_idx = max_idx  # default
+        closest_diff = None
+        for idx, w in widths:
+            diff = abs(w - center_width)
+            if closest_diff is None or diff < closest_diff:
+                closest_diff = diff
+                avg_idx = idx
+
+        # Extract chosen sprites as batch size 1 (B=1, H, W, C)
+        max_img = images[max_idx].unsqueeze(0)
+        min_img = images[min_idx].unsqueeze(0)
+        avg_img = images[avg_idx].unsqueeze(0)
+
+        return (max_img, min_img, avg_img)
+
+
+# Register node with ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "Get_Correct_Batch_Img": Get_Correct_Batch_Img,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Get_Correct_Batch_Img": "Get_Correct_Batch_Img (Salia)",
+}