search

Layout-Parser / layout-parser

A Unified Toolkit for Deep Learning Based Document Image Analysis
https://layout-parser.github.io/
Apache License 2.0
4.75k stars 456 forks source link

Setting the Color_map argument in 'draw_box' function does not change the color of bounding boxes #106

Closed venkatesh-kulkarni closed 2 years ago

venkatesh-kulkarni commented 2 years ago

I have tried setting the color_map argument in 'draw_box' function but it defaults to 'red' color. I believe the reason for the red color is that the DEFAULT_OUTLINE_COLOR in layoutparser.visualization file has been set to 'red'. The official documentation mentions the following to change the color_map:

A map from block.type to the colors, e.g., {1: 'red'}.

I have tried creating the dictionary in the same way but it does not work. This is the code segment that I am using:

model = lp.Detectron2LayoutModel('lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config', extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.80, "MODEL.ROI_HEADS.NMS_THRESH_TEST", 0.5], label_map = {0: "Text", 1: "Title", 2: "List", 3:"Table", 4:"Figure"})

color_map = {0:'green', 1:'red', 2:'white', 3:'black', 4:'blue'} for i in range(len(images_path)): image = cv2.imread(images_path[i]) image = image[..., ::-1]
layout = model.detect(image) image_new = lp.draw_box(image, layout, box_width=3, color_map = color_map) image_new.save('image_path.jpg')

lolipopshock commented 2 years ago

You might want to try use the following configuration: 

color_map  = {"Text":'green', "Title":'red', "List":'white', "Table":'black', "Figure":'blue'}

The label_map for the models is used to map the numeric model prediction index to the string name (and they are stored in block.type), and the color_map is used for assigning colors based on the block.types.