From 33c68fdd8e9bec8afdc67012a9bb7fc5d7c013a3 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Tue, 12 Dec 2023 18:39:47 +0900
Subject: [PATCH] feat: Set Default Image for SEGS

---
 README.md                    |  1 +
 __init__.py                  |  2 ++
 modules/impact/config.py     |  2 +-
 modules/impact/segs_nodes.py | 48 ++++++++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d77050a..a2337d3 100644
--- a/README.md
+++ b/README.md
@@ -86,6 +86,7 @@ This custom node helps to conveniently enhance images through Detector, Detailer
   * SEGS Filter (range) - This node retrieves only SEGs from SEGS that have a size and position within a certain range.
   * SEGSConcat - Concatenate segs1 and segs2. If source shape of segs1 and segs2 are different from segs2 will be ignored.
   * Picker (SEGS) - Among the input SEGS, you can select a specific SEG through a dialog. If no SEG is selected, it outputs an empty SEGS. Increasing the batch_size of SEGSDetailer can be used for the purpose of selecting from the candidates.
+  * Set Default Image For SEGS - Set a default image for SEGS. SEGS with images set this way do not need to have a fallback image set. When override is set to false, the original image is preserved. 
   * DecomposeSEGS - Decompose SEGS to allow for detailed manipulation.
   * AssembleSEGS - Reassemble the decomposed SEGS.
   * From SEG_ELT - Extract detailed information from SEG_ELT.
diff --git a/__init__.py b/__init__.py
index 8673175..0a0c940 100644
--- a/__init__.py
+++ b/__init__.py
@@ -195,6 +195,7 @@ NODE_CLASS_MAPPINGS = {
     "MasksToMaskList": MasksToMaskList,
     "MaskListToMaskBatch": MaskListToMaskBatch,
     "ImageListToImageBatch": ImageListToMaskBatch,
+    "SetDefaultImageForSEGS": DefaultImageForSEGS,
 
     "BboxDetectorSEGS": BboxDetectorForEach,
     "SegmDetectorSEGS": SegmDetectorForEach,
@@ -387,6 +388,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImpactMakeImageBatch": "Make Image Batch",
     "ImpactStringSelector": "String Selector",
     "ImpactIsNotEmptySEGS": "SEGS isn't Empty",
+    "SetDefaultImageForSEGS": "Set Default Image for SEGS",
 
     "RemoveNoiseMask": "Remove Noise Mask",
 
diff --git a/modules/impact/config.py b/modules/impact/config.py
index d530c5c..288fb9e 100644
--- a/modules/impact/config.py
+++ b/modules/impact/config.py
@@ -2,7 +2,7 @@ import configparser
 import os
 
 
-version = "V4.45.2"
+version = "V4.46"
 
 dependency_version = 19
 
diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py
index b3520cf..0260958 100644
--- a/modules/impact/segs_nodes.py
+++ b/modules/impact/segs_nodes.py
@@ -1096,3 +1096,51 @@ class SEGSPicker:
                 new_segs.append(segs[1][i])
 
         return ((segs[0], new_segs),)
+
+
+class DefaultImageForSEGS:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "segs": ("SEGS", ),
+                    "image": ("IMAGE", ),
+                    "override": ("BOOLEAN", {"default": True}),
+                }}
+
+    RETURN_TYPES = ("SEGS", )
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def doit(self, segs, image, override):
+        results = []
+
+        segs = core.segs_scale_match(segs, image.shape)
+
+        if len(segs[1]) > 0:
+            if segs[1][0].cropped_image is not None:
+                batch_count = len(segs[1][0].cropped_image)
+            else:
+                batch_count = len(image)
+
+            for seg in segs[1]:
+                if seg.cropped_image is not None and not override:
+                    cropped_image = seg.cropped_image
+                else:
+                    cropped_image = None
+                    for i in range(0, batch_count):
+                        # take from original image
+                        ref_image = image[i].unsqueeze(0)
+                        cropped_image2 = crop_image(ref_image, seg.crop_region)
+
+                        if cropped_image is None:
+                            cropped_image = cropped_image2
+                        else:
+                            torch.cat((cropped_image, cropped_image2), dim=0)
+
+                new_seg = SEG(cropped_image, seg.cropped_mask, seg.confidence, seg.crop_region, seg.bbox, seg.label, seg.control_net_wrapper)
+                results.append(new_seg)
+
+            return ((segs[0], results), )
+        else:
+            return (segs, )
\ No newline at end of file