ImageAI how to prevent detection box from being drawn

123 Views Asked by At

I'm using ImageAI and ResNet50 model to detect and extract pictures of people from photos. It all works very well, but I can't figure out how to disable the detection box from being drawn. Obviously, when multiple people are in the same shot the detection box from one then bleeds onto extracted images of other people.

Does anyone know how to do it? I've looked around and could only found one other question with the same issue and for some reason the answer is talking about OpenCV instead of imageai.

My code:

from pathlib import Path

import cv2
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.applications.resnet50 import ResNet50

from imageai.Detection import ObjectDetection

physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)

detector = ObjectDetection()
detector.setModelTypeAsRetinaNet()
detector.setModelPath('resnet50_coco_best_v2.1.0.h5')
detector.loadModel()
person = detector.CustomObjects(person=True)

source_folder = Path(r'd:/pics')
files = [f for f in source_folder.glob('*.jpg')]
output_folder = Path('c:/python/datasets/raw/extracted_images')

for idx, i in enumerate(files):
    print(f'Extracting images from: {i.name}')
    image = i 
    output_path = output_folder.joinpath(f'{idx}.jpg')
    detected_image, detections = detector.detectCustomObjectsFromImage(input_image=str(image),
                                                                       output_image_path=str(output_path),
                                                                       minimum_percentage_probability=20,
                                                                       input_type='file',
                                                                       output_type='file',
                                                                       display_percentage_probability=False,
                                                                       extract_detected_objects=True,
                                                                       custom_objects=person)

EDIT: I've found a hacky way to fix the issue - by going to the source of detectObjetcFromImage() function and commenting out this bit:

# image_copy = draw_boxes(image_copy, 
#                 box_points,
#                 display_box,
#                 label, 
#                 percentage_probability, 
#                 self.__box_color)

Is there a less intrusive way of dealing with this?

0

There are 0 best solutions below