@@ -349,17 +349,22 @@ def __init__(self, in_channels, dim_reduced, num_classes):
349349 # nn.init.constant_(param, 0)
350350
351351
352+ _COMMON_META = {
353+ "task" : "image_object_detection" ,
354+ "architecture" : "MaskRCNN" ,
355+ "categories" : _COCO_CATEGORIES ,
356+ "interpolation" : InterpolationMode .BILINEAR ,
357+ }
358+
359+
352360class MaskRCNN_ResNet50_FPN_Weights (WeightsEnum ):
353361 COCO_V1 = Weights (
354362 url = "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth" ,
355363 transforms = ObjectDetection ,
356364 meta = {
357- "task" : "image_object_detection" ,
358- "architecture" : "MaskRCNN" ,
365+ ** _COMMON_META ,
359366 "publication_year" : 2017 ,
360367 "num_params" : 44401393 ,
361- "categories" : _COCO_CATEGORIES ,
362- "interpolation" : InterpolationMode .BILINEAR ,
363368 "recipe" : "https://github.com/pytorch/vision/tree/main/references/detection#mask-r-cnn" ,
364369 "map" : 37.9 ,
365370 "map_mask" : 34.6 ,
@@ -369,7 +374,19 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
369374
370375
371376class MaskRCNN_ResNet50_FPN_V2_Weights (WeightsEnum ):
372- pass
377+ COCO_V1 = Weights (
378+ url = "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth" ,
379+ transforms = ObjectDetection ,
380+ meta = {
381+ ** _COMMON_META ,
382+ "publication_year" : 2021 ,
383+ "num_params" : 46359409 ,
384+ "recipe" : "https://github.com/pytorch/vision/pull/5773" ,
385+ "map" : 47.4 ,
386+ "map_mask" : 41.8 ,
387+ },
388+ )
389+ DEFAULT = COCO_V1
373390
374391
375392@handle_legacy_interface (
0 commit comments