Clipspace SAM Detector Feature added
This commit is contained in:
Dr.Lt.Data
@@ -4,7 +4,7 @@ This custom node helps to conveniently enhance images through Detector, Detailer
|
||||
|
||||
## Custom nodes pack for ComfyUI
|
||||
|
||||
# Features
|
||||
# Custom Nodes
|
||||
* SAMLoader - Load SAM model
|
||||
* MMDetDetectorProvider - Load MMDet model to provide BBOX_DETECTOR, SEGM_DETECTOR
|
||||
* ONNXDetectorProvider - Load ONNX model to provide SEGM_DETECTOR
|
||||
@@ -62,6 +62,10 @@ This takes latent as input and outputs latent as the result.
|
||||
|
||||
* PreviewBridge - This custom node can be used with a bridge when using the MaskEditor feature of Clipspace.
|
||||
|
||||
# Feature
|
||||
* SAM Detector (Clipspace) - When you right-click on a node that has 'MASK' and 'IMAGE' outputs, a context menu will open. From this menu, you can either open a dialog to create a SAM Mask using 'Open in SAM Detector', or copy the content (likely mask data) using 'Copy (Clipspace)' and generate a mask using 'Impact SAM Detector' from the clipspace menu, and then paste it using 'Paste (Clipspace)'.
|
||||
|
||||
|
||||
# Depercated
|
||||
* The following nodes have been kept only for compatibility with existing workflows, and are no longer supported. Please replace them with new nodes.
|
||||
* MMDetLoader -> MMDetDetectorProvider
|
||||
@@ -148,6 +152,21 @@ This takes latent as input and outputs latent as the result.
|
||||
 
