fix: debug 4

mindee · charlesmindee · Dec 26, 2021 · Dec 16, 2021 · Dec 17, 2021 · Dec 17, 2021
commit 183ce86c9c0517210eff0453643cc0db3b9a7fe4
diff --git a/doctr/models/_utils.py b/doctr/models/_utils.py
@@ -79,12 +79,13 @@ def extract_rcrops(
     for box in _boxes:
         src_pts = box[1:, :].astype(np.float32)
         # Preserve size
-        _, (w, h), _ = cv2.minAreaRect(box)
-        dst_pts = np.array([[0, 0], [w - 1, 0], [w - 1, h - 1]], dtype=dtype)
+        d1 = np.linalg.norm(src_pts[0, :] - src_pts[1, :], axis=-1)
+        d2 = np.linalg.norm(src_pts[1, :] - src_pts[2, :], axis=-1)
+        dst_pts = np.array([[0, 0], [d1 - 1, 0], [d1 - 1, d2 - 1]], dtype=dtype)
         # The transformation matrix
         M = cv2.getAffineTransform(src_pts, dst_pts)
         # Warp the rotated rectangle
-        crop = cv2.warpAffine(img if channels_last else img.transpose(1, 2, 0), M, (int(w), int(h)))
+        crop = cv2.warpAffine(img if channels_last else img.transpose(1, 2, 0), M, (int(d1), int(d2)))
         crops.append(crop)
 
     return crops

diff --git a/doctr/models/detection/differentiable_binarization/base.py b/doctr/models/detection/differentiable_binarization/base.py
@@ -133,7 +133,7 @@ def bitmap_to_boxes(
             if self.assume_straight_pages:
                 if _box is None or _box[2] < min_size_box or _box[3] < min_size_box:
                     continue
-            elif abs(_box[0, 0] - _box[2, 0]) < min_size_box or abs(_box[0, 1] - _box[2, 1]) < min_size_box:
+            elif np.linalg.norm(_box[2, :] - _box[0, :], axis=-1) < min_size_box:
                 continue
 
             if self.assume_straight_pages:

diff --git a/doctr/models/detection/differentiable_binarization/pytorch.py b/doctr/models/detection/differentiable_binarization/pytorch.py
@@ -277,8 +277,10 @@ def _dbnet(arch: str, pretrained: bool, pretrained_backbone: bool = True, **kwar
     model = DBNet(feat_extractor, cfg=default_cfgs[arch], **kwargs)
     # Load pretrained parameters
     if pretrained:
-        load_pretrained_params(model, default_cfgs[arch]['url'])
-
+        state_dict = torch.load("/home/laptopmindee/Téléchargements/db_resnet50_20211130-090054.pt", map_location='cpu')
+        # Load weights
+        model.load_state_dict(state_dict)
+
     return model
 
 

diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py
@@ -15,6 +15,7 @@
 from doctr.models.detection.predictor import DetectionPredictor
 from doctr.models.recognition.predictor import RecognitionPredictor
 from doctr.utils.geometry import rotate_boxes, rotate_image
+from ..classification import crop_orientation_predictor
 
 from .base import _OCRPredictor
 
@@ -52,6 +53,7 @@ def __init__(
         self.doc_builder = DocumentBuilder(export_as_straight_boxes=export_as_straight_boxes)
         self.assume_straight_pages = assume_straight_pages
         self.straighten_pages = straighten_pages
+        self.crop_orientation_predictor = crop_orientation_predictor(pretrained=True)
 
     @torch.no_grad()
     def forward(

diff --git a/doctr/models/predictor/tensorflow.py b/doctr/models/predictor/tensorflow.py
@@ -15,6 +15,7 @@
 from doctr.models.recognition.predictor import RecognitionPredictor
 from doctr.utils.geometry import rotate_boxes, rotate_image
 from doctr.utils.repr import NestedObject
+from ..classification import crop_orientation_predictor
 
 from .base import _OCRPredictor
 
@@ -52,6 +53,7 @@ def __init__(
         self.doc_builder = DocumentBuilder(export_as_straight_boxes=export_as_straight_boxes)
         self.assume_straight_pages = assume_straight_pages
         self.straighten_pages = straighten_pages
+        self.crop_orientation_predictor = crop_orientation_predictor(pretrained=True)
 
     def __call__(
         self,

diff --git a/doctr/models/zoo.py b/doctr/models/zoo.py
@@ -27,7 +27,7 @@ def _predictor(
         det_arch,
         pretrained=pretrained,
         batch_size=det_bs,
-        assume_straight_pages=assume_straight_pages
+        assume_straight_pages=assume_straight_pages,
     )
 
     # Recognition

diff --git a/doctr/utils/visualization.py b/doctr/utils/visualization.py
@@ -126,9 +126,8 @@ def create_obj_patch(
     if isinstance(geometry, tuple):
         if len(geometry) == 2:
             return rect_patch(geometry, page_dimensions, **kwargs)  # type: ignore[arg-type]
-    elif isinstance(geometry, np.ndarray):
-        return polygon_patch(geometry, page_dimensions, **kwargs)  # type: ignore[arg-type]
-
+        elif len(geometry) == 4:
+            return polygon_patch(np.asarray(geometry), page_dimensions, **kwargs)  # type: ignore[arg-type]
     raise ValueError("invalid geometry format")