How to see classification probabilities in the object detection

iamexperimentingnow: <p>currently, working on object detection model, I have built the model using Faster RCNN with ResNet50 as backbone and I have 4 target classes. Now, I would like to see the class probabilities while inferencing.<br>
eg: 4 classes[car, person, truck, train], when I pass the truck image in the model, it gives me output with bbox, predicted class(truck) and confidence score. How can I get the class probability [0.2, 0,0.75, 0.05].</p>
<p>For ‘car’ its probability is 0.2<br>
For ‘person’ its probability is 0<br>
For ‘truck’ its probability is 0.75<br>
For ‘train’ its probability is 0.05</p>
<p>Has anyone tried this? any articles or example script would help me to replicate the same.<br>
<a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Thanks in advance.</p>

bigbreadguy: <p>Normally, probability-like predictions over multi classes from a classification model is softmax result for each elements in prediction vector. You could pass the predictions into softmax and get probabilities.</p>

iamexperimentingnow: <p><a class="mention" href="/u/bigbreadguy">@bigbreadguy</a>, I’m using inbuilt fasterrcnn detector from pytorch. I want the extract the last layer from it. can you provide me any suggestion how to do that?</p>

bigbreadguy: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/4/b4cbe7dcd3f3946d3c32c533568cce93978e17a3.png" data-download-href="https://discuss.pytorch.org/uploads/default/b4cbe7dcd3f3946d3c32c533568cce93978e17a3" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/4/b4cbe7dcd3f3946d3c32c533568cce93978e17a3.png" alt="image" data-base62-sha1="pNoOgadm5w7U5xJUvPRvfWVLPbl" width="324" height="500" data-dominant-color="F4F4F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">976×1505 33.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>You can iterate through layers(modules, strictly) in the model.</p>
<p>Please refer to attached scripts.</p>

bigbreadguy: <p>My bad.</p>
<p>You need to entangle the layers as built-in methods like this.</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1sCL9HVl3EQXOdQbkevjpBsHnIHugZn3h?usp=sharing">
  <header class="source">
      <img src="https://colab.research.google.com/img/favicon.ico?vrz=colab-20211207-060055-RC00_414675949" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1sCL9HVl3EQXOdQbkevjpBsHnIHugZn3h?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="260" height="260">

<h3><a href="https://colab.research.google.com/drive/1sCL9HVl3EQXOdQbkevjpBsHnIHugZn3h?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


jonathan-ibex: <p>Hey <a class="mention" href="/u/bigbreadguy">@bigbreadguy</a> and <a class="mention" href="/u/iamexperimentingnow">@iamexperimentingnow</a><br>
I tried following your steps but couldn’t reach the wanted results… Any update on this?</p>

Alaric_Tong: <p>You can just overwrite the method and use monkey patch to replace the method, like<br>
model = fasterrcnn_resnet50_fpn(pretrained=True)<br>
model.roi_heads.postprocess_detections = MethodType(postprocess_detections, model.roi_heads)<br>
model.roi_heads.forward = MethodType(forward, model.roi_heads)</p>

Julian_Lehrer: <p><a class="mention" href="/u/jonathan-ibex">@jonathan-ibex</a> can you please be more specific? As in; what did you try, what were the results, and what results did you expect?</p>

annika_igl: <p>Using the answer by <a class="mention" href="/u/alaric_tong">@Alaric_Tong</a> as a starting point, this is my implementation:</p>
<pre><code class="lang-auto">import torch 
import torchvision
import torch.nn.functional as F
from torchvision.ops import boxes as box_ops
 
