Non-English texts still detected after fine-tuning text detection model

[Yuvaraj-off]

We fine-tuned db_mobilenet_v3_large using the pdfa-eng-wds dataset to detect only English text, as suggested in issue #1564. Though, we trained with 10 epoch and our accuracy results isnt the best for that, we expected the detections to be limited to english texts, but it is detected non english and other junks as texts. Are we missing something here.

Steps Taken
• Used db_mobilenet_v3_large as the base model and trained it on pdfa-eng-wds.
• The training results are below:

• Loaded it in our OCR pipeline and tested on a sample aadhaar card image and got the following result:

[felixdittrich92]

Hi @Yuvaraj-off 👋,

In general I don't think this will work this way - our dataset for the detection pre-training does also contain only latin text but the models learns really well to generalize for mostly any kind of text - keep in mind that's all CNN-based architectures so there is no "textual understanding" - It could maybe work by providing negative samples where different text is on the image but only the english annotated - but no guarante

Have you fine tuned our model or trained from scratch ?
Additional (not sure if the issues still exist) - the pdfa-eng-wds dataset had lots of issues wrong annotations / missing annotations (which is also really clear from the shared metrics) :)

Best,
Felix

[Yuvaraj-off]

Thanks for the quick response and insights, @felixdittrich92!

We've trained our model from scratch and are now focusing on using a our own dataset annotated exclusively for English text. We'll update you on the results once we test this approach.

Additionally, we're exploring the possibility of using a YOLO model for text detection. Would it be feasible to integrate a YOLO-based text detection model with our existing docTR pipeline? We'd love to hear your insights on this!

[Yuvaraj-off]

class YOLOProcessor:
"""YOLOv8 model wrapper using PyTorch (.pt model)."""

def __init__(self, model_path):
    print("Loading YOLOv8 PyTorch model...")
    self.model = YOLO(model_path)
    print("YOLOv8 PyTorch model loaded successfully.")

def forward(self, image: np.ndarray):
    # print(image.shape[:2],"image shape")
    input_tensor = self.preprocess_for_yolov8(image)
    results = self.model(input_tensor)[0]
    # print(results,"results")
    return self.parse_yolo_output(results, image.shape[:2])

@staticmethod
def preprocess_for_yolov8(image: np.ndarray) -> torch.Tensor:
    if image.dtype != np.uint8:  # Ensure correct dtype
        image = (image * 255).astype(np.uint8)

    transform = T.Compose([
        T.ToPILImage(),
        T.Resize((640, 640)),
        T.ToTensor(),  # Converts to (C, H, W), scales to [0,1]
    ])

    input_tensor = transform(image)

    return input_tensor.unsqueeze(0)  # Add batch dimension

@staticmethod
def parse_yolo_output(results, img_size):
    """Parse YOLOv8 output and normalize bounding boxes."""
    h, w = img_size
    boxes = results.boxes.xyxy.cpu().numpy()  # Bounding boxes (x1, y1, x2, y2)
    scores = results.boxes.conf.cpu().numpy()  # Confidence scores
    class_ids = results.boxes.cls.cpu().numpy()

    # 🔹 Filter out non-text classes (if applicable)
    TEXT_CLASS_ID = 0  # Change this to match your text class index
    valid_indices = class_ids == TEXT_CLASS_ID

    if np.sum(valid_indices) == 0:
        print("Warning: No text detections found.")
        return {"boxes": np.array([]), "scores": np.array([]), "class_ids": np.array([])}

    # Apply filtering
    boxes = boxes[valid_indices]
    scores = scores[valid_indices]
    class_ids = class_ids[valid_indices]

    # Normalize bounding boxes
    boxes[:, [0, 2]] /= w  # Normalize x-coordinates
    boxes[:, [1, 3]] /= h  # Normalize y-coordinates

    return {
        "boxes": boxes.astype(np.float32),
        "scores": scores.astype(np.float32),
        "class_ids": class_ids.astype(np.int32)
    }

class CustomDetectionPredictor(DetectionPredictor):
"""Custom DetectionPredictor for docTR using YOLOv8 PyTorch."""

def __init__(self, model_path: str):
    super().__init__(pre_processor=None, model=torch.nn.Module())
    self.yolo_model = YOLOProcessor(model_path)

