From dee4e8f1eb86809b026dd3d168184fe0e13e1008 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 2 Aug 2024 18:52:05 +0300 Subject: [PATCH] bbox to point editor --- nodes/batchcrop_nodes.py | 28 ++++++++++-- nodes/curve_nodes.py | 31 ++++++++++--- web/js/point_editor.js | 96 +++++++++++++++++++++++++++------------- 3 files changed, 114 insertions(+), 41 deletions(-) diff --git a/nodes/batchcrop_nodes.py b/nodes/batchcrop_nodes.py index 282dd93..61e7446 100644 --- a/nodes/batchcrop_nodes.py +++ b/nodes/batchcrop_nodes.py @@ -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),) \ No newline at end of file diff --git a/nodes/curve_nodes.py b/nodes/curve_nodes.py index 8cd2243..b4fecde 100644 --- a/nodes/curve_nodes.py +++ b/nodes/curve_nodes.py @@ -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)) \ No newline at end of file + + 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) \ No newline at end of file diff --git a/web/js/point_editor.js b/web/js/point_editor.js index 37e23a2..6c39652 100644 --- a/web/js/point_editor.js +++ b/web/js/point_editor.js @@ -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 }; } \ No newline at end of file