Segmenting and saving each class as image to allow OCR

[elnazsn1988]

🚀 Feature

Hi, is there an internal feature which lets each classed be saved as a seperate segment, or image? I am trying to identify tables, seperate and then run through a tabular data analyzer and ocr - so far am able to get the image predictions with your code, but not the actual annotations/segmented fields for further analysis/ocr.

Motivation

Pitch

Note that you can implement many features by extending detectron2.
See projects for some examples.
We would only consider adding new features if they are relevant to many users.

[hpanwar08]

@elnazsn1988 You can try the below code. It extracts the bounding boxes and then crop the image based on the bounding box and save.

classes = ['text', 'title', 'list', 'table', 'figure']
default_predictor = detectron2.engine.defaults.DefaultPredictor(cfg)
img = detectron2.data.detection_utils.read_image(path_to_image, format="BGR")
predictions = default_predictor(img)
instances = predictions["instances"].to('cpu')

pred_classes = instances.pred_classes
labels = [classes[i] for i in pred_classes]
boxes = instances.pred_boxes
if isinstance(boxes, detectron2.structures.boxes.Boxes):
    boxes = boxes.tensor.numpy()
else:
    boxes = np.asarray(boxes)

from PIL import Image
img = Image.fromarray(img)

for label, bbox in zip(labels, boxes):
    if label == "table":
        cropped_img = img.crop(bbox)
        croppped_img.save(f"{label}_{bbox}.png")

[elnazsn1988]

@hpanwar08 thanks for the above, it throws errors at img being extracted as a numpy object : ```

cropped_img = img.crop(bbox) AttributeError: 'numpy.ndarray' object has no attribute 'crop'

[hpanwar08]

@elnazsn1988 I have updated the code, it should work now

[akshay94950]

@hpanwar08 where to add the code ?which file?

[hpanwar08]

@akshay94950 You can add this code in a new python file and run it.

[akshay94950]

@hpanwar08
the above code is throwing error:
49 if label == "text":
50 cropped_img = img.crop(bbox)
---> 51 croppped_img.save(f"{label}_{bbox}.png")
52
53

NameError: name 'croppped_img' is not defined

[hpanwar08]

can you show the entire code.

[akshay94950]

import argparse
import glob
import multiprocessing as mp
import os
import time
import cv2
import tqdm
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detectron2.utils.logger import setup_logger
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg

from predictor import VisualizationDemo

constants

WINDOW_NAME = "COCO detections"
classes = ['text', 'title', 'list', 'table', 'figure']
cfg = get_cfg()
cfg.merge_from_file("/content/gdrive/My Drive/detectron2-master/configs/DLA_mask_rcnn_R_101_FPN_3x.yaml")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set threshold for this model
cfg.MODEL.WEIGHTS = os.path.join( '/content/gdrive/My Drive/Unt', 'model_final.pth')
cfg.MODEL.DEVICE='cpu'
default_predictor = detectron2.engine.defaults.DefaultPredictor(cfg)
img = detectron2.data.detection_utils.read_image("/content/gdrive/My Drive/unt/page1.jpg" , format="BGR")
predictions = default_predictor(img)
instances = predictions["instances"].to('cpu')

pred_classes = instances.pred_classes
labels = [classes[i] for i in pred_classes]
boxes = instances.pred_boxes
if isinstance(boxes, detectron2.structures.boxes.Boxes):
boxes = boxes.tensor.numpy()
else:
boxes = np.asarray(boxes)
print('b',boxes)
from PIL import Image
img = Image.fromarray(img)

for label, bbox in zip(labels, boxes):
if label == "text":
cropped_img = img.crop(bbox)
cropped_img.save(f"/content/gdrive/My Drive/unt/{label}_{bbox}.jpg")

[akshay94950]

it was because of spelling mistake in that above code .thanks @hpanwar08 but it shows this error:
AttributeError: 'numpy.ndarray' object has no attribute 'crop'

[akshay94950]

thanks bro its working.....@hpanwar08

[elnazsn1988]

@elnazsn1988 I have updated the code, it should work now

HI @hpanwar08 - codes working great thanks, do you happen to know how I can annotate a new image and retrain the existing weights for the new image? Ive been trying to donwload the full dataset from:
wget -O datz/publaynet.tar.gz https://dax-cdn.cdn.appdomain.cloud/dax-publaynet/1.0.0/publaynet.tar.gz
but somehow the download is unstable so I cant get the whole dataset to retrain and annotate/add my file.

[hpanwar08]

I did not understand your question. If you want to annotate new images then you can have a look at https://github.com/wkentaro/labelme
It's better to download the dataset in parts train-0.tar.gz, train-1.tar.gz etc rather than as one huge file.

[elnazsn1988]

I did not understand your question. If you want to annotate new images then you can have a look at https://github.com/wkentaro/labelme
It's better to download the dataset in parts train-0.tar.gz, train-1.tar.gz etc rather than as one huge file.

@hpanwar08 is there a way to add a new annotated image and train your pretrained model on it without retraining the whole thing? As in use your weights from the trained model and retrain on a new image set.

[hpanwar08]

Yes, you can annotated your custom dataset and save it in COCO format and train using train_net_dla.py

[akshay94950]

can we obtain predictions output image without figures class for using it in ocr?or how to use it for text extraction ?