Change what DensePoseResultExtractor returns in case of error
Summary: Change what DensePoseResultExtractor returns in case of error
# Copyright (c) Facebook, Inc. and its affiliates.
import logging
-from typing import Sequence
+from typing import List, Optional, Sequence, Tuple
import torch
from detectron2.layers.nms import batched_nms
from detectron2.structures.instances import Instances
from densepose.converters import ToChartResultConverterWithConfidences
+from densepose.structures import DensePoseChartResultWithConfidences
from densepose.vis.bounding_box import BoundingBoxVisualizer, ScoredBoundingBoxVisualizer
from densepose.vis.densepose_results import DensePoseResultsVisualizer
from .base import CompoundVisualizer
Scores = Sequence[float]
+DensePoseChartResultsWithConfidences = List[DensePoseChartResultWithConfidences]
def extract_scores_from_instances(instances: Instances, select=None):
@@ -77,10 +79,12 @@ class ScoredBoundingBoxExtractor(object):
class DensePoseResultExtractor(object):
"""
- Extracts DensePose result from instances
+ Extracts DensePose chart result with confidences from instances
"""
- def __call__(self, instances: Instances, select=None):
+ def __call__(
+ self, instances: Instances, select=None
+ ) -> Tuple[Optional[DensePoseChartResultsWithConfidences], Optional[torch.Tensor]]:
if instances.has("pred_densepose") and instances.has("pred_boxes"):
dpout = instances.pred_densepose
boxes_xyxy = instances.pred_boxes
@@ -92,7 +96,7 @@ class DensePoseResultExtractor(object):
results = [converter.convert(dpout[i], boxes_xyxy[[i]]) for i in range(len(dpout))]
return results, boxes_xywh
else:
- return None
+ return None, None
class CompoundExtractor(object):