@torch.inference_mode()
def forward(self, pages: list[np.ndarray], return_maps: bool = False, **kwargs):
    if any(page.ndim != 3 for page in pages):
        raise ValueError("Incorrect input shape: all pages must be multi-channel 2D images.")

    detections = [self.yolo_model.forward(page) for page in pages]
    feature_maps = [np.zeros((1, 1)) for _ in pages]  # Dummy feature maps

    return (detections, feature_maps) if return_maps else detections

[Yuvaraj-off]

Here's a refined version with improved clarity and structure:

The key challenge has been replicating the seg_maps output expected by DetectionPredictor, particularly in this snippet:

if return_maps:
    seg_maps = [pred.numpy() for batch in predicted_batches for pred in batch["out_map"]]
    return preds, seg_maps
return preds

For the YOLO-based model, I need to generate a compatible seg_maps output to align with docTR’s expected structure. Additionally, I need guidance on setting up a suitable postprocessor to ensure seamless integration with docTR’s pipeline. Any insights on structuring this for YOLO-based text detection would be helpful! 🚀

[felixdittrich92]

Hi @Yuvaraj-off 👋,

Correct the seg_maps are of shape (H, W, 1) and contain the prob scores for each pixel so everything you need to do is to create a zero filled mask (dtype = float) and fill all the detected areas from your yolo model with 1.0 for example

A yeah sry I missed this part:

doctr/doctr/models/predictor/pytorch.py

Line 85 in de2c359

np.where(out_map > getattr(self.det_predictor.model.postprocessor, "bin_thresh"), 255, 0).astype(np.uint8)

from doctr.models.detection.core import DetectionPostProcessor

class CustomDetectionPredictor(DetectionPredictor):
"""Custom DetectionPredictor for docTR using YOLOv8 PyTorch."""

    def __init__(self, model_path: str):
        super().__init__(pre_processor=None, model=torch.nn.Module())
        self.yolo_model = YOLOProcessor(model_path)
        self.model.postprocessor = DetectionPostProcessor(box_thresh=0.5, bin_thresh=0.5, assume_straight_pages=True)

As mentioned it is possible but not designed for the use case to plug in any model - especially if you want to use it with the whole pipeline 😅

The format for the preds with assume_straight_pages=True is [xmin, ymin, xmax, ymax, box_score]

[{'words': array([[0.90203401, 0.93945312, 0.91446805, 0.95214844, 0.33106691],
       [0.83157444, 0.85449219, 0.91308649, 0.86523438, 0.49566621],
       [0.78183828, 0.85449219, 0.83157444, 0.86621094, 0.46532482],
       ...,
       [0.44473759, 0.03320312, 0.50138156, 0.04492188, 0.57190824],
       [0.39776455, 0.03320312, 0.42815887, 0.04785156, 0.51862246],
       [0.24302981, 0.03320312, 0.27066101, 0.04589844, 0.55790615]])}]

and with assume_straight_pages=False

[[x1, y1], [x2, y2], [x3, y3], [x4, y4], [0, box_score]]

[{'words': array([[[0.90203401, 0.93945312],
        [0.91170493, 0.93945312],
        [0.91170493, 0.95019531],
        [0.90203401, 0.95019531],
        [0.        , 0.42683684]],

       [[0.83157444, 0.85449219],
        [0.91032337, 0.85449219],
        [0.91032337, 0.86328125],
        [0.83157444, 0.86328125],
        [0.        , 0.5850419 ]],

       [[0.78183828, 0.85449219],
        [0.82881132, 0.85449219],
        [0.82881132, 0.86425781],
        [0.78183828, 0.86425781],
        [0.        , 0.5922716 ]],

       ...,

       [[0.44473759, 0.03320312],
        [0.49861844, 0.03320312],
        [0.49861844, 0.04296875],
        [0.44473759, 0.04296875],
        [0.        , 0.65185542]],

       [[0.39638299, 0.03222656],
        [0.42539575, 0.03222656],
        [0.42539575, 0.04589844],
        [0.39638299, 0.04589844],
        [0.        , 0.66118372]],

       [[0.24302981, 0.03320312],
        [0.26789789, 0.03320312],
        [0.26789789, 0.04394531],
        [0.24302981, 0.04394531],
        [0.        , 0.66179151]]])}]

where box_score is the fast computed objectness_score so you can fill this value with any value you want 0.0-1.0
And all coords are relative

[felixdittrich92]

Hope this helps :)

[Yuvaraj-off]

This should really help. Thanks for your insights, will look into and revert back