SmartTools/SmartLoad.py at main · slikvik55/SmartTools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#
# SmartLoad.py
#
# Load Image equivalent with optional crop rectangle (interactive UI in smart_load.js).
#

import torch

from nodes import LoadImage


class SmartLoad(LoadImage):
    """Same as Load Image; outputs are cropped to crop_x/y/w/h (image pixel space)."""

    @classmethod
    def INPUT_TYPES(cls):
        parent = LoadImage.INPUT_TYPES()
        req = dict(parent["required"])
        req["crop_x"] = ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1})
        req["crop_y"] = ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1})
        req["crop_w"] = ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1})
        req["crop_h"] = ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1})
        out = dict(parent)
        out["required"] = req
        return out

    CATEGORY = "slikvik/Image"
    DISPLAY_NAME = "Smart Load"
    DESCRIPTION = (
        "Loads an image like Load Image. Crop writes a cropped copy as the node's image "
        "(preview and workflow restore); Clear Crop brings back the stored original."
    )

    def load_image(self, image, crop_x, crop_y, crop_w, crop_h):
        img, mask = super().load_image(image)
        _, H0, W0, _ = img.shape

        if crop_w <= 0 or crop_h <= 0:
            return (img, mask)

        x, y, w, h = self._clamp_crop(crop_x, crop_y, crop_w, crop_h, W0, H0)

        img = img[:, y : y + h, x : x + w, :]

        if mask.shape[1] == H0 and mask.shape[2] == W0:
            mask = mask[:, y : y + h, x : x + w]
        else:
            b = mask.shape[0]
            mask = torch.zeros(
                (b, h, w),
                dtype=mask.dtype,
                device=mask.device,
            )

        return (img, mask)

    @staticmethod
    def _clamp_crop(x, y, w, h, W, H):
        """Match CropImage-style bounds: minimum 1 px, stay inside image."""
        min_size = 1
        xi = max(0, min(int(x), W - min_size))
        yi = max(0, min(int(y), H - min_size))
        wi = max(min_size, min(int(w), W - xi))
        hi = max(min_size, min(int(h), H - yi))
        return xi, yi, wi, hi

    @classmethod
    def IS_CHANGED(cls, image, crop_x, crop_y, crop_w, crop_h):
        return (super().IS_CHANGED(image), crop_x, crop_y, crop_w, crop_h)


NODE_CLASS_MAPPINGS = {
    "SmartLoad": SmartLoad,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SmartLoad": "Smart Load",
}