From fe69b9c4a8428cc6a957f2b40c5aa559dd25416b Mon Sep 17 00:00:00 2001
From: vahidrezanezhad <vahid631983@gmail.com>
Date: Fri, 21 Jun 2024 23:42:25 +0200
Subject: [PATCH] update inference

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

diff --git a/train/inference.py b/train/inference.py
index c7a8b02..3fec9c2 100644
--- a/train/inference.py
+++ b/train/inference.py
@@ -557,6 +557,10 @@ class sbb_predict:
         res=self.predict()
         if (self.task == 'classification' or self.task == 'reading_order'):
             pass
+        elif self.task == 'enhancement':
+            if self.save:
+                print(self.save)
+                cv2.imwrite(self.save,res)
         else:
             img_seg_overlayed = self.visualize_model_output(res, self.img_org, self.task)
             if self.save: