Repository: pytorch/vision
Pull request: Modifying keypoint_rcnn.py for keypoint_predictor issue
Fixes #5126

Hi @datumbox, I have created a PR for the issue [#5126](https://github.com/pytorch/vision/issues/5126). This is my first PR. Please review and comment on it. Thank you!


cc @datumbox

Problem statement:
Bug in keypoint_rcnn initializer
### 🐛 Describe the bug

Line 1890 in keypoint_rcnn
"
        if num_classes is not None:
            if keypoint_predictor is not None:
                raise ValueError("num_classes should be None when keypoint_predictor is specified")
"
Shouldn't we change num_classes to num_keypoints in this case?

### Versions

Newest

cc @datumbox

Hints from issue discussion:
@JunzheJosephZhu You are right, this is a bug. Instead of `num_classes` we should be checking the `num_keypoints`.

The problem now is that the default value of `num_keypoints=17`, so checking for `None` makes no sense. One way to address this is to change the default value to `None` and for maintaining backwards compatibility overwrite the None value to 17 after the highlighted-if. Let me know if you would like to send a PR to fix this.

Commit message: Modifying keypoint_rcnn.py

Patch to review:
diff --git a/torchvision/models/detection/keypoint_rcnn.py b/torchvision/models/detection/keypoint_rcnn.py
index 71893b3b0d0..1a0586b9d1c 100644
--- a/torchvision/models/detection/keypoint_rcnn.py
+++ b/torchvision/models/detection/keypoint_rcnn.py
@@ -188,16 +188,19 @@ def __init__(
         keypoint_roi_pool=None,
         keypoint_head=None,
         keypoint_predictor=None,
-        num_keypoints=17,
+        num_keypoints=None,
     ):
 
         assert isinstance(keypoint_roi_pool, (MultiScaleRoIAlign, type(None)))
         if min_size is None:
             min_size = (640, 672, 704, 736, 768, 800)
 
-        if num_classes is not None:
+        if num_keypoints is not None:
             if keypoint_predictor is not None:
-                raise ValueError("num_classes should be None when keypoint_predictor is specified")
+                raise ValueError("num_keypoints should be None when keypoint_predictor is specified")
+
+        if num_keypoints is None:
+            num_keypoints = 17
 
         out_channels = backbone.out_channels
 


Code review comments:
diff_hunk: @@ -188,16 +188,19 @@ def __init__(
         keypoint_roi_pool=None,
         keypoint_head=None,
         keypoint_predictor=None,
-        num_keypoints=17,
+        num_keypoints=None,
     ):
 
         assert isinstance(keypoint_roi_pool, (MultiScaleRoIAlign, type(None)))
         if min_size is None:
             min_size = (640, 672, 704, 736, 768, 800)
 
-        if num_classes is not None:
+        if num_keypoints is not None:
             if keypoint_predictor is not None:
-                raise ValueError("num_classes should be None when keypoint_predictor is specified")
+                raise ValueError("num_keypoints should be None when keypoint_predictor is specified")
+
+        if num_keypoints is None:
path: torchvision/models/detection/keypoint_rcnn.py
text: @user1:
Nit: Simplify by
```suggestion
        else:
```