open-mmlab / mmdetection

OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io
Apache License 2.0
28.52k stars 9.29k forks source link

How to select which classes to detect on pretrained model? #11700

Open SFSimDev opened 1 month ago

SFSimDev commented 1 month ago

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.