mypy

onnxtr/models/builder.py

@@ -74,7 +74,7 @@ def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: list[int]) -> list[li
         """
         lines = []
         # Sort words horizontally
-        word_idcs = [word_idcs[idx] for idx in boxes[word_idcs, 0].argsort().tolist()]
+        word_idcs = [word_idcs[idx] for idx in boxes[word_idcs, 0].argsort().tolist()]  # type: ignore[call-overload]
 
         # Eventually split line horizontally
         if len(word_idcs) < 2:

onnxtr/models/recognition/models/master.py

@@ -106,7 +106,7 @@ def __call__(self, logits: np.ndarray) -> list[tuple[str, float]]:
             "".join(self._embedding[idx] for idx in encoded_seq).split("<eos>")[0] for encoded_seq in out_idxs
         ]
 
-        return list(zip(word_values, np.clip(probs, 0, 1).astype(float).tolist()))
+        return list(zip(word_values, np.clip(probs, 0, 1).astype(float).tolist()))  # type: ignore[arg-type]
 
 
 def _master(

onnxtr/utils/visualization.py

@@ -274,9 +274,13 @@ def draw_boxes(boxes: np.ndarray, image: np.ndarray, color: tuple[int, int, int]
     _boxes[:, [1, 3]] *= h
     _boxes = _boxes.astype(np.int32)
     for box in _boxes.tolist():
-        xmin, ymin, xmax, ymax = box
+        xmin, ymin, xmax, ymax = box  # type: ignore[misc]
         image = cv2.rectangle(
-            image, (xmin, ymin), (xmax, ymax), color=color if isinstance(color, tuple) else (0, 0, 255), thickness=2
+            image,
+            (xmin, ymin),
+            (xmax, ymax),
+            color=color if isinstance(color, tuple) else (0, 0, 255),
+            thickness=2,  # type: ignore[arg-type]
         )
     plt.imshow(image)
     plt.plot(**kwargs)