I have object detection model with 2 classes and I want to visualize only one class per image.

Here is the code with two images one_image and two_image. With class one the visualization is only show on one_image and two_image show only class two visualization.

The problem is in both images it show same result how to differentiate it?

`def one_class_show_inference(self, model, image_np, category_index, vis_thresh, max_boxes_to_draw):

    output_dict = self.run_inference_for_single_image(model, image_np)
    one_img = image_np.copy()
    two_image = image_np.copy()

    for res in output_dict['detection_classes']:
        if res == 1:
            # Visualization of the results of a detection.
            visualize_boxes_and_labels_on_image_array(
                one_img,
                output_dict['detection_boxes'],
                output_dict['detection_classes'],
                output_dict['detection_scores'],
                category_index,
                instance_masks=output_dict.get('detection_masks_reframed', None),
                use_normalized_coordinates=True,
                min_score_thresh = vis_thresh,
                max_boxes_to_draw = max_boxes_to_draw,
                line_thickness=4)
        else:
            visualize_boxes_and_labels_on_image_array(
                two_image,
                output_dict['detection_boxes'],
                output_dict['detection_classes'],
                output_dict['detection_scores'],
                category_index,
                instance_masks=output_dict.get('detection_masks_reframed', None),
                use_normalized_coordinates=True,
                min_score_thresh = vis_thresh,
                max_boxes_to_draw = max_boxes_to_draw,
                line_thickness=4)

    return output_dict, one_img, two_image`

I tried to get index of every class and then with respect to index, I tried to get boxes, classes using loop. Is there any simplest way to solve it?