How to select which classes to detect on pretrained model?
SFSimDev opened this issue · 0 comments
SFSimDev commented
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.