class RoIHeadsCustom(torchvision.models.detection.roi_heads.RoIHeads):
    def postprocess_detections(
        self,
        class_logits,  # type: Tensor
        box_regression,  # type: Tensor
        proposals,  # type: List[Tensor]
        image_shapes,  # type: List[Tuple[int, int]]
    ):
         # type: (...) -&gt; Tuple[List[Tensor], List[Tensor], List[Tensor], List[Tensor]]
        device = class_logits.device
        num_classes = class_logits.shape[-1]

        boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals]
        pred_boxes = self.box_coder.decode(box_regression, proposals)

        pred_scores = F.softmax(class_logits, -1)

        pred_boxes_list = pred_boxes.split(boxes_per_image, 0)
        pred_scores_list = pred_scores.split(boxes_per_image, 0)

        all_boxes = []
        all_scores = []
        all_labels = []
        all_softmax_scores = []  # List to store softmax scores for each detected object per image

        for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes):
            boxes = box_ops.clip_boxes_to_image(boxes, image_shape)

            # create labels for each prediction
            labels = torch.arange(num_classes, device=device)
            labels = labels.view(1, -1).expand_as(scores)

            # Remove background class at this point
            boxes = boxes[:, 1:]
            scores = scores[:, 1:]
            labels = labels[:, 1:]

            # Keep the original class scores (softmax) for output without the background class
            original_class_scores = scores.detach().clone()
            original_class_scores = original_class_scores.unsqueeze(1).repeat(1, num_classes-1, 1)

            # Flatten everything to filter them
            boxes = boxes.reshape(-1, 4)
            scores = scores.reshape(-1)
            labels = labels.reshape(-1)
            original_class_scores = original_class_scores.reshape(-1, num_classes - 1)

            # Remove low scoring boxes
            inds = torch.where(scores &gt; self.score_thresh)[0]
            boxes, scores, labels, original_class_scores = boxes[inds], scores[inds], labels[inds], original_class_scores[inds]
            # Remove empty boxes
            keep = box_ops.remove_small_boxes(boxes, min_size=1e-2)
            boxes, scores, labels, original_class_scores = boxes[keep], scores[keep], labels[keep], original_class_scores[keep]

            # non-maximum suppression, independently done per class
            keep = box_ops.batched_nms(boxes, scores, labels, self.nms_thresh)
            keep = keep[: self.detections_per_img]
            boxes, scores, labels, original_class_scores = boxes[keep], scores[keep], labels[keep], original_class_scores[keep]

            all_boxes.append(boxes)
            all_scores.append(scores)
            all_labels.append(labels)
            all_softmax_scores.append(original_class_scores)  # Adding softmax scores to the list

        return all_boxes, all_softmax_scores, all_labels #return softmax instead of scores
</code></pre>
<p>The method is an adjusted version of the official implementation of postprocess_detections that can be found here: <a href="https://github.com/pytorch/vision/blob/main/torchvision/models/detection/roi_heads.py#L668" class="inline-onebox" rel="noopener nofollow ugc">vision/torchvision/models/detection/roi_heads.py at main · pytorch/vision · GitHub</a><br>
This ROI head returns a tensor of shape [predictions, number_of_classes] as scores instead of a tensor of shape [predictions] (as the default implementation of postprocess_detections would do).</p>
<p>I then use my custom roi head like this in my model declaration:</p>
<pre><code class="lang-auto">from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor 

def init_object_detection_model(num_classes):

    num_classes = num_classes + 1
    model = fasterrcnn_resnet50_fpn(weights = "DEFAULT") 

    in_features  = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    r = model.roi_heads 
    new_roi_heads = RoIHeadsCustom(r.box_roi_pool, r.box_head, r.box_predictor, 
              r.proposal_matcher.high_threshold, r.proposal_matcher.low_threshold, 
              r.fg_bg_sampler.batch_size_per_image, r.fg_bg_sampler.positive_fraction,
              r.box_coder.weights, r.score_thresh, r.nms_thresh, r.detections_per_img,
              r.mask_roi_pool, r.mask_head, r.mask_predictor,
              r.keypoint_roi_pool, r.keypoint_head, r.keypoint_predictor)
    model.roi_heads = new_roi_heads

    return model
</code></pre>
<p>Hope this helps other people who are looking for a solution to this problem! <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>