nodes/ltdrdata-ComfyUI-Impact-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

MaskPainter, EmptySEGS node added.

Browse Source 
bug patch about crop_factor on MaskToSEGS node
This commit is contained in:
Dr.Lt.Data
2023-04-16 22:46:09 +09:00
parent cc5cf656e8
commit e227cfcee1
8 changed files with 1868 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+15
View File
@@ -23,6 +23,8 @@
* Segs & Masks - Perform a bitwise AND operation on SEGS and MASK.
* MaskToSegs - This node generates SEGS based on the mask.
* ToBinaryMask - This node separates the mask generated with alpha values between 0 and 255 into 0 and 255. The non-zero parts are always set to 255.
* EmptySEGS - This node provides a empty SEGS.
* MaskPainter - This node provides a feature to draw masks.
# Installation
@@ -106,6 +108,19 @@
* Unlike in face detection, for non-rigid objects, the center point may not always be the segmentation area, so be careful not to assume that the center point is always the segmentation area.
#### Mask Painter
![example](misc/maskpainter1.png)
* Click "Edit mask" button
* **Don't connect to 'mask_image' input**
![example](misc/maskpainter2.png)
* You can draw mask on Mask Painter
* Currently, this editor only provides basic functionalities
![example](misc/maskpainter-original.png) ![example](misc/maskpainter-result.png)
* When used together, SAMDetector and MaskPainter can be used to enhance specific elements of an image.
# Others Tutorials
* [ONNX Tutorial](tutorial/ONNX.md)
comfyui-impact-pack.py
+154 -11
View File
@@ -1,14 +1,54 @@
import os, sys, subprocess
from torchvision.datasets.utils import download_url
import platform
import configparser
import folder_paths
# INSTALL
print("### Loading: ComfyUI-Impact-Pack")
comfy_path = os.path.dirname(folder_paths.__file__)
config_path = os.path.join(comfy_path, "custom_nodes", "impact-pack.ini")
js_path = os.path.join(comfy_path, "web", "extensions", "core")
js_version = 1
js_url = "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Impact-Pack/Main/js/impact-pack.js"
def read_js_version():
try:
config = configparser.ConfigParser()
config.read(config_path)
return int(config['default']['js_version'])
except:
return 0
def write_js_version():
config = configparser.ConfigParser()
config['default'] = {
'js_version': js_version
}
with open(config_path, 'w') as configfile:
config.write(configfile)
# Download js
if not os.path.exists(os.path.join(js_path, "impact-pack.js")):
download_url(js_url, js_path)
write_js_version()
else:
# update js
js = read_js_version()
if js < js_version:
download_url(js_url, js_path)
write_js_version()
# ----- SETUP --------------------------------------------------------------
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
sys.path.append('../ComfyUI')
# INSTALL
print("Loading: ComfyUI-Impact-Pack")
print("### ComfyUI-Impact-Pack: Check dependencies")
@@ -57,17 +97,12 @@ def ensure_mmdet_package():
subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmdet==3.0.0'])
subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmengine==0.7.2'])
ensure_pip_packages()
ensure_mmdet_package()
# Download model
print("### ComfyUI-Impact-Pack: Check basic models")
import folder_paths
comfy_path = os.path.dirname(folder_paths.__file__)
model_path = folder_paths.models_dir
bbox_path = os.path.join(model_path, "mmdets", "bbox")
@@ -313,14 +348,14 @@ class NO_SEGM_MODEL:
def normalize_region(limit, startp, size):
if startp < 0:
new_endp = size
new_endp = min(limit, size)
new_startp = 0
elif startp + size > limit:
new_startp = limit - size
new_endp = limit
else:
new_startp = startp
new_endp = startp+size
new_endp = min(limit, startp+size)
return int(new_startp), int(new_endp)
@@ -705,6 +740,19 @@ class DetailerForEachTest(DetailerForEach):
else:
return enhanced_img, cropped, cropped_enhanced,
class EmptySEGS:
@classmethod
def INPUT_TYPES(s):
return { }
RETURN_TYPES = ("SEGS",)
FUNCTION = "doit"
CATEGORY = "ImpactPack/Util"
def doit(self):
return ([],)
class SegsMaskCombine:
@classmethod
@@ -1119,11 +1167,11 @@ class MaskToSEGS:
# iterate over the regions and print their bounding boxes
for region in regions:
y1, x1, y2, x2 = region.bbox
bbox = x1, x2, y1, y2
bbox = x1, y1, x2, y2
crop_region = make_crop_region(mask.shape[1], mask.shape[0], bbox, crop_factor)
if x2 - x1 > 0 and y2 - y1 > 0:
cropped_mask = mask[y1:y2, x1:x2]
cropped_mask = mask[crop_region[1]:crop_region[3], crop_region[0]:crop_region[2]]
result.append((None, cropped_mask, 1.0, crop_region, bbox))
if not result:
@@ -1240,6 +1288,62 @@ class SubtractMask:
return (mask,)
import nodes
class MaskPainter(nodes.PreviewImage):
@classmethod
def INPUT_TYPES(s):
return {"required":
{"images": ("IMAGE", ), },
"hidden": {
"prompt": "PROMPT",
"extra_pnginfo": "EXTRA_PNGINFO",
},
"optional": { "mask_image": ("IMAGE_PATH", ), },
}
RETURN_TYPES = ("MASK", )
FUNCTION = "save_painted_images"
CATEGORY = "ImpactPack"
def load_mask(self, imagepath):
if imagepath['type'] == "temp":
input_dir = folder_paths.get_temp_directory()
else:
input_dir = folder_paths.get_input_directory()
image_path = os.path.join(input_dir, imagepath['filename'])
if os.path.exists(image_path):
i = Image.open(image_path)
if 'A' in i.getbands():
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
return (mask, )
def save_painted_images(self, images, filename_prefix="impact-mask",
prompt=None, extra_pnginfo=None, mask_image=None):
res = self.save_images(images, filename_prefix, prompt, extra_pnginfo)
if mask_image is not None:
res['result'] = self.load_mask(mask_image)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
res['result'] = (mask, )
return res
NODE_CLASS_MAPPINGS = {
"MMDetLoader": MMDetLoader,
"SAMLoader": SAMLoader,
@@ -1262,7 +1366,46 @@ NODE_CLASS_MAPPINGS = {
"SubtractMask": SubtractMask,
"Segs & Mask": SegsBitwiseAndMask,
"SegsMaskCombine": SegsMaskCombine,
"EmptySegs": EmptySEGS,
"MaskToSEGS": MaskToSEGS,
"ToBinaryMask": ToBinaryMask,
"MaskPainter": MaskPainter,
}
import server
from aiohttp import web
@server.PromptServer.instance.routes.post("/upload/temp")
async def upload_image(request):
upload_dir = folder_paths.get_temp_directory()
if not os.path.exists(upload_dir):
os.makedirs(upload_dir)
post = await request.post()
image = post.get("image")
if image and image.file:
filename = image.filename
if not filename:
return web.Response(status=400)
split = os.path.splitext(filename)
i = 1
while os.path.exists(os.path.join(upload_dir, filename)):
filename = f"{split[0]} ({i}){split[1]}"
i += 1
filepath = os.path.join(upload_dir, filename)
with open(filepath, "wb") as f:
f.write(image.file.read())
return web.json_response({"name" : filename})
else:
return web.Response(status=400)
js/impact-pack.js
+247
View File
@@ -0,0 +1,247 @@
import { app } from "/scripts/app.js";
import { ComfyDialog, $el } from "/scripts/ui.js";
// Helper function to convert a data URL to a Blob object
function dataURLToBlob(dataURL) {
const parts = dataURL.split(';base64,');
const contentType = parts[0].split(':')[1];
const byteString = atob(parts[1]);
const arrayBuffer = new ArrayBuffer(byteString.length);
const uint8Array = new Uint8Array(arrayBuffer);
for (let i = 0; i < byteString.length; i++) {
uint8Array[i] = byteString.charCodeAt(i);
}
return new Blob([arrayBuffer], { type: contentType });
}
async function invalidateImage(node, formData) {
const filepath = node.images[0];
await fetch('/upload/temp', {
method: 'POST',
body: formData
}).then(response => {
}).catch(error => {
console.error('Error:', error);
});
const img = new Image();
img.onload = () => {
node.imgs = [img];
app.graph.setDirtyCanvas(true);
};
img.src = `view?filename=${filepath.filename}&type=${filepath.type}`;;
}
class ImpactInpaintDialog extends ComfyDialog {
constructor() {
super();
this.element = $el("div.comfy-modal", { parent: document.body },
[
$el("div.comfy-modal-content",
[
...this.createButtons()]),
]);
}
createButtons() {
return [
$el("button", {
type: "button",
textContent: "Save",
onclick: () => {
const backupCtx = this.backupCanvas.getContext('2d', {transparent: true});
backupCtx.clearRect(0,0,this.backupCanvas.width,this.backupCanvas.height);
backupCtx.drawImage(this.maskCanvas,
0, 0, this.maskCanvas.width, this.maskCanvas.height,
0, 0, this.backupCanvas.width, this.backupCanvas.height);
// paste mask data into alpha channel
const backupData = backupCtx.getImageData(0, 0, this.backupCanvas.width, this.backupCanvas.height);
for (let i = 0; i < backupData.data.length; i += 4) {
if(backupData.data[i+3] == 255)
backupData.data[i+3] = 0;
else
backupData.data[i+3] = 255;
backupData.data[i] = 0;
backupData.data[i+1] = 0;
backupData.data[i+2] = 0;
}
backupCtx.globalCompositeOperation = 'source-over';
backupCtx.putImageData(backupData, 0, 0);
const dataURL = this.backupCanvas.toDataURL();
const blob = dataURLToBlob(dataURL);
/*
// copy image data
backupCtx.globalCompositeOperation = 'copy';
backupCtx.globalAlpha = 1.0;
backupCtx.drawImage(this.image, 0, 0);
backupCtx.globalCompositeOperation = 'source-over';
const backupData2 = backupCtx.getImageData(0, 0, this.backupCanvas.width, this.backupCanvas.height);
// restore alpha channel
var cnt_r = 0;
for (let i = 0; i < backupData2.data.length; i += 4) {
if(backupData2.data[i] == 0) {
cnt_r++;
}
backupData2.data[i + 3] = backupData.data[i + 3];
}
// I don't know why RGB channel is effected by this code....
backupCtx.putImageData(backupData2, 0, 0);
const dataURL2 = this.backupCanvas.toDataURL();
const blob2 = dataURLToBlob(dataURL2);
*/
const formData = new FormData();
const filename = "impact-mask-" + performance.now() + ".png";
const item =
{
"filename": filename,
"subfolder": "",
"type": "temp",
};
this.node.images[0] = item;
this.node.widgets[1].value = item;
formData.append('image', blob, filename);
invalidateImage(this.node, formData);
this.close();
}
}),
$el("button", {
type: "button",
textContent: "Cancel",
onclick: () => this.close(),
}),
$el("button", {
type: "button",
textContent: "Clear",
onclick: () => {
this.maskCtx.clearRect(0, 0, maskCanvas.width, maskCanvas.height);
},
}),
];
}
show() {
const imgCanvas = document.createElement('canvas');
const maskCanvas = document.createElement('canvas');
const backupCanvas = document.createElement('canvas');
imgCanvas.id = "imageCanvas";
maskCanvas.id = "maskCanvas";
backupCanvas.id = "backupCanvas";
this.element.appendChild(imgCanvas);
this.element.appendChild(maskCanvas);
this.node.widgets[1].value = null;
this.element.style.display = "block";
imgCanvas.style.position = "relative";
imgCanvas.style.top = "200";
imgCanvas.style.left = "0";
maskCanvas.style.position = "absolute";
const imgCtx = imgCanvas.getContext('2d');
const maskCtx = maskCanvas.getContext('2d');
const backupCtx = backupCanvas.getContext('2d');
this.maskCanvas = maskCanvas;
this.maskCtx = maskCtx;
this.backupCanvas = backupCanvas;
window.addEventListener("resize", () => {
// repositioning
imgCanvas.width = window.innerWidth - 250;
imgCanvas.height = window.innerHeight - 300;
// redraw image
let drawWidth = image.width;
let drawHeight = image.height;
if (image.width > imgCanvas.width) {
drawWidth = imgCanvas.width;
drawHeight = (drawWidth / image.width) * image.height;
}
if (drawHeight > imgCanvas.height) {
drawHeight = imgCanvas.height;
drawWidth = (drawHeight / image.height) * image.width;
}
imgCtx.drawImage(image, 0, 0, drawWidth, drawHeight);
// update mask
backupCtx.drawImage(maskCanvas, 0, 0, maskCanvas.width, maskCanvas.height, 0, 0, backupCanvas.width, backupCanvas.height);
maskCanvas.width = drawWidth;
maskCanvas.height = drawHeight;
maskCanvas.style.top = imgCanvas.offsetTop + "px";
maskCanvas.style.left = imgCanvas.offsetLeft + "px";
maskCtx.drawImage(backupCanvas, 0, 0, backupCanvas.width, backupCanvas.height, 0, 0, maskCanvas.width, maskCanvas.height);
});
// image load
const image = new Image();
image.onload = function() {
backupCanvas.width = image.width;
backupCanvas.height = image.height;
window.dispatchEvent(new Event('resize'));
};
const filepath = this.node.images[0];
image.src = `view?filename=${filepath.filename}&type=${filepath.type}`;
this.image = image;
// event handler for user drawing ------
let isDrawing = false;
function draw_move(event) {
event.preventDefault();
const maskRect = maskCanvas.getBoundingClientRect();
const x = event.offsetX || event.targetTouches[0].clientX - maskRect.left;
const y = event.offsetY || event.targetTouches[0].clientY - maskRect.top;
maskCtx.beginPath();
maskCtx.fillStyle = "rgb(0,0,0)";
maskCtx.arc(x, y, 10, 0, Math.PI * 2, false);
maskCtx.fill();
}
maskCanvas.addEventListener('mousedown', function(event) { event.preventDefault(); isDrawing = true; } );
maskCanvas.addEventListener('mousemove', function(event) { if(isDrawing) { draw_move(event); } } );
maskCanvas.addEventListener('mouseup', function(event) { isDrawing = false; } );
maskCanvas.addEventListener('touchmove', draw_move);
}
}
app.registerExtension({
name: "Comfy.Impack",
nodeCreated(node, app) {
if(node.comfyClass == "MaskPainter") {
node.addWidget("button", "Edit mask", null, () => {
this.dlg = new ImpactInpaintDialog(app);
this.dlg.node = node;
this.dlg.show();
});
node.addWidget("hidden", "mask_image", null, null);
}
}
});
misc/maskpainter-original.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

misc/maskpainter-result.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

misc/maskpainter.json
+1452
View File
File diff suppressed because it is too large Load Diff
misc/maskpainter1.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 108 KiB

misc/maskpainter2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB