fix: linknet hyperparameters postprocessing + demo for rotation model (…

…#865)

* fix: linknet parameters

* feat: add demo rotation

* feat: add rotation in demo

demo/app.py

@@ -21,7 +21,7 @@
 from doctr.models import ocr_predictor
 from doctr.utils.visualization import visualize_page
 
-DET_ARCHS = ["db_resnet50", "db_mobilenet_v3_large"]
+DET_ARCHS = ["db_resnet50", "db_mobilenet_v3_large", "linknet_resnet18_rotation"]
 RECO_ARCHS = ["crnn_vgg16_bn", "crnn_mobilenet_v3_small", "master", "sar_resnet31"]
 
 
@@ -73,7 +73,10 @@ def main():
 
  else:
  with st.spinner('Loading model...'):
- predictor = ocr_predictor(det_arch, reco_arch, pretrained=True)
+ predictor = ocr_predictor(
+ det_arch, reco_arch, pretrained=True,
+ assume_straight_pages=(det_arch != "linknet_resnet18_rotation")
+ )
 
  with st.spinner('Analyzing...'):
 
@@ -97,8 +100,9 @@ def main():
 
  # Page reconsitution under input page
  page_export = out.pages[0].export()
- img = out.pages[0].synthesize()
- cols[3].image(img, clamp=True)
+ if det_arch != "linknet_resnet18_rotation":
+ img = out.pages[0].synthesize()
+ cols[3].image(img, clamp=True)
 
  # Display JSON
  st.markdown("\nHere are your analysis results in JSON format:")

doctr/models/detection/linknet/base.py

@@ -30,7 +30,7 @@ class LinkNetPostProcessor(DetectionPostProcessor):
  """
  def __init__(
  self,
- bin_thresh: float = 0.5,
+ bin_thresh: float = 0.1,
  box_thresh: float = 0.1,
  assume_straight_pages: bool = True,
  ) -> None:
@@ -39,7 +39,7 @@ def __init__(
  bin_thresh,
  assume_straight_pages
  )
- self.unclip_ratio = 1.5
+ self.unclip_ratio = 1.2
 
  def polygon_to_box(
  self,
@@ -103,13 +103,12 @@ def bitmap_to_boxes(
  containing x, y, w, h, alpha, score for the box
  """
  height, width = bitmap.shape[:2]
- min_size_box = 1 + int(height / 512)
  boxes = []
  # get contours from connected components on the bitmap
  contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  for contour in contours:
  # Check whether smallest enclosing bounding box is not too small
- if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < min_size_box):
+ if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
  continue
  # Compute objectness
  if self.assume_straight_pages:

doctr/utils/visualization.py

@@ -218,13 +218,16 @@ def visualize_page(
  int(page['dimensions'][1] * word['geometry'][0][0]),
  int(page['dimensions'][0] * word['geometry'][0][1])
  )
- ax.text(
- *text_loc,
- word['value'],
- size=10,
- alpha=0.5,
- color=(0, 0, 1),
- )
+
+ if len(word['geometry']) == 2:
+ # We draw only if boxes are in straight format
+ ax.text(
+ *text_loc,
+ word['value'],
+ size=10,
+ alpha=0.5,
+ color=(0, 0, 1),
+ )
 
  if display_artefacts:
  for artefact in block['artefacts']:
@@ -251,7 +254,6 @@ def visualize_page(
 def synthesize_page(
  page: Dict[str, Any],
  draw_proba: bool = False,
- font_size: int = 13,
  font_family: Optional[str] = None,
 ) -> np.ndarray:
  """Draw a the content of the element page (OCR response) on a blank page.