# How do I apply non-max suppression for 2-classes problems?

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