|
||||
|
||||
|
||||
#### 5. SAM Detector (Clipspace)
|
||||
|
||||
* When you right-click on the node that outputs 'MASK' and 'IMAGE', a menu called "Open in SAM Detector" appears, as shown in the following picture. Clicking on the menu opens a dialog in SAM's functionality, allowing you to generate a segment mask.
|
||||
|
||||
|
||||
* By clicking the left mouse button on a coordinate, a positive prompt in blue color is entered, indicating the area that should be included. Clicking the right mouse button on a coordinate enters a negative prompt in red color, indicating the area that should be excluded. Positive prompts represent the areas that should be included, while negative prompts represent the areas that should be excluded.
|
||||
* You can remove the points that were added by using the "undo" button. After selecting the points, pressing the "detect" button generates the mask. Additionally, you can adjust the fidelity slider to determine the extent to which the mask belongs to the confidence region.
|
||||
|
||||
|
||||
|
||||
* If you opened the dialog through "Open in SAM Detector" from the node, you can directly apply the changes by clicking the "Save to node" button. However, if you opened the dialog through the "clipspace" menu, you can save it to clipspace by clicking the "Save" button.
|
||||
|
||||
|
||||
|
||||
* When you execute using the reflected mask in the node, you can observe that the image and mask are displayed separately.
|
||||
|
||||
# Others Tutorials
|
||||
* [ComfyUI-extension-tutorials/ComfyUI-Impact-Pack](https://github.com/ltdrdata/ComfyUI-extension-tutorials/tree/Main/ComfyUI-Impact-Pack) - You can find various tutorials and workflows on this page.
|
||||
|
||||
@@ -46,8 +46,12 @@ def setup_js():
|
||||
js_dest_path = os.path.join(comfy_path, "web", "extensions", "impact-pack")
|
||||
if not os.path.exists(js_dest_path):
|
||||
os.makedirs(js_dest_path)
|
||||
|
||||
js_src_path = os.path.join(impact_path, "js", "impact-pack.js")
|
||||
shutil.copy(js_src_path, js_dest_path)
|
||||
|
||||
js_src_path = os.path.join(impact_path, "js", "impact-sam-editor.js")
|
||||
shutil.copy(js_src_path, js_dest_path)
|
||||
|
||||
setup_js()
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
import configparser
|
||||
import os
|
||||
|
||||
version = "V2.5.1"
|
||||
version = "V2.6"
|
||||
|
||||
dependency_version = 1
|
||||
|
||||
|
||||
@@ -1,8 +1,18 @@
|
||||
import os
|
||||
import threading
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
import folder_paths
|
||||
|
||||
import impact_core as core
|
||||
import impact_pack
|
||||
from segment_anything import SamPredictor, sam_model_registry
|
||||
import numpy as np
|
||||
import nodes
|
||||
import torch
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
@server.PromptServer.instance.routes.post("/upload/temp")
|
||||
async def upload_image(request):
|
||||
@@ -33,3 +43,107 @@ async def upload_image(request):
|
||||
return web.json_response({"name": filename})
|
||||
else:
|
||||
return web.Response(status=400)
|
||||
|
||||
|
||||
sam_predictor = None
|
||||
default_sam_model_name = os.path.join(impact_pack.model_path, "sams", "sam_vit_b_01ec64.pth")
|
||||
|
||||
sam_lock = threading.Condition()
|
||||
|
||||
last_prepare_data = None
|
||||
|
||||
@server.PromptServer.instance.routes.post("/sam/prepare")
|
||||
async def load_sam_model(request):
|
||||
global sam_predictor
|
||||
global last_prepare_data
|
||||
data = await request.json()
|
||||
|
||||
with sam_lock:
|
||||
if last_prepare_data is not None and last_prepare_data == data:
|
||||
# already loaded: skip -- prevent redundant loading
|
||||
return web.Response(status=200)
|
||||
|
||||
last_prepare_data = data
|
||||
|
||||
model_name = os.path.join(impact_pack.model_path, "sams", data['sam_model_name'])
|
||||
|
||||
print(f"ComfyUI-Impact-Pack: Loading SAM model '{impact_pack.model_path}'")
|
||||
|
||||
filename, image_dir = folder_paths.annotated_filepath(data["filename"])
|
||||
|
||||
if image_dir is None:
|
||||
typ = data['type'] if data['type'] != '' else 'output'
|
||||
image_dir = folder_paths.get_directory_by_type(typ)
|
||||
|
||||
if image_dir is None:
|
||||
return web.Response(status=400)
|
||||
|
||||
if 'vit_h' in model_name:
|
||||
model_kind = 'vit_h'
|
||||
elif 'vit_l' in model_name:
|
||||
model_kind = 'vit_l'
|
||||
else:
|
||||
model_kind = 'vit_b'
|
||||
|
||||
sam_model = sam_model_registry[model_kind](checkpoint=model_name)
|
||||
sam_predictor = SamPredictor(sam_model)
|
||||
|
||||
image_path = os.path.join(image_dir, filename)
|
||||
image = nodes.LoadImage().load_image(image_path)[0]
|
||||
image = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
||||
|
||||
sam_predictor.set_image(image, "RGB")
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/sam/release")
|
||||
async def release_sam(request):
|
||||
global sam_predictor
|
||||
|
||||
with sam_lock:
|
||||
sam_predictor = None
|
||||
|
||||
print(f"ComfyUI-Impact-Pack: unloading SAM model")
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/sam/detect")
|
||||
async def sam_detect(request):
|
||||
global sam_predictor
|
||||
with sam_lock:
|
||||
if sam_predictor is not None:
|
||||
data = await request.json()
|
||||
|
||||
positive_points = data['positive_points']
|
||||
negative_points = data['negative_points']
|
||||
threshold = data['threshold']
|
||||
|
||||
points = []
|
||||
plabs = []
|
||||
|
||||
for p in positive_points:
|
||||
points.append(p)
|
||||
plabs.append(1)
|
||||
|
||||
for p in negative_points:
|
||||
points.append(p)
|
||||
plabs.append(0)
|
||||
|
||||
detected_masks = core.sam_predict(sam_predictor, points, plabs, None, threshold)
|
||||
mask = core.combine_masks2(detected_masks)
|
||||
|
||||
if mask is None:
|
||||
return web.Response(status=400)
|
||||
|
||||
image = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
||||
i = 255. * image.cpu().numpy()
|
||||
|
||||
img = Image.fromarray(np.clip(i[0], 0, 255).astype(np.uint8))
|
||||
|
||||
img_buffer = io.BytesIO()
|
||||
img.save(img_buffer, format='png')
|
||||
|
||||
headers = {'Content-Type': 'image/png'}
|
||||
|
||||
return web.Response(body=img_buffer.getvalue(), headers=headers)
|
||||
|
||||
else:
|
||||
return web.Response(status=400)
|
||||
|
||||
@@ -265,9 +265,7 @@ function executeHandler(event) {
|
||||
if(event.detail.output.aux){
|
||||
const id = event.detail.node;
|
||||
if(input_tracking.hasOwnProperty(id)) {
|
||||
console.log(input_tracking[id]);
|
||||
if(input_tracking.hasOwnProperty(id) && input_tracking[id][0] != event.detail.output.aux[0]) {
|
||||
console.log(`dirty(${id}) = ${input_tracking[id][0]} != ${event.detail.output.aux[0]}`);
|
||||
input_dirty[id] = true;
|
||||
}
|
||||
else{
|
||||
|
||||
@@ -0,0 +1,626 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { ComfyDialog, $el } from "/scripts/ui.js";
|
||||
import { ComfyApp } from "/scripts/app.js";
|
||||
import { ClipspaceDialog } from "/extensions/core/clipspace.js";
|
||||
|
||||
function addMenuHandler(nodeType, cb) {
|
||||
const getOpts = nodeType.prototype.getExtraMenuOptions;
|
||||
nodeType.prototype.getExtraMenuOptions = function () {
|
||||
const r = getOpts.apply(this, arguments);
|
||||
cb.apply(this, arguments);
|
||||
return r;
|
||||
};
|
||||
}
|
||||
|
||||
// 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 });
|
||||
}
|
||||
|
||||
function loadedImageToBlob(image) {
|
||||
const canvas = document.createElement('canvas');
|
||||
|
||||
canvas.width = image.width;
|
||||
canvas.height = image.height;
|
||||
|
||||
const ctx = canvas.getContext('2d');
|
||||
|
||||
ctx.drawImage(image, 0, 0);
|
||||
|
||||
const dataURL = canvas.toDataURL('image/png', 1);
|
||||
const blob = dataURLToBlob(dataURL);
|
||||
|
||||
return blob;
|
||||
}
|
||||
|
||||
async function uploadMask(filepath, formData) {
|
||||
await fetch('/upload/mask', {
|
||||
method: 'POST',
|
||||
body: formData
|
||||
}).then(response => {}).catch(error => {
|
||||
console.error('Error:', error);
|
||||
});
|
||||
|
||||
ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']] = new Image();
|
||||
ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src = `view?filename=${filepath.filename}&type=${filepath.type}`;
|
||||
|
||||
if(ComfyApp.clipspace.images)
|
||||
ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']] = filepath;
|
||||
|
||||
ClipspaceDialog.invalidatePreview();
|
||||
}
|
||||
|
||||
class ImpactSamEditorDialog extends ComfyDialog {
|
||||
static instance = null;
|
||||
|
||||
static getInstance() {
|
||||
if(!ImpactSamEditorDialog.instance) {
|
||||
ImpactSamEditorDialog.instance = new ImpactSamEditorDialog();
|
||||
}
|
||||
|
||||
return ImpactSamEditorDialog.instance;
|
||||
}
|
||||
|
||||
constructor() {
|
||||
super();
|
||||
this.element = $el("div.comfy-modal", { parent: document.body },
|
||||
[ $el("div.comfy-modal-content",
|
||||
[...this.createButtons()]),
|
||||
]);
|
||||
}
|
||||
|
||||
createButtons() {
|
||||
return [];
|
||||
}
|
||||
|
||||
createButton(name, callback) {
|
||||
var button = document.createElement("button");
|
||||
button.innerText = name;
|
||||
button.addEventListener("click", callback);
|
||||
return button;
|
||||
}
|
||||
|
||||
createLeftButton(name, callback) {
|
||||
var button = this.createButton(name, callback);
|
||||
button.style.cssFloat = "left";
|
||||
button.style.marginRight = "4px";
|
||||
return button;
|
||||
}
|
||||
|
||||
createRightButton(name, callback) {
|
||||
var button = this.createButton(name, callback);
|
||||
button.style.cssFloat = "right";
|
||||
button.style.marginLeft = "4px";
|
||||
return button;
|
||||
}
|
||||
|
||||
createLeftSlider(self, name, callback) {
|
||||
const divElement = document.createElement('div');
|
||||
divElement.id = "sam-fidelity-slider";
|
||||
divElement.style.cssFloat = "left";
|
||||
divElement.style.fontFamily = "sans-serif";
|
||||
divElement.style.marginRight = "4px";
|
||||
divElement.style.color = "var(--input-text)";
|
||||
divElement.style.backgroundColor = "var(--comfy-input-bg)";
|
||||
divElement.style.borderRadius = "8px";
|
||||
divElement.style.borderColor = "var(--border-color)";
|
||||
divElement.style.borderStyle = "solid";
|
||||
divElement.style.fontSize = "15px";
|
||||
divElement.style.height = "21px";
|
||||
divElement.style.padding = "1px 6px";
|
||||
divElement.style.display = "flex";
|
||||
divElement.style.position = "relative";
|
||||
divElement.style.top = "2px";
|
||||
self.fidelity_slider_input = document.createElement('input');
|
||||
self.fidelity_slider_input.setAttribute('type', 'range');
|
||||
self.fidelity_slider_input.setAttribute('min', '0');
|
||||
self.fidelity_slider_input.setAttribute('max', '100');
|
||||
self.fidelity_slider_input.setAttribute('value', '70');
|
||||
const labelElement = document.createElement("label");
|
||||
labelElement.textContent = name;
|
||||
|
||||
divElement.appendChild(labelElement);
|
||||
divElement.appendChild(self.fidelity_slider_input);
|
||||
|
||||
self.fidelity_slider_input.addEventListener("change", callback);
|
||||
|
||||
return divElement;
|
||||
}
|
||||
|
||||
async detect_and_invalidate_mask_canvas(self) {
|
||||
const mask_img = await self.detect(self);
|
||||
|
||||
const canvas = self.maskCtx.canvas;
|
||||
const ctx = self.maskCtx;
|
||||
|
||||
ctx.clearRect(0, 0, canvas.width, canvas.height);
|
||||
|
||||
await new Promise((resolve, reject) => {
|
||||
self.mask_image = new Image();
|
||||
self.mask_image.onload = function() {
|
||||
ctx.drawImage(self.mask_image, 0, 0, canvas.width, canvas.height);
|
||||
resolve();
|
||||
};
|
||||
self.mask_image.onerror = reject;
|
||||
self.mask_image.src = mask_img.src;
|
||||
});
|
||||
}
|
||||
|
||||
setlayout(imgCanvas, maskCanvas, pointsCanvas) {
|
||||
const self = this;
|
||||
|
||||
// If it is specified as relative, using it only as a hidden placeholder for padding is recommended
|
||||
// to prevent anomalies where it exceeds a certain size and goes outside of the window.
|
||||
var placeholder = document.createElement("div");
|
||||
placeholder.style.position = "relative";
|
||||
placeholder.style.height = "50px";
|
||||
|
||||
var bottom_panel = document.createElement("div");
|
||||
bottom_panel.style.position = "absolute";
|
||||
bottom_panel.style.bottom = "0px";
|
||||
bottom_panel.style.left = "20px";
|
||||
bottom_panel.style.right = "20px";
|
||||
bottom_panel.style.height = "50px";
|
||||
|
||||
var brush = document.createElement("div");
|
||||
brush.id = "sam-brush";
|
||||
brush.style.backgroundColor = "blue";
|
||||
brush.style.outline = "2px solid pink";
|
||||
brush.style.borderRadius = "50%";
|
||||
brush.style.MozBorderRadius = "50%";
|
||||
brush.style.WebkitBorderRadius = "50%";
|
||||
brush.style.position = "absolute";
|
||||
brush.style.zIndex = 100;
|
||||
brush.style.pointerEvents = "none";
|
||||
this.brush = brush;
|
||||
this.element.appendChild(imgCanvas);
|
||||
this.element.appendChild(maskCanvas);
|
||||
this.element.appendChild(pointsCanvas);
|
||||
this.element.appendChild(placeholder); // must below z-index than bottom_panel to avoid covering button
|
||||
this.element.appendChild(bottom_panel);
|
||||
document.body.appendChild(brush);
|
||||
this.brush_size = 5;
|
||||
|
||||
var fidelity_slider = this.createLeftSlider(self, "Fidelity", (event) => {
|
||||
self.fidelity = event.target.value;
|
||||
});
|
||||
|
||||
var clearButton = this.createLeftButton("Clear", () => {
|
||||
self.maskCtx.clearRect(0, 0, self.maskCanvas.width, self.maskCanvas.height);
|
||||
self.pointsCtx.clearRect(0, 0, self.pointsCanvas.width, self.pointsCanvas.height);
|
||||
|
||||
self.prompt_points = [];
|
||||
|
||||
self.invalidatePointsCanvas(self);
|
||||
});
|
||||
|
||||
var detectButton = this.createLeftButton("Detect", () => self.detect_and_invalidate_mask_canvas(self));
|
||||
|
||||
var cancelButton = this.createRightButton("Cancel", () => {
|
||||
document.removeEventListener("mouseup", ImpactSamEditorDialog.handleMouseUp);
|
||||
document.removeEventListener("keydown", ImpactSamEditorDialog.handleKeyDown);
|
||||
self.close();
|
||||
});
|
||||
|
||||
self.saveButton = this.createRightButton("Save", () => {
|
||||
document.removeEventListener("mouseup", ImpactSamEditorDialog.handleMouseUp);
|
||||
document.removeEventListener("keydown", ImpactSamEditorDialog.handleKeyDown);
|
||||
self.save(self);
|
||||
});
|
||||
|
||||
var undoButton = this.createLeftButton("Undo", () => {
|
||||
if(self.prompt_points.length > 0) {
|
||||
self.prompt_points.pop();
|
||||
self.pointsCtx.clearRect(0, 0, self.pointsCanvas.width, self.pointsCanvas.height);
|
||||
self.invalidatePointsCanvas(self);
|
||||
}
|
||||
});
|
||||
|
||||
bottom_panel.appendChild(clearButton);
|
||||
bottom_panel.appendChild(detectButton);
|
||||
bottom_panel.appendChild(self.saveButton);
|
||||
bottom_panel.appendChild(cancelButton);
|
||||
bottom_panel.appendChild(fidelity_slider);
|
||||
bottom_panel.appendChild(undoButton);
|
||||
|
||||
imgCanvas.style.position = "relative";
|
||||
imgCanvas.style.top = "200";
|
||||
imgCanvas.style.left = "0";
|
||||
|
||||
maskCanvas.style.position = "absolute";
|
||||
maskCanvas.style.opacity = 0.5;
|
||||
pointsCanvas.style.position = "absolute";
|
||||
}
|
||||
|
||||
show() {
|
||||
this.mask_image = null;
|
||||
self.prompt_points = [];
|
||||
|
||||
this.message_box = $el("p", ["Please wait a moment while the SAM model and the image are being loaded."]);
|
||||
this.element.appendChild(this.message_box);
|
||||
|
||||
if(self.imgCtx) {
|
||||
self.imgCtx.clearRect(0, 0, self.imageCanvas.width, self.imageCanvas.height);
|
||||
}
|
||||
|
||||
const target_image_path = ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src;
|
||||
this.load_sam(target_image_path);
|
||||
|
||||
if(!this.is_layout_created) {
|
||||
// layout
|
||||
const imgCanvas = document.createElement('canvas');
|
||||
const maskCanvas = document.createElement('canvas');
|
||||
const pointsCanvas = document.createElement('canvas');
|
||||
|
||||
imgCanvas.id = "imageCanvas";
|
||||
maskCanvas.id = "maskCanvas";
|
||||
pointsCanvas.id = "pointsCanvas";
|
||||
|
||||
this.setlayout(imgCanvas, maskCanvas, pointsCanvas);
|
||||
|
||||
// prepare content
|
||||
this.imgCanvas = imgCanvas;
|
||||
this.maskCanvas = maskCanvas;
|
||||
this.pointsCanvas = pointsCanvas;
|
||||
this.maskCtx = maskCanvas.getContext('2d');
|
||||
this.pointsCtx = pointsCanvas.getContext('2d');
|
||||
|
||||
this.is_layout_created = true;
|
||||
|
||||
// replacement of onClose hook since close is not real close
|
||||
const self = this;
|
||||
const observer = new MutationObserver(function(mutations) {
|
||||
mutations.forEach(function(mutation) {
|
||||
if (mutation.type === 'attributes' && mutation.attributeName === 'style') {
|
||||
if(self.last_display_style && self.last_display_style != 'none' && self.element.style.display == 'none') {
|
||||
ComfyApp.onClipspaceEditorClosed();
|
||||
}
|
||||
|
||||
self.last_display_style = self.element.style.display;
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
const config = { attributes: true };
|
||||
observer.observe(this.element, config);
|
||||
}
|
||||
|
||||
this.setImages(target_image_path, this.imgCanvas, this.pointsCanvas);
|
||||
|
||||
if(ComfyApp.clipspace_return_node) {
|
||||
this.saveButton.innerText = "Save to node";
|
||||
}
|
||||
else {
|
||||
this.saveButton.innerText = "Save";
|
||||
}
|
||||
this.saveButton.disabled = true;
|
||||
|
||||
this.element.style.display = "block";
|
||||
this.element.style.zIndex = 8888; // NOTE: alert dialog must be high priority.
|
||||
}
|
||||
|
||||
updateBrushPreview(self, event) {
|
||||
event.preventDefault();
|
||||
|
||||
const centerX = event.pageX;
|
||||
const centerY = event.pageY;
|
||||
|
||||
const brush = self.brush;
|
||||
|
||||
brush.style.width = self.brush_size * 2 + "px";
|
||||
brush.style.height = self.brush_size * 2 + "px";
|
||||
brush.style.left = (centerX - self.brush_size) + "px";
|
||||
brush.style.top = (centerY - self.brush_size) + "px";
|
||||
}
|
||||
|
||||
setImages(target_image_path, imgCanvas, pointsCanvas) {
|
||||
const imgCtx = imgCanvas.getContext('2d');
|
||||
const maskCtx = this.maskCtx;
|
||||
const maskCanvas = this.maskCanvas;
|
||||
|
||||
const self = this;
|
||||
|
||||
// image load
|
||||
const orig_image = new Image();
|
||||
window.addEventListener("resize", () => {
|
||||
// repositioning
|
||||
imgCanvas.width = window.innerWidth - 250;
|
||||
imgCanvas.height = window.innerHeight - 200;
|
||||
|
||||
// redraw image
|
||||
let drawWidth = orig_image.width;
|
||||
let drawHeight = orig_image.height;
|
||||
|
||||
if (orig_image.width > imgCanvas.width) {
|
||||
drawWidth = imgCanvas.width;
|
||||
drawHeight = (drawWidth / orig_image.width) * orig_image.height;
|
||||
}
|
||||
|
||||
if (drawHeight > imgCanvas.height) {
|
||||
drawHeight = imgCanvas.height;
|
||||
drawWidth = (drawHeight / orig_image.height) * orig_image.width;
|
||||
}
|
||||
|
||||
imgCtx.drawImage(orig_image, 0, 0, drawWidth, drawHeight);
|
||||
|
||||
// update mask
|
||||
pointsCanvas.width = drawWidth;
|
||||
pointsCanvas.height = drawHeight;
|
||||
pointsCanvas.style.top = imgCanvas.offsetTop + "px";
|
||||
pointsCanvas.style.left = imgCanvas.offsetLeft + "px";
|
||||
|
||||
maskCanvas.width = drawWidth;
|
||||
maskCanvas.height = drawHeight;
|
||||
maskCanvas.style.top = imgCanvas.offsetTop + "px";
|
||||
maskCanvas.style.left = imgCanvas.offsetLeft + "px";
|
||||
|
||||
self.invalidateMaskCanvas(self);
|
||||
self.invalidatePointsCanvas(self);
|
||||
});
|
||||
|
||||
// original image load
|
||||
orig_image.onload = () => self.onLoaded(self);
|
||||
const rgb_url = new URL(target_image_path);
|
||||
rgb_url.searchParams.delete('channel');
|
||||
rgb_url.searchParams.set('channel', 'rgb');
|
||||
orig_image.src = rgb_url;
|
||||
self.image = orig_image;
|
||||
}
|
||||
|
||||
onLoaded(self) {
|
||||
if(self.message_box) {
|
||||
self.element.removeChild(self.message_box);
|
||||
self.message_box = null;
|
||||
}
|
||||
|
||||
window.dispatchEvent(new Event('resize'));
|
||||
|
||||
self.setEventHandler(pointsCanvas);
|
||||
self.saveButton.disabled = false;
|
||||
}
|
||||
|
||||
setEventHandler(targetCanvas) {
|
||||
targetCanvas.addEventListener("contextmenu", (event) => {
|
||||
event.preventDefault();
|
||||
});
|
||||
|
||||
const self = this;
|
||||
targetCanvas.addEventListener('pointermove', (event) => this.updateBrushPreview(self,event));
|
||||
targetCanvas.addEventListener('pointerdown', (event) => this.handlePointerDown(self,event));
|
||||
targetCanvas.addEventListener('pointerover', (event) => { this.brush.style.display = "block"; });
|
||||
targetCanvas.addEventListener('pointerleave', (event) => { this.brush.style.display = "none"; });
|
||||
document.addEventListener('keydown', ImpactSamEditorDialog.handleKeyDown);
|
||||
}
|
||||
|
||||
static handleKeyDown(event) {
|
||||
const self = ImpactSamEditorDialog.instance;
|
||||
if (event.key === '=') { // positive
|
||||
brush.style.backgroundColor = "blue";
|
||||
brush.style.outline = "2px solid pink";
|
||||
self.is_positive_mode = true;
|
||||
} else if (event.key === '-') { // negative
|
||||
brush.style.backgroundColor = "red";
|
||||
brush.style.outline = "2px solid skyblue";
|
||||
self.is_positive_mode = false;
|
||||
}
|
||||
}
|
||||
|
||||
is_positive_mode = true;
|
||||
prompt_points = [];
|
||||
fidelity = 70;
|
||||
|
||||
invalidatePointsCanvas(self) {
|
||||
const ctx = self.pointsCtx;
|
||||
|
||||
for (const i in self.prompt_points) {
|
||||
const [is_positive, x, y] = self.prompt_points[i];
|
||||
|
||||
const scaledX = x * ctx.canvas.width / self.image.width;
|
||||
const scaledY = y * ctx.canvas.height / self.image.height;
|
||||
|
||||
if(is_positive)
|
||||
ctx.fillStyle = "blue";
|
||||
else
|
||||
ctx.fillStyle = "red";
|
||||
ctx.beginPath();
|
||||
ctx.arc(scaledX, scaledY, 3, 0, 3 * Math.PI);
|
||||
ctx.fill();
|
||||
}
|
||||
}줘
|
||||
|
||||
invalidateMaskCanvas(self) {
|
||||
if(self.mask_image) {
|
||||
self.maskCtx.clearRect(0, 0, self.maskCanvas.width, self.maskCanvas.height);
|
||||
self.maskCtx.drawImage(self.mask_image, 0, 0, self.maskCanvas.width, self.maskCanvas.height);
|
||||
}
|
||||
}
|
||||
|
||||
async load_sam(url) {
|
||||
const parsedUrl = new URL(url);
|
||||
const searchParams = new URLSearchParams(parsedUrl.search);
|
||||
|
||||
const filename = searchParams.get("filename") || "";
|
||||
const fileType = searchParams.get("type") || "";
|
||||
const subfolder = searchParams.get("subfolder") || "";
|
||||
|
||||
const data = {
|
||||
sam_model_name: "sam_vit_b_01ec64.pth",
|
||||
filename: filename,
|
||||
type: fileType,
|
||||
subfolder: subfolder
|
||||
};
|
||||
|
||||
fetch('/sam/prepare', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify(data)
|
||||
});
|
||||
}
|
||||
|
||||
async detect(self) {
|
||||
const positive_points = [];
|
||||
const negative_points = [];
|
||||
|
||||
for(const i in self.prompt_points) {
|
||||
const [is_positive, x, y] = self.prompt_points[i];
|
||||
const point = [x,y];
|
||||
if(is_positive)
|
||||
positive_points.push(point);
|
||||
else
|
||||
negative_points.push(point);
|
||||
}
|
||||
|
||||
const data = {
|
||||
positive_points: positive_points,
|
||||
negative_points: negative_points,
|
||||
threshold: self.fidelity/100
|
||||
};
|
||||
|
||||
const response = await fetch('/sam/detect', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'image/png' },
|
||||
body: JSON.stringify(data)
|
||||
});
|
||||
|
||||
const blob = await response.blob();
|
||||
const url = URL.createObjectURL(blob);
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
const image = new Image();
|
||||
image.onload = () => resolve(image);
|
||||
image.onerror = reject;
|
||||
image.src = url;
|
||||
});
|
||||
}
|
||||
|
||||
handlePointerDown(self, event) {
|
||||
if ([0, 2, 5].includes(event.button)) {
|
||||
event.preventDefault();
|
||||
const x = event.offsetX || event.targetTouches[0].clientX - maskRect.left;
|
||||
const y = event.offsetY || event.targetTouches[0].clientY - maskRect.top;
|
||||
|
||||
const originalX = x * self.image.width / self.pointsCanvas.width;
|
||||
const originalY = y * self.image.height / self.pointsCanvas.height;
|
||||
|
||||
var point = null;
|
||||
if (event.button == 0) {
|
||||
// positive
|
||||
point = [true, originalX, originalY];
|
||||
} else {
|
||||
// negative
|
||||
point = [false, originalX, originalY];
|
||||
}
|
||||
|
||||
self.prompt_points.push(point);
|
||||
|
||||
self.invalidatePointsCanvas(self);
|
||||
}
|
||||
}
|
||||
|
||||
async save(self) {
|
||||
if(!self.mask_image) {
|
||||
this.close();
|
||||
return;
|
||||
}
|
||||
|
||||
const save_canvas = document.createElement('canvas');
|
||||
|
||||
const save_ctx = save_canvas.getContext('2d', {willReadFrequently:true});
|
||||
save_canvas.width = self.mask_image.width;
|
||||
save_canvas.height = self.mask_image.height;
|
||||
|
||||
save_ctx.drawImage(self.mask_image, 0, 0, save_canvas.width, save_canvas.height);
|
||||
|
||||
const save_data = save_ctx.getImageData(0, 0, save_canvas.width, save_canvas.height);
|
||||
|
||||
// refine mask image
|
||||
for (let i = 0; i < save_data.data.length; i += 4) {
|
||||
if(save_data.data[i]) {
|
||||
save_data.data[i+3] = 0;
|
||||
}
|
||||
else {
|
||||
save_data.data[i+3] = 255;
|
||||
}
|
||||
|
||||
save_data.data[i] = 0;
|
||||
save_data.data[i+1] = 0;
|
||||
save_data.data[i+2] = 0;
|
||||
}
|
||||
|
||||
save_ctx.globalCompositeOperation = 'source-over';
|
||||
save_ctx.putImageData(save_data, 0, 0);
|
||||
|
||||
const formData = new FormData();
|
||||
const filename = "clipspace-mask-" + performance.now() + ".png";
|
||||
|
||||
const item =
|
||||
{
|
||||
"filename": filename,
|
||||
"subfolder": "",
|
||||
"type": "temp",
|
||||
};
|
||||
|
||||
if(ComfyApp.clipspace.images)
|
||||
ComfyApp.clipspace.images[0] = item;
|
||||
|
||||
if(ComfyApp.clipspace.widgets) {
|
||||
const index = ComfyApp.clipspace.widgets.findIndex(obj => obj.name === 'image');
|
||||
|
||||
if(index >= 0)
|
||||
ComfyApp.clipspace.widgets[index].value = item;
|
||||
}
|
||||
|
||||
const dataURL = save_canvas.toDataURL();
|
||||
const blob = dataURLToBlob(dataURL);
|
||||
|
||||
const original_blob = loadedImageToBlob(this.image);
|
||||
|
||||
formData.append('image', blob, filename);
|
||||
formData.append('original_image', original_blob);
|
||||
formData.append('type', "temp");
|
||||
|
||||
await uploadMask(item, formData);
|
||||
ComfyApp.onClipspaceEditorSave();
|
||||
this.close();
|
||||
}
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "Comfy.Impact.SAMEditor",
|
||||
init(app) {
|
||||
const callback =
|
||||
function () {
|
||||
let dlg = ImpactSamEditorDialog.getInstance();
|
||||
dlg.show();
|
||||
};
|
||||
|
||||
const context_predicate = () => ComfyApp.clipspace && ComfyApp.clipspace.imgs && ComfyApp.clipspace.imgs.length > 0
|
||||
ClipspaceDialog.registerButton("Impact SAM Detector", context_predicate, callback);
|
||||
},
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
if (nodeData.output.includes("MASK") && nodeData.output.includes("IMAGE")) {
|
||||
addMenuHandler(nodeType, function (_, options) {
|
||||
options.unshift({
|
||||
content: "Open in SAM Detector",
|
||||
callback: () => {
|
||||
ComfyApp.copyToClipspace(this);
|
||||
ComfyApp.clipspace_return_node = this;
|
||||
|
||||
let dlg = ImpactSamEditorDialog.getInstance();
|
||||
dlg.show();
|
||||
},
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
main_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
|
||||
sys.path.append(os.path.dirname(__file__))
|
||||
sys.path.append(main_dir)
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user