fix: Fixed compatibility with FP16

mindee · Nov 10, 2021 · a4de6c4 · a4de6c4
1 parent 1af2e73
commit a4de6c4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/doctr/models/detection/core.py b/doctr/models/detection/core.py
@@ -98,7 +98,7 @@ def __call__(
 
         for p_, bitmap_ in zip(proba_map, bitmap):
             # Perform opening (erosion + dilatation)
-            bitmap_ = cv2.morphologyEx(bitmap_, cv2.MORPH_OPEN, kernel)
+            bitmap_ = cv2.morphologyEx(bitmap_.astype(np.float32), cv2.MORPH_OPEN, kernel)
             # Rotate bitmap and proba_map
             angle = get_bitmap_angle(bitmap_)
             angles_batch.append(angle)

diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py
@@ -208,7 +208,7 @@ def rotate_image(
 
     height, width = exp_img.shape[:2]
     rot_mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
-    rot_img = cv2.warpAffine(exp_img, rot_mat, (width, height))
+    rot_img = cv2.warpAffine(exp_img.astype(np.float32), rot_mat, (width, height))
     if expand:
         # Pad to get the same aspect ratio
         if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):