Update inference.py to check if save_layout was passed as argument otherwise can give an cv2 error

2 months ago · df4a47ae6f
parent cca4d17823
commit df4a47ae6f
1 changed files with 1 additions and 0 deletions
--- a/inference.py
+++ b/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: