fix(nodes): handle no detected line segments

invoke-ai · Sep 11, 2024 · b422968 · b422968
1 parent 66489b7
commit b422968
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/invokeai/backend/image_util/mlsd/utils.py b/invokeai/backend/image_util/mlsd/utils.py
@@ -79,11 +79,15 @@ def pred_lines(image, model,
             y_end = y + disp_y_end
             segments_list.append([x_start, y_start, x_end, y_end])
 
-    lines = 2 * np.array(segments_list)  # 256 > 512
-    lines[:, 0] = lines[:, 0] * w_ratio
-    lines[:, 1] = lines[:, 1] * h_ratio
-    lines[:, 2] = lines[:, 2] * w_ratio
-    lines[:, 3] = lines[:, 3] * h_ratio
+    if segments_list:
+        lines = 2 * np.array(segments_list)  # 256 > 512
+        lines[:, 0] = lines[:, 0] * w_ratio
+        lines[:, 1] = lines[:, 1] * h_ratio
+        lines[:, 2] = lines[:, 2] * w_ratio
+        lines[:, 3] = lines[:, 3] * h_ratio
+    else:
+        # No segments detected - return empty array
+        lines = np.array([])
 
     return lines