From cca73af8c12baba352970e936c462a49b785a090 Mon Sep 17 00:00:00 2001 From: TTPlanetPig <152850462+TTPlanetPig@users.noreply.github.com> Date: Thu, 12 Jun 2025 15:41:39 +0800 Subject: [PATCH 1/4] Add padding option to smart tiler --- README.md | 17 +++++ TTP_smart_tile.py | 111 +++++++++++++++++++++++++++++++ __init__.py | 19 +++++- examples/smart_tile_example.json | 32 +++++++++ pyproject.toml | 12 +++- 5 files changed, 187 insertions(+), 4 deletions(-) create mode 100644 TTP_smart_tile.py create mode 100644 examples/smart_tile_example.json diff --git a/README.md b/README.md index 99205ea..cc5b108 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,20 @@ This node merges all tiled conditions into one and prepares them for building th ![Condition Merge Node](https://github.com/user-attachments/assets/3039c8a3-8284-4b71-a9de-4120723258c7) +--- +### **7. Smart Tile Nodes** +These nodes allow object-aware tiling using a YOLO detection or segmentation model. Use **TTP Smart Tile Batch** to crop detected regions and **TTP Smart Image Assy** to paste the edited tiles back. + +`TTP Smart Tile Batch` now accepts an optional *padding* value to expand bounding boxes before cropping. + +The nodes depend on the `ultralytics` package: + +```bash +pip install ultralytics +``` + +See `examples/smart_tile_example.json` for a minimal workflow using these nodes. + --- ## **Examples** @@ -116,6 +130,9 @@ This node merges all tiled conditions into one and prepares them for building th ### **Latent Example** ![Latent Example Workflow](https://github.com/TTPlanetPig/Comfyui_TTP_Toolset/blob/main/examples/Flux_8Mega_Pixel_image_upscale_process.png) +### **Smart Tile Example** +`examples/smart_tile_example.json` demonstrates object-aware tiling and recombination. + --- diff --git a/TTP_smart_tile.py b/TTP_smart_tile.py new file mode 100644 index 0000000..f44103c --- /dev/null +++ b/TTP_smart_tile.py @@ -0,0 +1,111 @@ +import torch +import numpy as np +from PIL import Image +from typing import List, Tuple + +from .TTP_toolsets import pil2tensor, tensor2pil + +try: + from ultralytics import YOLO +except Exception: # pragma: no cover - ultralytics optional + YOLO = None + + +class TTP_Smart_Tile_Batch: + """Split image into smart tiles using detection or segmentation.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_path": ("STRING", {"default": "yolov8n-seg.pt"}), + }, + "optional": { + "task": (["detect", "segment"], {"default": "detect"}), + "conf": ("FLOAT", {"default": 0.25, "min": 0.01, "max": 1.0}), + "iou": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0}), + "padding": ("INT", {"default": 0, "min": 0, "max": 256}), + }, + } + + RETURN_TYPES = ("IMAGE", "LIST", "LIST") + RETURN_NAMES = ("tiles", "infos", "ids") + FUNCTION = "tile_image" + CATEGORY = "TTP/Image" + + def tile_image(self, image, model_path="yolov8n-seg.pt", task="detect", conf=0.25, iou=0.45, padding=0): + if YOLO is None: + raise RuntimeError("ultralytics package is required for smart tiling") + + pil_img = tensor2pil(image.squeeze(0)) + model = YOLO(model_path) + results = model.predict(pil_img, conf=conf, iou=iou, verbose=False) + result = results[0] + + tiles: List[torch.Tensor] = [] + infos: List = [] + ids: List[int] = [] + + boxes = result.boxes.xyxy.cpu().numpy().astype(int) + classes = result.boxes.cls.cpu().numpy().astype(int) + masks = None + if task == "segment" and result.masks is not None: + masks = result.masks.data.cpu().numpy() + + width, height = pil_img.size + for idx, box in enumerate(boxes): + x1, y1, x2, y2 = box.tolist() + x1 = max(0, x1 - padding) + y1 = max(0, y1 - padding) + x2 = min(width, x2 + padding) + y2 = min(height, y2 + padding) + tile = pil_img.crop((x1, y1, x2, y2)) + tiles.append(pil2tensor(tile)) + if masks is not None: + mask = masks[idx][y1:y2, x1:x2] + infos.append(torch.from_numpy(mask).unsqueeze(0)) + else: + infos.append((x1, y1, x2, y2)) + ids.append(int(classes[idx])) + + tiles_tensor = torch.stack(tiles, dim=0).squeeze(1) if tiles else torch.empty((0, 3, 0, 0)) + return tiles_tensor, infos, ids + + +class TTP_Smart_Image_Assy: + """Reassemble edited smart tiles back to the original image.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "base_image": ("IMAGE",), + "tiles": ("IMAGE", {"forceInput": True}), + "infos": ("LIST",), + "task": (["detect", "segment"], {"default": "detect"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "assemble_image" + CATEGORY = "TTP/Image" + + def assemble_image(self, base_image, tiles, infos, task="detect"): + canvas = tensor2pil(base_image.squeeze(0)).copy() + for idx, tile in enumerate(tiles): + tile_img = tensor2pil(tile.unsqueeze(0)) + info = infos[idx] + if task == "segment" and not isinstance(info, tuple): + mask = tensor2pil(info) + bbox = mask.getbbox() + if bbox: + x1, y1, x2, y2 = bbox + mask_resized = mask.crop(bbox) + tile_crop = tile_img.crop((0, 0, x2 - x1, y2 - y1)) + canvas.paste(tile_crop, (x1, y1, x2, y2), mask_resized) + else: + x1, y1, x2, y2 = info + canvas.paste(tile_img, (x1, y1, x2, y2)) + return pil2tensor(canvas) diff --git a/__init__.py b/__init__.py index 12efac0..5af0929 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,16 @@ -from .TTP_toolsets import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS - -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file +from .TTP_toolsets import NODE_CLASS_MAPPINGS as TOOLSET_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as TOOLSET_DISPLAY +from .TTP_smart_tile import TTP_Smart_Tile_Batch, TTP_Smart_Image_Assy + +NODE_CLASS_MAPPINGS = { + **TOOLSET_MAPPINGS, + "TTP_Smart_Tile_Batch": TTP_Smart_Tile_Batch, + "TTP_Smart_Image_Assy": TTP_Smart_Image_Assy, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + **TOOLSET_DISPLAY, + "TTP_Smart_Tile_Batch": "TTP Smart Tile Batch", + "TTP_Smart_Image_Assy": "TTP Smart Image Assy", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/examples/smart_tile_example.json b/examples/smart_tile_example.json new file mode 100644 index 0000000..2a03839 --- /dev/null +++ b/examples/smart_tile_example.json @@ -0,0 +1,32 @@ +{ + "last_node_id": 4, + "last_link_id": 4, + "nodes": [ + { + "id": 1, + "type": "ImageInput", + "pos": [0, 0], + "size": [210, 70], + "widgets_values": ["input.png"] + }, + { + "id": 2, + "type": "TTP_Smart_Tile_Batch", + "pos": [250, 0], + "size": [250, 120], + "widgets_values": ["yolov8n-seg.pt", "detect", 0.25, 0.45, 10] + }, + { + "id": 3, + "type": "TTP_Smart_Image_Assy", + "pos": [500, 0], + "size": [250, 120], + "widgets_values": ["detect"] + } + ], + "links": [ + [1, 0, 2, 0], + [2, 0, 3, 1], + [2, 1, 3, 2] + ] +} diff --git a/pyproject.toml b/pyproject.toml index 64d3289..5f48d16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,8 +1,9 @@ [project] name = "comfyui_ttp_toolset" description = "This is a workflow for my simple logic amazing upscale node for DIT model. it can be common use for Flux,Hunyuan,SD3 It can simple tile the initial image into pieces and then use image-interrogator to get each tile prompts for more accurate upscale process. The condition will be properly handled and the hallucination will be significantly eliminated." -version = "1.0.5" +version = "1.0.6" license = {file = "LICENSE"} +dependencies = ["ultralytics>=8.1"] [project.urls] Repository = "https://github.com/TTPlanetPig/Comfyui_TTP_Toolset" @@ -12,3 +13,12 @@ Repository = "https://github.com/TTPlanetPig/Comfyui_TTP_Toolset" PublisherId = "ttplanet" DisplayName = "Comfyui_TTP_Toolset" Icon = "🪐" + +[build-system] +requires = ["setuptools>=66", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["Comfyui_TTP_Toolset"] +[tool.setuptools.package-dir] +Comfyui_TTP_Toolset = "." From dc2f43cb87aa56015e47a2f27120e805f14d5672 Mon Sep 17 00:00:00 2001 From: TTPlanetPig <152850462+TTPlanetPig@users.noreply.github.com> Date: Thu, 12 Jun 2025 16:31:51 +0800 Subject: [PATCH 2/4] fix smart tile padding --- TTP_smart_tile.py | 40 ++++++++++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/TTP_smart_tile.py b/TTP_smart_tile.py index f44103c..8d07b67 100644 --- a/TTP_smart_tile.py +++ b/TTP_smart_tile.py @@ -46,6 +46,8 @@ def tile_image(self, image, model_path="yolov8n-seg.pt", task="detect", conf=0.2 tiles: List[torch.Tensor] = [] infos: List = [] ids: List[int] = [] + max_w = 0 + max_h = 0 boxes = result.boxes.xyxy.cpu().numpy().astype(int) classes = result.boxes.cls.cpu().numpy().astype(int) @@ -61,15 +63,29 @@ def tile_image(self, image, model_path="yolov8n-seg.pt", task="detect", conf=0.2 x2 = min(width, x2 + padding) y2 = min(height, y2 + padding) tile = pil_img.crop((x1, y1, x2, y2)) - tiles.append(pil2tensor(tile)) + tile_tensor = pil2tensor(tile) + _, h, w, _ = tile_tensor.shape + max_h = max(max_h, h) + max_w = max(max_w, w) + tiles.append(tile_tensor) if masks is not None: mask = masks[idx][y1:y2, x1:x2] - infos.append(torch.from_numpy(mask).unsqueeze(0)) + infos.append((x1, y1, x2, y2, torch.from_numpy(mask).unsqueeze(0))) else: infos.append((x1, y1, x2, y2)) ids.append(int(classes[idx])) - tiles_tensor = torch.stack(tiles, dim=0).squeeze(1) if tiles else torch.empty((0, 3, 0, 0)) + padded_tiles = [] + for tile in tiles: + _, h, w, _ = tile.shape + pad_w = max_w - w + pad_h = max_h - h + if pad_w or pad_h: + tile_chw = tile.permute(0, 3, 1, 2) + tile_chw = torch.nn.functional.pad(tile_chw, (0, pad_w, 0, pad_h)) + tile = tile_chw.permute(0, 2, 3, 1) + padded_tiles.append(tile) + tiles_tensor = torch.cat(padded_tiles, dim=0) if padded_tiles else torch.empty((0, max_h, max_w, 3)) return tiles_tensor, infos, ids @@ -97,15 +113,15 @@ def assemble_image(self, base_image, tiles, infos, task="detect"): for idx, tile in enumerate(tiles): tile_img = tensor2pil(tile.unsqueeze(0)) info = infos[idx] - if task == "segment" and not isinstance(info, tuple): - mask = tensor2pil(info) - bbox = mask.getbbox() - if bbox: - x1, y1, x2, y2 = bbox - mask_resized = mask.crop(bbox) - tile_crop = tile_img.crop((0, 0, x2 - x1, y2 - y1)) - canvas.paste(tile_crop, (x1, y1, x2, y2), mask_resized) + if task == "segment": + x1, y1, x2, y2, mask_tensor = info + width, height = x2 - x1, y2 - y1 + tile_crop = tile_img.crop((0, 0, width, height)) + mask = tensor2pil(mask_tensor) + canvas.paste(tile_crop, (x1, y1, x2, y2), mask) else: x1, y1, x2, y2 = info - canvas.paste(tile_img, (x1, y1, x2, y2)) + width, height = x2 - x1, y2 - y1 + tile_crop = tile_img.crop((0, 0, width, height)) + canvas.paste(tile_crop, (x1, y1, x2, y2)) return pil2tensor(canvas) From 7eab116be490709e87f36e189ba122a1f4981b8a Mon Sep 17 00:00:00 2001 From: TTPlanetPig <152850462+TTPlanetPig@users.noreply.github.com> Date: Thu, 12 Jun 2025 17:01:55 +0800 Subject: [PATCH 3/4] Add manual tile nodes --- README.md | 14 ++--- TTP_smart_tile.py | 98 +++++++++----------------------- examples/smart_tile_example.json | 4 +- pyproject.toml | 4 +- 4 files changed, 36 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index cc5b108..21c1111 100644 --- a/README.md +++ b/README.md @@ -106,16 +106,10 @@ This node merges all tiled conditions into one and prepares them for building th ![Condition Merge Node](https://github.com/user-attachments/assets/3039c8a3-8284-4b71-a9de-4120723258c7) --- -### **7. Smart Tile Nodes** -These nodes allow object-aware tiling using a YOLO detection or segmentation model. Use **TTP Smart Tile Batch** to crop detected regions and **TTP Smart Image Assy** to paste the edited tiles back. +### **7. Manual Tile Nodes** +Use **TTP Smart Tile Batch** to crop arbitrary regions defined by a list of bounding boxes. After processing the tiles with any other ComfyUI nodes, **TTP Smart Image Assy** pastes them back in place. -`TTP Smart Tile Batch` now accepts an optional *padding* value to expand bounding boxes before cropping. - -The nodes depend on the `ultralytics` package: - -```bash -pip install ultralytics -``` +Bounding boxes should be provided as `[x1, y1, x2, y2]` coordinates. See `examples/smart_tile_example.json` for a minimal workflow using these nodes. @@ -131,7 +125,7 @@ See `examples/smart_tile_example.json` for a minimal workflow using these nodes. ![Latent Example Workflow](https://github.com/TTPlanetPig/Comfyui_TTP_Toolset/blob/main/examples/Flux_8Mega_Pixel_image_upscale_process.png) ### **Smart Tile Example** -`examples/smart_tile_example.json` demonstrates object-aware tiling and recombination. +`examples/smart_tile_example.json` demonstrates manual tiling and recombination. --- diff --git a/TTP_smart_tile.py b/TTP_smart_tile.py index 8d07b67..8b04ac6 100644 --- a/TTP_smart_tile.py +++ b/TTP_smart_tile.py @@ -1,81 +1,47 @@ import torch -import numpy as np from PIL import Image from typing import List, Tuple from .TTP_toolsets import pil2tensor, tensor2pil -try: - from ultralytics import YOLO -except Exception: # pragma: no cover - ultralytics optional - YOLO = None - class TTP_Smart_Tile_Batch: - """Split image into smart tiles using detection or segmentation.""" + """Manually crop image regions based on provided bounding boxes.""" @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "model_path": ("STRING", {"default": "yolov8n-seg.pt"}), - }, - "optional": { - "task": (["detect", "segment"], {"default": "detect"}), - "conf": ("FLOAT", {"default": 0.25, "min": 0.01, "max": 1.0}), - "iou": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0}), - "padding": ("INT", {"default": 0, "min": 0, "max": 256}), - }, + "boxes": ("LIST",), + } } - RETURN_TYPES = ("IMAGE", "LIST", "LIST") - RETURN_NAMES = ("tiles", "infos", "ids") + RETURN_TYPES = ("IMAGE", "LIST") + RETURN_NAMES = ("tiles", "positions") FUNCTION = "tile_image" CATEGORY = "TTP/Image" - def tile_image(self, image, model_path="yolov8n-seg.pt", task="detect", conf=0.25, iou=0.45, padding=0): - if YOLO is None: - raise RuntimeError("ultralytics package is required for smart tiling") - + def tile_image(self, image, boxes: List[Tuple[int, int, int, int]]): pil_img = tensor2pil(image.squeeze(0)) - model = YOLO(model_path) - results = model.predict(pil_img, conf=conf, iou=iou, verbose=False) - result = results[0] - tiles: List[torch.Tensor] = [] - infos: List = [] - ids: List[int] = [] + tiles = [] + positions = [] max_w = 0 max_h = 0 - - boxes = result.boxes.xyxy.cpu().numpy().astype(int) - classes = result.boxes.cls.cpu().numpy().astype(int) - masks = None - if task == "segment" and result.masks is not None: - masks = result.masks.data.cpu().numpy() - - width, height = pil_img.size - for idx, box in enumerate(boxes): - x1, y1, x2, y2 = box.tolist() - x1 = max(0, x1 - padding) - y1 = max(0, y1 - padding) - x2 = min(width, x2 + padding) - y2 = min(height, y2 + padding) - tile = pil_img.crop((x1, y1, x2, y2)) - tile_tensor = pil2tensor(tile) + for box in boxes: + if len(box) != 4: + raise ValueError(f"Each box must contain 4 values, got {box}") + x1, y1, x2, y2 = map(int, box) + crop = pil_img.crop((x1, y1, x2, y2)) + tile_tensor = pil2tensor(crop) _, h, w, _ = tile_tensor.shape max_h = max(max_h, h) max_w = max(max_w, w) tiles.append(tile_tensor) - if masks is not None: - mask = masks[idx][y1:y2, x1:x2] - infos.append((x1, y1, x2, y2, torch.from_numpy(mask).unsqueeze(0))) - else: - infos.append((x1, y1, x2, y2)) - ids.append(int(classes[idx])) + positions.append((x1, y1, x2, y2)) - padded_tiles = [] + padded = [] for tile in tiles: _, h, w, _ = tile.shape pad_w = max_w - w @@ -84,13 +50,14 @@ def tile_image(self, image, model_path="yolov8n-seg.pt", task="detect", conf=0.2 tile_chw = tile.permute(0, 3, 1, 2) tile_chw = torch.nn.functional.pad(tile_chw, (0, pad_w, 0, pad_h)) tile = tile_chw.permute(0, 2, 3, 1) - padded_tiles.append(tile) - tiles_tensor = torch.cat(padded_tiles, dim=0) if padded_tiles else torch.empty((0, max_h, max_w, 3)) - return tiles_tensor, infos, ids + padded.append(tile) + + tiles_tensor = torch.cat(padded, dim=0) if padded else torch.empty((0, max_h, max_w, 3)) + return tiles_tensor, positions class TTP_Smart_Image_Assy: - """Reassemble edited smart tiles back to the original image.""" + """Reassemble manually cropped tiles back onto the base image.""" @classmethod def INPUT_TYPES(cls): @@ -98,8 +65,7 @@ def INPUT_TYPES(cls): "required": { "base_image": ("IMAGE",), "tiles": ("IMAGE", {"forceInput": True}), - "infos": ("LIST",), - "task": (["detect", "segment"], {"default": "detect"}), + "positions": ("LIST",), } } @@ -108,20 +74,12 @@ def INPUT_TYPES(cls): FUNCTION = "assemble_image" CATEGORY = "TTP/Image" - def assemble_image(self, base_image, tiles, infos, task="detect"): + def assemble_image(self, base_image, tiles, positions): canvas = tensor2pil(base_image.squeeze(0)).copy() for idx, tile in enumerate(tiles): - tile_img = tensor2pil(tile.unsqueeze(0)) - info = infos[idx] - if task == "segment": - x1, y1, x2, y2, mask_tensor = info - width, height = x2 - x1, y2 - y1 - tile_crop = tile_img.crop((0, 0, width, height)) - mask = tensor2pil(mask_tensor) - canvas.paste(tile_crop, (x1, y1, x2, y2), mask) - else: - x1, y1, x2, y2 = info - width, height = x2 - x1, y2 - y1 - tile_crop = tile_img.crop((0, 0, width, height)) - canvas.paste(tile_crop, (x1, y1, x2, y2)) + x1, y1, x2, y2 = positions[idx] + width = x2 - x1 + height = y2 - y1 + tile_img = tensor2pil(tile.unsqueeze(0)).crop((0, 0, width, height)) + canvas.paste(tile_img, (x1, y1, x2, y2)) return pil2tensor(canvas) diff --git a/examples/smart_tile_example.json b/examples/smart_tile_example.json index 2a03839..c47bd64 100644 --- a/examples/smart_tile_example.json +++ b/examples/smart_tile_example.json @@ -14,14 +14,14 @@ "type": "TTP_Smart_Tile_Batch", "pos": [250, 0], "size": [250, 120], - "widgets_values": ["yolov8n-seg.pt", "detect", 0.25, 0.45, 10] + "widgets_values": [[ [10, 10, 110, 110], [130, 130, 220, 220] ]] }, { "id": 3, "type": "TTP_Smart_Image_Assy", "pos": [500, 0], "size": [250, 120], - "widgets_values": ["detect"] + "widgets_values": [] } ], "links": [ diff --git a/pyproject.toml b/pyproject.toml index 5f48d16..18410a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,9 @@ [project] name = "comfyui_ttp_toolset" description = "This is a workflow for my simple logic amazing upscale node for DIT model. it can be common use for Flux,Hunyuan,SD3 It can simple tile the initial image into pieces and then use image-interrogator to get each tile prompts for more accurate upscale process. The condition will be properly handled and the hallucination will be significantly eliminated." -version = "1.0.6" +version = "1.0.7" license = {file = "LICENSE"} -dependencies = ["ultralytics>=8.1"] +dependencies = [] [project.urls] Repository = "https://github.com/TTPlanetPig/Comfyui_TTP_Toolset" From f69476be377609beee7ee9880a54a352f7dc6c5e Mon Sep 17 00:00:00 2001 From: TTPlanetPig <152850462+TTPlanetPig@users.noreply.github.com> Date: Thu, 12 Jun 2025 17:06:12 +0800 Subject: [PATCH 4/4] Add interactive crop tool --- README.md | 11 ++++++ manual_crop_tool.py | 84 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 manual_crop_tool.py diff --git a/README.md b/README.md index 21c1111..641b25e 100644 --- a/README.md +++ b/README.md @@ -113,6 +113,17 @@ Bounding boxes should be provided as `[x1, y1, x2, y2]` coordinates. See `examples/smart_tile_example.json` for a minimal workflow using these nodes. +You can also draw bounding boxes interactively using `manual_crop_tool.py`: + +```bash +python manual_crop_tool.py input.jpg output_dir --tile-size 512 +``` + +Draw rectangles on the displayed image. Press `q` when finished. The script saves +all tiles to `output_dir` and writes `boxes.json` containing the coordinates. +Large selections are automatically split into smaller tiles using the specified +tile size. + --- ## **Examples** diff --git a/manual_crop_tool.py b/manual_crop_tool.py new file mode 100644 index 0000000..0b8bb22 --- /dev/null +++ b/manual_crop_tool.py @@ -0,0 +1,84 @@ +import argparse +import json +import os +from PIL import Image +import matplotlib.pyplot as plt +from matplotlib.widgets import RectangleSelector + + +def split_box(box, tile_w, tile_h): + x1, y1, x2, y2 = box + boxes = [] + for top in range(y1, y2, tile_h): + for left in range(x1, x2, tile_w): + boxes.append([ + left, + top, + min(left + tile_w, x2), + min(top + tile_h, y2), + ]) + return boxes + + +def crop_tiles(img, boxes, out_dir): + tiles = [] + for idx, b in enumerate(boxes): + tile = img.crop(b) + tiles.append(tile) + tile.save(os.path.join(out_dir, f"tile_{idx}.png")) + return tiles + + +def main(): + parser = argparse.ArgumentParser(description="Draw rectangles on an image and export tiles") + parser.add_argument("image", help="Path to image") + parser.add_argument("output", help="Directory to save tiles") + parser.add_argument("--tile-size", type=int, default=512, help="Max tile size when splitting selections") + args = parser.parse_args() + + os.makedirs(args.output, exist_ok=True) + img = Image.open(args.image).convert("RGB") + + fig, ax = plt.subplots() + ax.imshow(img) + ax.set_title("Drag rectangles, press 'q' when done") + + boxes = [] + + def onselect(eclick, erelease): + x1, y1 = int(eclick.xdata), int(eclick.ydata) + x2, y2 = int(erelease.xdata), int(erelease.ydata) + if x2 < x1: + x1, x2 = x2, x1 + if y2 < y1: + y1, y2 = y2, y1 + ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='red')) + boxes.append([x1, y1, x2, y2]) + fig.canvas.draw() + + toggle_selector = RectangleSelector(ax, onselect, drawtype="box", useblit=True) + + def on_key(event): + if event.key == 'q': + plt.close(event.canvas.figure) + fig.canvas.mpl_connect('key_press_event', on_key) + plt.show() + + final_boxes = [] + for b in boxes: + width = b[2] - b[0] + height = b[3] - b[1] + if width > args.tile_size or height > args.tile_size: + final_boxes.extend(split_box(b, args.tile_size, args.tile_size)) + else: + final_boxes.append(b) + + with open(os.path.join(args.output, "boxes.json"), "w", encoding="utf-8") as f: + json.dump(final_boxes, f) + + crop_tiles(img, final_boxes, args.output) + print(f"Saved {len(final_boxes)} tiles to {args.output}") + + +if __name__ == "__main__": + main()