.

README.md

@@ -4,7 +4,7 @@ emoji: 📊
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 5.0.2
+sdk_version: 4.33
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -59,7 +59,7 @@ def reset_image(image, state):
     return img, state
 
 def auto_point_detect(image):
-    out_image, msk_np = docscanner_rec(image, docscanner)
+    out_image, msk_np = docscanner_rec(image, docscanner, cuda)
     state = list(get_corner(mask2point(mask=msk_np)))
     
     img = Image.fromarray(image)
@@ -126,7 +126,7 @@ def sort_corners(corners):
 def convert(image, state):
     h,w = image.shape[:2]
     if len(state) < 4:
-        out_image, msk_np = docscanner_rec(image, docscanner)
+        out_image, msk_np = docscanner_rec(image, docscanner, cuda)
         out_image = out_image[:,:,::-1]
     elif len(state) ==4:
         state = list(sort_corners(state))

rect_main.py

@@ -52,11 +52,11 @@ def preprocess_image(img, target_size=[288, 288]):
     return im_ori, im, h_, w_
 
 
-def geotrp_rec(img, model):
+def geotrp_rec(img, model, cuda):
     im_ori, im, h_, w_ = preprocess_image(img)
 
     with torch.no_grad():
-        bm = model(im.cuda())
+        bm = model(im.to(cuda))
         bm = bm.cpu().numpy()[0]
         bm0 = bm[0, :, :]
         bm1 = bm[1, :, :]
@@ -69,11 +69,11 @@ def geotrp_rec(img, model):
         return img_geo
 
 
-def docscanner_get_mask(img, model):
+def docscanner_get_mask(img, model, cuda):
     _, im, h, w = preprocess_image(img)
 
     with torch.no_grad():
-        _, msk = model(im.cuda())
+        _, msk = model(im.to(cuda))
         msk = msk.cpu()
 
         mask_np = (msk[0, 0].numpy() * 255).astype(np.uint8)
@@ -82,11 +82,11 @@ def docscanner_get_mask(img, model):
     return mask_resized
 
 
-def docscanner_rec_img(img, model):
+def docscanner_rec_img(img, model, cuda):
     im_ori, im, h, w = preprocess_image(img)
 
     with torch.no_grad():
-        bm = model(im.cuda())
+        bm = model(im.to(cuda))
         bm = bm.cpu()
 
         # save rectified image
@@ -106,7 +106,7 @@ def docscanner_rec_img(img, model):
 
 
 
-def docscanner_rec(img, model):
+def docscanner_rec(img, model, cuda):
     im_ori = img[:, :, :3] / 255.0
     h, w, _ = im_ori.shape
     im = cv2.resize(im_ori, (288, 288))
@@ -114,7 +114,7 @@ def docscanner_rec(img, model):
     im = torch.from_numpy(im).float().unsqueeze(0)
 
     with torch.no_grad():
-        bm, msk = model(im.cuda())
+        bm, msk = model(im.to(cuda))
         bm = bm.cpu()
         msk = msk.cpu()
 
@@ -165,7 +165,7 @@ def main():
     )
     doctr = load_geotrp_model(cuda, path=config.get_geotr_model_path)
 
-    mask = docscanner_get_mask(img, docscanner)
+    mask = docscanner_get_mask(img, docscanner, cuda)
     mask_dict.add(get_mask_white_area(mask))