mirror of
https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git
synced 2025-12-09 21:04:41 +08:00
bbox to point editor
This commit is contained in:
kijai
parent
5ccd427bd8
commit
dee4e8f1eb
@ -710,28 +710,48 @@ Visualizes the specified bbox on the image.
|
||||
image_list = []
|
||||
for image, bbox in zip(images, bboxes):
|
||||
x_min, y_min, width, height = bbox
|
||||
|
||||
# Ensure bbox coordinates are integers
|
||||
x_min = int(x_min)
|
||||
y_min = int(y_min)
|
||||
width = int(width)
|
||||
height = int(height)
|
||||
|
||||
# Permute the image dimensions
|
||||
image = image.permute(2, 0, 1)
|
||||
|
||||
# Clone the image to draw bounding boxes
|
||||
img_with_bbox = image.clone()
|
||||
|
||||
# Define the color for the bbox, e.g., red
|
||||
color = torch.tensor([1, 0, 0], dtype=torch.float32)
|
||||
|
||||
# Ensure color tensor matches the image channels
|
||||
if color.shape[0] != img_with_bbox.shape[0]:
|
||||
color = color.unsqueeze(1).expand(-1, line_width)
|
||||
|
||||
# Draw lines for each side of the bbox with the specified line width
|
||||
for lw in range(line_width):
|
||||
# Top horizontal line
|
||||
img_with_bbox[:, y_min + lw, x_min:x_min + width] = color[:, None]
|
||||
if y_min + lw < img_with_bbox.shape[1]:
|
||||
img_with_bbox[:, y_min + lw, x_min:x_min + width] = color[:, None]
|
||||
|
||||
# Bottom horizontal line
|
||||
img_with_bbox[:, y_min + height - lw, x_min:x_min + width] = color[:, None]
|
||||
if y_min + height - lw < img_with_bbox.shape[1]:
|
||||
img_with_bbox[:, y_min + height - lw, x_min:x_min + width] = color[:, None]
|
||||
|
||||
# Left vertical line
|
||||
img_with_bbox[:, y_min:y_min + height, x_min + lw] = color[:, None]
|
||||
if x_min + lw < img_with_bbox.shape[2]:
|
||||
img_with_bbox[:, y_min:y_min + height, x_min + lw] = color[:, None]
|
||||
|
||||
# Right vertical line
|
||||
img_with_bbox[:, y_min:y_min + height, x_min + width - lw] = color[:, None]
|
||||
if x_min + width - lw < img_with_bbox.shape[2]:
|
||||
img_with_bbox[:, y_min:y_min + height, x_min + width - lw] = color[:, None]
|
||||
|
||||
# Permute the image dimensions back
|
||||
img_with_bbox = img_with_bbox.permute(1, 2, 0).unsqueeze(0)
|
||||
image_list.append(img_with_bbox)
|
||||
|
||||
return (torch.cat(image_list, dim=0),)
|
||||
|
||||
return (torch.cat(image_list, dim=0),)
|
||||
@ -1254,13 +1254,15 @@ class PointsEditor:
|
||||
"required": {
|
||||
"points_store": ("STRING", {"multiline": False}),
|
||||
"coordinates": ("STRING", {"multiline": False}),
|
||||
"bbox_store": ("STRING", {"multiline": False}),
|
||||
"bboxes": ("STRING", {"multiline": False}),
|
||||
"width": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}),
|
||||
"height": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING",)
|
||||
RETURN_NAMES = ("coord_str", "normalized_str",)
|
||||
RETURN_TYPES = ("STRING", "STRING", "BBOX", "MASK")
|
||||
RETURN_NAMES = ("coord_str", "normalized_str", "bbox", "bbox_mask")
|
||||
FUNCTION = "splinedata"
|
||||
CATEGORY = "KJNodes/weights"
|
||||
DESCRIPTION = """
|
||||
@ -1293,7 +1295,7 @@ output types:
|
||||
example compatible nodes: unknown
|
||||
"""
|
||||
|
||||
def splinedata(self, points_store, width, height, coordinates):
|
||||
def splinedata(self, points_store, bbox_store, width, height, coordinates, bboxes):
|
||||
|
||||
coordinates = json.loads(coordinates)
|
||||
normalized = []
|
||||
@ -1305,6 +1307,23 @@ output types:
|
||||
norm_y = (1.0 - (coord['y'] / height) - 0.0)
|
||||
normalized_y_values.append(norm_y)
|
||||
normalized.append({'x':norm_x, 'y':norm_y})
|
||||
|
||||
|
||||
return (json.dumps(coordinates), json.dumps(normalized))
|
||||
|
||||
bboxes = json.loads(bboxes)
|
||||
bboxes = [(int(bboxes["x"]), int(bboxes["y"]), int(bboxes["width"]), int(bboxes["height"]))]
|
||||
|
||||
# Create a blank mask
|
||||
mask = np.zeros((height, width), dtype=np.uint8)
|
||||
|
||||
# Draw the bounding box on the mask
|
||||
for bbox in bboxes:
|
||||
x_min, y_min, w, h = bbox
|
||||
x_max = x_min + w
|
||||
y_max = y_min + h
|
||||
mask[y_min:y_max, x_min:x_max] = 1 # Fill the bounding box area with 1s
|
||||
|
||||
mask_tensor = torch.from_numpy(mask)
|
||||
mask_tensor = mask_tensor.unsqueeze(0).float().cpu()
|
||||
#mask_tensor = mask_tensor[:,:,0]
|
||||
print(mask_tensor.shape)
|
||||
|
||||
return (json.dumps(coordinates), json.dumps(normalized), bboxes, mask_tensor)
|
||||
@ -106,6 +106,7 @@ app.registerExtension({
|
||||
chainCallback(nodeType.prototype, "onNodeCreated", function () {
|
||||
|
||||
hideWidgetForGood(this, this.widgets.find(w => w.name === "coordinates"))
|
||||
hideWidgetForGood(this, this.widgets.find(w => w.name === "bboxes"))
|
||||
|
||||
var element = document.createElement("div");
|
||||
this.uuid = makeUUID()
|
||||
@ -277,7 +278,7 @@ function createPointsEditor(context, reset=false) {
|
||||
context.setSize([context.size[0], h + 230]);
|
||||
vis.width(w);
|
||||
vis.height(h);
|
||||
updatePath();
|
||||
updateData();
|
||||
backgroundImage.url(imageUrl).visible(true).root.render();
|
||||
};
|
||||
}
|
||||
@ -319,7 +320,7 @@ function createPointsEditor(context, reset=false) {
|
||||
context.setSize([context.size[0], this.height + 230]);
|
||||
vis.width(w);
|
||||
vis.height(h);
|
||||
updatePath();
|
||||
updateData();
|
||||
backgroundImage.url(imageUrl).visible(true).root.render();
|
||||
};
|
||||
}
|
||||
@ -341,33 +342,23 @@ function createPointsEditor(context, reset=false) {
|
||||
}
|
||||
|
||||
createContextMenu();
|
||||
function updatePath() {
|
||||
function updateData() {
|
||||
if (points.length == 0) {
|
||||
console.log("no points")
|
||||
return
|
||||
}
|
||||
let coords = points
|
||||
|
||||
// if (pointsLayer) {
|
||||
// // Update the data of the existing points layer
|
||||
// pointsLayer.data(coords);
|
||||
// } else {
|
||||
// // Create the points layer if it doesn't exist
|
||||
// pointsLayer = vis.add(pv.Dot)
|
||||
// .data(coords)
|
||||
// .left(function(d) { return d.x; })
|
||||
// .top(function(d) { return d.y; })
|
||||
// .radius(3) // Adjust the radius as needed
|
||||
// .fillStyle("red") // Change the color as needed
|
||||
// .strokeStyle("black") // Change the stroke color as needed
|
||||
// .lineWidth(1); // Adjust the line width as needed
|
||||
// }
|
||||
|
||||
let coordsString = JSON.stringify(coords);
|
||||
let bbox = calculateBBox(box_startX, box_startY, box_endX, box_endY)
|
||||
let bboxString = JSON.stringify(bbox);
|
||||
pointsStoreWidget.value = JSON.stringify(points);
|
||||
bboxStoreWidget.value = JSON.stringify(bboxString);
|
||||
if (coordWidget) {
|
||||
coordWidget.value = coordsString;
|
||||
}
|
||||
if (bboxWidget) {
|
||||
bboxWidget.value = bboxString;
|
||||
}
|
||||
vis.render();
|
||||
}
|
||||
|
||||
@ -378,8 +369,8 @@ function createPointsEditor(context, reset=false) {
|
||||
const pointsStoreWidget = context.widgets.find(w => w.name === "points_store");
|
||||
const widthWidget = context.widgets.find(w => w.name === "width");
|
||||
const heightWidget = context.widgets.find(w => w.name === "height");
|
||||
|
||||
var pointsLayer = null;
|
||||
const bboxStoreWidget = context.widgets.find(w => w.name === "bbox_store");
|
||||
const bboxWidget = context.widgets.find(w => w.name === "bboxes");
|
||||
|
||||
widthWidget.callback = () => {
|
||||
w = widthWidget.value;
|
||||
@ -387,17 +378,17 @@ function createPointsEditor(context, reset=false) {
|
||||
context.setSize([w + 45, context.size[1]]);
|
||||
}
|
||||
vis.width(w);
|
||||
updatePath();
|
||||
updateData();
|
||||
}
|
||||
heightWidget.callback = () => {
|
||||
h = heightWidget.value
|
||||
vis.height(h)
|
||||
context.setSize([context.size[0], h + 430]);
|
||||
updatePath();
|
||||
updateData();
|
||||
}
|
||||
pointsStoreWidget.callback = () => {
|
||||
points = JSON.parse(pointsStoreWidget.value);
|
||||
updatePath();
|
||||
updateData();
|
||||
}
|
||||
|
||||
// Initialize or reset points array
|
||||
@ -407,6 +398,9 @@ function createPointsEditor(context, reset=false) {
|
||||
var h = heightWidget.value;
|
||||
var i = 3;
|
||||
let points = [];
|
||||
let bbox = [];
|
||||
var box_startX, box_startY, box_endX, box_endY;
|
||||
var drawing = false;
|
||||
|
||||
if (!reset && pointsStoreWidget.value != "") {
|
||||
points = JSON.parse(pointsStoreWidget.value);
|
||||
@ -420,7 +414,7 @@ function createPointsEditor(context, reset=false) {
|
||||
];
|
||||
pointsStoreWidget.value = JSON.stringify(points);
|
||||
}
|
||||
|
||||
|
||||
var vis = new pv.Panel()
|
||||
.width(w)
|
||||
.height(h)
|
||||
@ -436,11 +430,14 @@ function createPointsEditor(context, reset=false) {
|
||||
y: this.mouse().y / app.canvas.ds.scale
|
||||
};
|
||||
i = points.push(scaledMouse) - 1;
|
||||
updatePath();
|
||||
updateData();
|
||||
return this;
|
||||
}
|
||||
else if (pv.event.ctrlKey) {
|
||||
|
||||
console.log("start drawing at " + this.mouse().x / app.canvas.ds.scale + ", " + this.mouse().y / app.canvas.ds.scale);
|
||||
drawing = true;
|
||||
box_startX = this.mouse().x / app.canvas.ds.scale;
|
||||
box_startY = this.mouse().y / app.canvas.ds.scale;
|
||||
}
|
||||
else if (pv.event.button === 2) {
|
||||
context.contextMenu.style.display = 'block';
|
||||
@ -448,8 +445,37 @@ function createPointsEditor(context, reset=false) {
|
||||
context.contextMenu.style.top = `${pv.event.clientY}px`;
|
||||
}
|
||||
})
|
||||
var backgroundImage = vis.add(pv.Image)
|
||||
.visible(false)
|
||||
.event("mousemove", function() {
|
||||
if (drawing) {
|
||||
box_endX = this.mouse().x / app.canvas.ds.scale;
|
||||
box_endY = this.mouse().y / app.canvas.ds.scale;
|
||||
vis.render();
|
||||
}
|
||||
})
|
||||
.event("mouseup", function() {
|
||||
console.log("end drawing at " + this.mouse().x / app.canvas.ds.scale + ", " + this.mouse().y / app.canvas.ds.scale);
|
||||
drawing = false;
|
||||
updateData();
|
||||
});
|
||||
|
||||
var backgroundImage = vis.add(pv.Image)
|
||||
.visible(false)
|
||||
|
||||
vis.add(pv.Area)
|
||||
.data(function() {
|
||||
return drawing || bbox ? [box_startX, box_endX] : [];
|
||||
})
|
||||
.bottom(function() {
|
||||
return h - Math.max(box_startY, box_endY);
|
||||
})
|
||||
.left(function(d) {
|
||||
return d;
|
||||
})
|
||||
.height(function() {
|
||||
return Math.abs(box_startY - box_endY);
|
||||
})
|
||||
.fillStyle("rgba(70, 130, 180, 0.5)")
|
||||
.strokeStyle("steelblue");
|
||||
|
||||
vis.add(pv.Dot)
|
||||
.data(() => points)
|
||||
@ -475,7 +501,7 @@ function createPointsEditor(context, reset=false) {
|
||||
points.splice(i--, 1);
|
||||
|
||||
}
|
||||
updatePath();
|
||||
updateData();
|
||||
isDragging = false;
|
||||
|
||||
})
|
||||
@ -539,7 +565,7 @@ function createPointsEditor(context, reset=false) {
|
||||
context.setSize([w + 45, context.size[1]]);
|
||||
}
|
||||
context.setSize([context.size[0], h + 430]);
|
||||
updatePath();
|
||||
updateData();
|
||||
|
||||
if (context.properties.imgData && context.properties.imgData.base64) {
|
||||
const base64String = context.properties.imgData.base64;
|
||||
@ -572,4 +598,12 @@ export function hideWidgetForGood(node, widget, suffix = '') {
|
||||
hideWidgetForGood(node, w, ':' + widget.name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function calculateBBox(x1, y1, x2, y2) {
|
||||
var x = Math.min(x1, x2);
|
||||
var y = Math.min(y1, y2);
|
||||
var width = Math.abs(x2 - x1);
|
||||
var height = Math.abs(y2 - y1);
|
||||
return { x: x, y: y, width: width, height: height };
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user