Detection predictor results are not ordered properly

[dhea1323]

Hello @fg-mindee, @charlesmindee , and community,

I used detection predictor results from doctr and send them to trocr recognition for handwritten. After I put all the words together, I found that the text was not in ordered properly. This is caused by the results of the detection predictor from doctr.

Can anyone help solve this problem ?
I used the following code,

doc = DocumentFile.from_pdf("path/to/pdf/file")
#Word detection -> localizing words
det_model = db_resnet50(pretrained=False)
det_params = torch.load(".../db_resnet50.pt", map_location="cpu")
det_model.load_state_dict(det_params)
det_predictor = DetectionPredictor(PreProcessor((1024, 1024), batch_size=1, mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287)), det_model)
detection = det_predictor(doc)

#Word recognition
model_path = '..../trocr-base-handwritten.pt'
model, cfg, task, generator, bpe, img_transform, device = init(model_path, beam)
bucket_doc = []
for page in range(len(detection)):
  bucket_word = []
  for word in range(len(detection[page])):
    xmin, ymin, xmax, ymax, score = detection[page][word]
    dy,dx,d = doc[page].shape
    xmin_ = int(xmin*dx)
    xmax_ = int(xmax*dx)
    ymin_ = int(ymin*dy)
    ymax_ = int(ymax*dy)
    word_img = doc[page][ymin_:ymax_, xmin_:xmax_]
    pre_processing = preprocess(word_img, img_transform)
    text = get_text(cfg, generator, model, pre_processing, bpe)
    bucket_word.append(text)
  bucket_doc.append(bucket_word)

[dhea1323]

Any updates? @frgfm

[dhea1323]

I want to know how to sort the words from the output generated from the DetectionPredictor in doctr. The model recognition that I will use later is not from doctr but trocr. Is there a separate or other callable function in doctr to sort words without having to use the OCRPredictor ?

[frgfm]

As it is, the feature is not developed yet for the DetectionPredictor. But I guess we could move the sorting part earlier or make it compatible with the detection output. What do you think @charlesmindee ?

[dhea1323]

thank you very much for the response, I hope it will be possible soon

[dhea1323]

Hi, any updates?

[frgfm]

Hey @dhea1323 👋

Sorry about the late reply! The feature hasn't been integrated yet, but I have a way for you to handle this:

from doctr.models.builder import DocumentBuilder
# Everything to define your det_predictor
det_predictor = ...
detection = det_predictor(doc)
# Sort the boxes using the builder private method
ordered_detection = [boxes[DocumentBuilder._sort_boxes(boxes[:, :4])[0]] for boxes in detection]

Let me know if you still encounter some problems :)