0
$\begingroup$

I have basic knowledge about non-max suppression and I know how it works for multiple classes, but what if I want to get a prediction for two classification problems?

I give you an example. So I have a board (an image) that could contain a couple of small images of animals. An animal can be a dog or can be a cat, so I have two classes. So, I feed my model with input, and as output I got two tensors: one is $(N, 4)$ where $N$ is a number of anchors and 4 points of anchor(expected to be in $(x1, x2, y1, y2)$) and my second tensor is $(N, 2)$ where $(N, 0)$ is the probability that object is a cat, and $(N, 1)$ is the probability of dog.

So, when I want to get predictions of my bounding boxes, I've just applied independently anchors, score_cats and anchor, score_dogs to pytorch.nms, and it returns me indexes_cats of anchors for cat and indexes_dogs for dogs. Next, I can return a list of bounding boxes where a single bounding box contains $(x1, y1, x2, y2)$ and (cat/dog (0/1)) - (cat when anchors index was in indexes_cast and dog when anchor index was in indexes_dogs).

But what if I want to classify if an image of a cat/dog is rotated? (flipped by 90 degrees). I extended my model to give me the 3rd tensor as an output, which looks like $(N, 2)$ where $(N, 0)$ is the probability of being rotated and $(N, 1)$ is the probability of not rotated. And my idea is to apply nms again, but this time for score_rotated and score_norot. And when creating predictions, check if the anchor is in indexes_cat and in indexes_rotated, and so on.

Pseudocode for my intuition:

# bounding box is a tuple of size 6 - (x1, x2, y1, y2, is_dog, is_rotated) 
# is_dog == 0 -> cat, is_dog == 1 -> dog
# is_rotated == 0 -> no_rotated, is_rotated == 1 -> rotated

scores_cats = animal_scores(:, 0)
scores_dogs = animal_scores(:, 1)
indexs_cats = nms(anchors, scores_cats, iou_treshold = 0.5)
indexs_dogs = nms(anchors, scores_dogs, iou_treshold = 0.5)

scores_rotated = is_rotated(:, 0)
scores_norot = is_rotated(:, 1)
indexs_rotated = nms(anchors, scores_rotated, iou_treshold = 0.5)
indexs_norot = nms(anchors, scores_norot, iou_treshold = 0.5)

# return list of bounding boxes
L = []
for i in scores_cats:
    if i in indexs_rotated:
        box = (anchors[i].x1, anchors[i].x2, anchors[i].y1, anchors[i].y2, 0, 1)
        L.append(box)
    elif i in indexs_norot
        box = (anchors[i].x1, anchors[i].x2, anchors[i].y1, anchors[i].y2,, 0, 0)
        L.append(box)

for i in scores_dogs:
    if i in indexs_rotated:
        box = (anchors[i].x1, anchors[i].x2, anchors[i].y1, anchors[i].y2,, 1, 1)
        L.append(box)
    elif i in indexs_norot:
        box = (anchors[i].x1, anchors[i].x2, anchors[i].y1, anchors[i].y2,, 1, 0)
        L.append(box)

where L is a list of bounding boxes and each bounding box contains information about anchor points and animal_class and rotated_class

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.