CODEWITHSUNDEEP
pythontensorflowcomputer-visiondata-scienceobject-detection
CodeMONSTER
CodeMONSTER

Reputation: 21

How to add custom colors to Tensorflow object detection python model bound boxes?

Currently i am using the standard Tensorflow object detection script which is custom trained, i would like to change the colors of the bound boxes to suit the nature of my application. However i cannot seem to find a way to do so. For example instead of drawing a green box id like to draw a red box around a detected object. Thank you in advance!

Upvotes: 0

Views: 2086

Answers (2)

Yusuf
Yusuf

Reputation: 743

  1. add “color” attribute in the label map pbtxt file. i.e.

    item { name: "/m/01g317" id: 1 display_name: "person" color: "Pink" }

  2. Open file “research/object_detection/protos/string_int_label_map.proto”. Add the following line.

    optional string color = 4;

    Be aware about the semicolon, curly braces.

  3. then it is required to serialize the data, so run the following command from research folder

    protoc object_detection/protos/*.proto --python_out=.

    Before that you must install protobuf based on your OS.

  4. change the code of function “convert_label_map_to_categories” of the file object_detection/utils/lable_map_util.py

    categories = [] list_of_ids_already_added = [] if not label_map: label_id_offset = 1 for class_id in range(max_num_classes): categories.append({ 'id': class_id + label_id_offset, 'name': 'category_{}'.format(class_id + label_id_offset) }) return categories for item in label_map.item:

    if not 0 < item.id <= max_num_classes:
  logging.info(
      'Ignore item %d since it falls outside of requested '
      'label range.', item.id)
  continue
if use_display_name and item.HasField('display_name'):
  name = item.display_name
else:
  name = item.name

if use_display_name and item.HasField('color'):
  color = item.color
else:
  color = ''

if item.id not in list_of_ids_already_added:
  list_of_ids_already_added.append(item.id)
  categories.append({'id': item.id, 'name': name, 'color': color})          
return categories

5.Open the file “object_detection/utils/visualization_utils.py”. Go to the function named “visualize_boxes_and_labels_on_image_array”. Add the following code 

else:
         if classes[i] in category_index.keys():
            class_color = category_index[classes[i]]['color']
            box_to_color_map[box] = class_color

after the code 

if agnostic_mode:
          box_to_color_map[box] = 'DarkOrange'

Upvotes: 0

Chris Norris
Chris Norris

Reputation: 197

I sort of found a way - after much trouble. I found nothing documenting how to do this. Sort of, as some colors don't seem to work.

Open "visualizations_utils.py". Should be in Lib\site-packages\utils. Rows 41 to 63 are your colors.

Directly under row 164,

draw = ImageDraw.Draw(image),

enter a new row

color = 'Pink'

Save it, and you have now changed the color to a pinkish color. Row 175, you can make the label text smaller.

Some colors don't seem to work, like "Red".

Upvotes: 1

Related Questions