open-mmlab/mmdetection

How to select which classes to detect on pretrained model?

SFSimDev opened this issue · 0 comments

Hi,

I'm using the pretrained rtmdet-r_l_syncbn_fast_2xb4-aug-100e_dota model to scan some images and check for just one type of class.

I only want to detect one of the many classes it is trained on. How can I do that within my code?

def prepare_inference():

    config_file = 'models/rtmdet-r_l_syncbn_fast_2xb4-aug-100e_dota.py'
    checkpoint_file = 'models/rtmdet-r_l_syncbn_fast_2xb4-aug-100e_dota_20230224_124735-ed4ea966.pth'

    model = init_detector(config_file, checkpoint_file, device='cuda:0')  # or device='cuda:0'
    return model

def start_inference(model, file):

    result = inference_detector(model, file)

Thanks for any help you can give.