From a524f8b1a7e5e68219cdcb12e239bc6ae8a1391c Mon Sep 17 00:00:00 2001
From: johnlockejrr <johnlockejrr@gmail.com>
Date: Sat, 19 Oct 2024 13:21:29 -0700
Subject: [PATCH] Update inference.py to check if save_layout was passed as
 argument otherwise can give an cv2 error

---
 train/inference.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/train/inference.py b/train/inference.py
index 8d0a572..89d32de 100644
--- a/train/inference.py
+++ b/train/inference.py
@@ -566,6 +566,7 @@ class sbb_predict:
             img_seg_overlayed, only_layout  = self.visualize_model_output(res, self.img_org, self.task)
             if self.save:
                 cv2.imwrite(self.save,img_seg_overlayed)
+            if self.save_layout:
                 cv2.imwrite(self.save_layout, only_layout)
                 
         if self.ground_truth: