From 06a905414e0f7a6f324815c00ece9bb8f525ee46 Mon Sep 17 00:00:00 2001
From: Kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 6 May 2024 16:55:52 +0300
Subject: [PATCH] Add fit_in_frame option for InstanceDiffusion bboxes

---
 nodes/curve_nodes.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/nodes/curve_nodes.py b/nodes/curve_nodes.py
index 00c081c..12e42e8 100644
--- a/nodes/curve_nodes.py
+++ b/nodes/curve_nodes.py
@@ -718,10 +718,11 @@ for example:
         },
         "optional": {
             "size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True}),
+            "fit_in_frame": ("BOOLEAN", {"default": True}),
         }
     } 
 
-    def tracking(self, coordinates, class_name, class_id, width, height, bbox_width, bbox_height, prompt, size_multiplier=[1.0]):
+    def tracking(self, coordinates, class_name, class_id, width, height, bbox_width, bbox_height, prompt, size_multiplier=[1.0], fit_in_frame=True):
         # Define the number of images in the batch
         coordinates = coordinates.replace("'", '"')
         coordinates = json.loads(coordinates)
@@ -744,9 +745,16 @@ for example:
             bottom_right_x = x + adjusted_bbox_width // 2
             bottom_right_y = y + adjusted_bbox_height // 2
 
+            if fit_in_frame:
+                # Clip the coordinates to the frame boundaries
+                top_left_x = max(0, top_left_x)
+                top_left_y = max(0, top_left_y)
+                bottom_right_x = min(width, bottom_right_x)
+                bottom_right_y = min(height, bottom_right_y)
+
             # Append the top left and bottom right coordinates to the list for the current ID
             id_coordinates.append([top_left_x, top_left_y, bottom_right_x, bottom_right_y, width, height])
-          
+        
         class_id = int(class_id)
         # Assign the list of coordinates to the specified ID within the class_id dictionary
         tracked[class_name][class_id] = id_coordinates
@@ -897,13 +905,11 @@ CreateInstanceDiffusionTracking -node.
         
         colormap = cm.get_cmap('rainbow', len(tracking))
         if draw_text:
-            #font = ImageFont.load_default()
             font = ImageFont.truetype(font, font_size)
 
         # Iterate over each image in the batch
         for i in range(image.shape[0]):
             # Extract the current image and convert it to a PIL image
-            # Adjust the tensor to (C, H, W) for ToPILImage
             current_image = image[i, :, :, :].permute(2, 0, 1)
             pil_image = transforms.ToPILImage()(current_image)