Process only a subset of classes?

[Trotts]

Hi,

I'm wondering if it is possible in the current codebase to postprocess boxes only from a certain subset of classes? In my work, I have some classes which are usually small and close together as well as some large classes which have big distances between them (and may be larger than the patch size used to train the model). Because of this, it would make sense in my case for SAHI to only postprocess boxes of the larger classes, and leave the small, close together class boxes untouched.

Thanks :)

[pwernette]

I was just coming here to ask the same thing. Searching only for a subset of classes within the images would be extremely beneficial.

[pwernette]

After doing a bit of digging into the code, it looks like modifying the perform_inference function kwargs object to include a classes argument would allow the user to specify what classes they want to predict/identify in the input image/image slice. In yolov8.py this corresponds to kwargs dictionary on line 68 and prediction_result object on line 73. It wouldn't take much to propagate another user argument (similar to confidence threshold up to the get_sliced_prediction() function. I'd be interested in hearing more from the authors though.

[pwernette]

Hi,

I'm wondering if it is possible in the current codebase to postprocess boxes only from a certain subset of classes? In my work, I have some classes which are usually small and close together as well as some large classes which have big distances between them (and may be larger than the patch size used to train the model). Because of this, it would make sense in my case for SAHI to only postprocess boxes of the larger classes, and leave the small, close together class boxes untouched.

Thanks :)

My impatience and curiosity got the better of me and I couldn't wait for a solution to this issue, so I figured out a workaround for now that takes the result object from the sliced prediction and filters through all objects, keeping only those in a list you specify. Here is the function, in case you or anybody else is interested:

def filter_objects(results_obj, keep_classes=['person','umbrella'], keep_list_type='name', verbose=1):
    keep_list = []
    for i,v in enumerate(results_obj.object_prediction_list):
        if keep_list_type == 'name':
            if v.category.name in keep_classes:
                keep_list.append(i)
        else:
            if v.category.id in keep_classes:
                keep_list.append(i)
    if verbose == 1:
        print('   keeping {} of {} objects'.format(len(keep_list), len(results_obj.object_prediction_list)))
    tmp = []
    [tmp.append(results_obj.object_prediction_list[i]) for i in keep_list]
    results_obj.object_prediction_list = tmp

Admittedly it's not the cleanest, but it does the trick and is at least efficient enough for the time being. It would still be ideal if we could specify the classes of interest as you would normally with the YOLO_model.predict() function.

[HuangYison]

In batched_greedy_nmm, line 168

category_ids = object_predictions_as_tensor[:, 5].squeeze()
keep_to_merge_list = {}
for category_id in torch.unique(category_ids):

i think you can do some filters here, exclude the small classes or include the large classes

[fcakyon]

Unfortunately, it's not supported atm, but it would make a good PR 👍🏻 We are open for PRs.