From 5a0b6cc72f31b8adde930c2ab67c3b3b90ae4db0 Mon Sep 17 00:00:00 2001
From: liuyebo <1515783401@qq.com>
Date: Thu, 29 Aug 2024 11:02:11 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=8E=BB=E6=89=AD=E6=9B=B2?=
 =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E7=9A=84=E5=A3=B0=E6=98=8E?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 doc_dewarp/__init__.py | 10 ++++++++++
 doc_dewarp/dewarp.py   | 12 ++----------
 docker-compose.yml     |  2 +-
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/doc_dewarp/__init__.py b/doc_dewarp/__init__.py
index e69de29..3f477b4 100644
--- a/doc_dewarp/__init__.py
+++ b/doc_dewarp/__init__.py
@@ -0,0 +1,10 @@
+import paddle
+
+from doc_dewarp.GeoTr import GeoTr
+
+model_path = "model/dewarp_model/best.ckpt"
+checkpoint = paddle.load(model_path)
+state_dict = checkpoint["model"]
+DEWARP = GeoTr()
+DEWARP.set_state_dict(state_dict)
+DEWARP.eval()
diff --git a/doc_dewarp/dewarp.py b/doc_dewarp/dewarp.py
index 1decad9..34548bd 100644
--- a/doc_dewarp/dewarp.py
+++ b/doc_dewarp/dewarp.py
@@ -1,23 +1,15 @@
 import cv2
 import paddle
 
-from .GeoTr import GeoTr
+from . import DEWARP
 from .utils import to_tensor, to_image
 
 
 def dewarp_image(image):
-    model_path = "model/dewarp_model/best.ckpt"
-
-    checkpoint = paddle.load(model_path)
-    state_dict = checkpoint["model"]
-    model = GeoTr()
-    model.set_state_dict(state_dict)
-    model.eval()
-
     img = cv2.resize(image, (288, 288))
     x = to_tensor(img)
     y = to_tensor(image)
-    bm = model(x)
+    bm = DEWARP(x)
     bm = paddle.nn.functional.interpolate(
         bm, y.shape[2:], mode="bilinear", align_corners=False
     )
diff --git a/docker-compose.yml b/docker-compose.yml
index 7f26ee8..d3599f0 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,6 +1,6 @@
 x-env:
   &template
-  image: fcb_photo_review:1.13.9
+  image: fcb_photo_review:1.13.10
   restart: always
 
 services: