From f8d7bb71da1d587b15839e2c40e632bef65caf3c Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Wed, 13 Sep 2023 16:58:51 +0900
Subject: [PATCH] feat: ImpactSEGSToMaskBatch, ImpactScaleBy_BBOX_SEG_ELT fix:
 MaskListToMaskBatch, type fix MASK -> MASKS

---
 __init__.py                   |  4 ++
 modules/impact/config.py      |  2 +-
 modules/impact/impact_pack.py |  2 +-
 modules/impact/segs_nodes.py  | 70 ++++++++++++++++++++++++++++++++++-
 4 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/__init__.py b/__init__.py
index cf80046..10d858d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -177,6 +177,7 @@ NODE_CLASS_MAPPINGS = {
     "ImpactEdit_SEG_ELT": Edit_SEG_ELT,
     "ImpactDilate_Mask_SEG_ELT": Dilate_SEG_ELT,
     "ImpactDilateMask": DilateMask,
+    "ImpactScaleBy_BBOX_SEG_ELT": SEG_ELT_BBOX_ScaleBy,
 
     "BboxDetectorCombined_v2": BboxDetectorCombined,
     "SegmDetectorCombined_v2": SegmDetectorCombined,
@@ -211,6 +212,7 @@ NODE_CLASS_MAPPINGS = {
     "SEGSPreview": SEGSPreview,
     "SEGSToImageList": SEGSToImageList,
     "ImpactSEGSToMaskList": SEGSToMaskList,
+    "ImpactSEGSToMaskBatch": SEGSToMaskBatch,
     "ImpactSEGSConcat": SEGSConcat,
 
     "ImpactKSamplerBasicPipe": KSamplerBasicPipe,
@@ -303,12 +305,14 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImpactSEGSOrderedFilter": "SEGS Filter (ordered)",
     "ImpactSEGSConcat": "SEGS Concat",
     "ImpactSEGSToMaskList": "SEGS to Mask List",
+    "ImpactSEGSToMaskBatch": "SEGS to Mask Batch",
 
     "ImpactDecomposeSEGS": "Decompose (SEGS)",
     "ImpactAssembleSEGS": "Assemble (SEGS)",
     "ImpactFrom_SEG_ELT": "From SEG_ELT",
     "ImpactEdit_SEG_ELT": "Edit SEG_ELT",
     "ImpactDilate_Mask_SEG_ELT": "Dilate Mask (SEG_ELT)",
+    "ImpactScaleBy_BBOX_SEG_ELT": "ScaleBy BBOX (SEG_ELT)",
     "ImpactDilateMask": "Dilate Mask",
 
     "PreviewBridge": "Preview Bridge",
diff --git a/modules/impact/config.py b/modules/impact/config.py
index 017c5d4..b271d4b 100644
--- a/modules/impact/config.py
+++ b/modules/impact/config.py
@@ -2,7 +2,7 @@ import configparser
 import os
 
 
-version = "V4.5"
+version = "V4.6"
 
 dependency_version = 11
 
diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py
index 30059bd..785152e 100644
--- a/modules/impact/impact_pack.py
+++ b/modules/impact/impact_pack.py
@@ -1524,7 +1524,7 @@ class MaskListToMaskBatch:
 
     INPUT_IS_LIST = True
 
-    RETURN_TYPES = ("MASK", )
+    RETURN_TYPES = ("MASKS", )
     FUNCTION = "doit"
 
     CATEGORY = "ImpactPack/Operation"
diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py
index c791228..7df37cc 100644
--- a/modules/impact/segs_nodes.py
+++ b/modules/impact/segs_nodes.py
@@ -416,6 +416,26 @@ class SEGSToMaskList:
         return (masks,)
 
 
+class SEGSToMaskBatch:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                     "segs": ("SEGS", ),
+                     },
+                }
+
+    RETURN_TYPES = ("MASKS",)
+    OUTPUT_IS_LIST = (True,)
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def doit(self, segs):
+        masks = core.segs_to_masklist(segs)
+        mask_batch = torch.stack(masks, dim=0)
+        return (mask_batch,)
+
+
 class SEGSConcat:
     @classmethod
     def INPUT_TYPES(s):
@@ -574,7 +594,55 @@ class Dilate_SEG_ELT:
     def doit(self, seg, dilation):
         mask = core.dilate_mask(seg.cropped_mask, dilation)
         seg = SEG(seg.cropped_image, mask, seg.confidence, seg.crop_region, seg.bbox, seg.label, seg.control_net_wrapper)
-        return seg
+        return (seg,)
+
+
+class SEG_ELT_BBOX_ScaleBy:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                     "seg": ("SEG_ELT", ),
+                     "scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01}), }
+                }
+
+    RETURN_TYPES = ("SEG_ELT", )
+
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    @staticmethod
+    def fill_zero_outside_bbox(mask, crop_region, bbox):
+        cx1, cy1, _, _ = crop_region
+        x1, y1, x2, y2 = bbox
+        x1, y1, x2, y2 = x1-cx1, y1-cy1, x2-cx1, y2-cy1
+        h, w = mask.shape
+
+        x1 = min(w-1, max(0, x1))
+        x2 = min(w-1, max(0, x2))
+        y1 = min(h-1, max(0, y1))
+        y2 = min(h-1, max(0, y2))
+
+        mask_cropped = mask.copy()
+        mask_cropped[:, :x1] = 0  # zero fill left side
+        mask_cropped[:, x2:] = 0  # zero fill right side
+        mask_cropped[:y1, :] = 0  # zero fill top side
+        mask_cropped[y2:, :] = 0  # zero fill bottom side
+        return mask_cropped
+
+    def doit(self, seg, scale_by):
+        x1, y1, x2, y2 = seg.bbox
+        w = x2-x1
+        h = y2-y1
+
+        dw = int((w * scale_by - w)/2)
+        dh = int((h * scale_by - h)/2)
+
+        bbox = (x1-dw, y1-dh, x2+dw, y2+dh)
+
+        cropped_mask = SEG_ELT_BBOX_ScaleBy.fill_zero_outside_bbox(seg.cropped_mask, seg.crop_region, bbox)
+        seg = SEG(seg.cropped_image, cropped_mask, seg.confidence, seg.crop_region, bbox, seg.label, seg.control_net_wrapper)
+        return (seg,)
 
 
 class EmptySEGS: