def __call__(
self,
instances: Instances,
select=None,
) -> Tuple[Optional[DensePoseEmbeddingPredictorOutput],
Optional[torch.Tensor], Optional[List[int]]]:
if not (instances.has("pred_densepose")
and instances.has("pred_boxes")):
return None, None, None
dpout = instances.pred_densepose
boxes_xyxy = instances.pred_boxes
boxes_xywh = extract_boxes_xywh_from_instances(instances)
if instances.has("pred_classes"):
classes = instances.pred_classes.tolist()
else:
classes = None
if select is not None:
dpout = dpout[select]
boxes_xyxy = boxes_xyxy[select]
if classes is not None:
classes = classes[select]
return dpout, boxes_xywh, classes
def postprocess(self,
results,
output_height,
output_width,
padded_im_h,
padded_im_w,
mask_threshold=0.5):
"""
Resize the output instances.
The input images are often resized when entering an object detector.
As a result, we often need the outputs of the detector in a different
resolution from its inputs.
This function will resize the raw outputs of an R-CNN detector
to produce outputs according to the desired output resolution.
Args:
results (Instances): the raw outputs from the detector.
`results.image_size` contains the input image resolution the detector sees.
This object might be modified in-place.
output_height, output_width: the desired output resolution.
Returns:
Instances: the resized output from the model, based on the output resolution
"""
scale_x, scale_y = (output_width / results.image_size[1],
output_height / results.image_size[0])
resized_im_h, resized_im_w = results.image_size
results = Instances((output_height, output_width),
**results.get_fields())
if results.has("pred_boxes"):
output_boxes = results.pred_boxes
elif results.has("proposal_boxes"):
output_boxes = results.proposal_boxes
output_boxes.scale(scale_x, scale_y)
output_boxes.clip(results.image_size)
results = results[output_boxes.nonempty()]
if results.has("pred_global_masks"):
mask_h, mask_w = results.pred_global_masks.size()[-2:]
factor_h = padded_im_h // mask_h
factor_w = padded_im_w // mask_w
assert factor_h == factor_w
factor = factor_h
pred_global_masks = aligned_bilinear(results.pred_global_masks,
factor)
pred_global_masks = pred_global_masks[:, :, :resized_im_h, :
resized_im_w]
pred_global_masks = F.interpolate(pred_global_masks,
size=(output_height,
output_width),
mode="bilinear",
align_corners=False)
pred_global_masks = pred_global_masks[:, 0, :, :]
results.pred_masks = (pred_global_masks > mask_threshold).float()
return results
def execute_on_outputs(cls: type, context: Dict[str, Any],
entry: Dict[str, Any], outputs: Instances):
image_fpath = entry["file_name"]
logger.info(f"Processing {image_fpath}")
result = {"file_name": image_fpath}
if outputs.has("scores"):
result["scores"] = outputs.get("scores").cpu()
if outputs.has("pred_boxes"):
result["pred_boxes_XYXY"] = outputs.get("pred_boxes").tensor.cpu()
if outputs.has("pred_densepose"):
result["pred_densepose"], _ = DensePoseResultExtractor()(
outputs)
context["results"].append(result)
def execute_on_outputs(context: Dict[str, Any], entry: Dict[str, Any],
outputs: Instances):
image_fpath = entry["file_name"]
print(f"Processing {image_fpath}")
result = {"file_name": image_fpath}
if outputs.has("scores"):
result["scores"] = outputs.get("scores").cpu()
if outputs.has("pred_boxes"):
result["pred_boxes_XYXY"] = outputs.get("pred_boxes").tensor.cpu()
if outputs.has("pred_densepose"):
boxes_XYWH = BoxMode.convert(result["pred_boxes_XYXY"],
BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
result["pred_densepose"] = outputs.get("pred_densepose").to_result(
boxes_XYWH)
context["results"].append(result)
def extract_boxes_xywh_from_instances(instances: Instances, select=None):
if instances.has("pred_boxes"):
boxes_xywh = instances.pred_boxes.tensor.clone()
boxes_xywh[:, 2] -= boxes_xywh[:, 0]
boxes_xywh[:, 3] -= boxes_xywh[:, 1]
return boxes_xywh if select is None else boxes_xywh[select]
return None
def __call__(self, instances: Instances, select=None):
if instances.has("pred_densepose") and instances.has("pred_boxes"):
dpout = instances.pred_densepose
boxes_xyxy = instances.pred_boxes
boxes_xywh = extract_boxes_xywh_from_instances(instances)
if select is not None:
dpout = dpout[select]
boxes_xyxy = boxes_xyxy[select]
converter = ToChartResultConverterWithConfidences()
results = [
converter.convert(dpout[i], boxes_xyxy[[i]])
for i in range(len(dpout))
]
return results, boxes_xywh
else:
return None
def __call__(self, instances: Instances, select=None):
boxes_xywh = extract_boxes_xywh_from_instances(instances)
if instances.has("pred_keypoints") and (boxes_xywh is not None):
kpout = instances.pred_keypoints
kpout = kpout.cpu().numpy()
return kpout
else:
return None
def execute_on_outputs(
cls: type, context: Dict[str, Any], entry: Dict[str, Any], outputs: Instances
):
image_fpath = entry["file_name"]
logger.info(f"Processing {image_fpath}")
result = {"file_name": image_fpath}
if outputs.has("scores"):
result["scores"] = outputs.get("scores").cpu()
if outputs.has("pred_boxes"):
result["pred_boxes_XYXY"] = outputs.get("pred_boxes").tensor.cpu()
if outputs.has("pred_densepose"):
if isinstance(outputs.pred_densepose, DensePoseChartPredictorOutput):
extractor = DensePoseResultExtractor()
elif isinstance(outputs.pred_densepose, DensePoseEmbeddingPredictorOutput):
extractor = DensePoseOutputsExtractor()
result["pred_densepose"] = extractor(outputs)[0]
context["results"].append(result)
def __call__(self, instances: Instances, select=None):
boxes_xywh = extract_boxes_xywh_from_instances(instances)
if instances.has("pred_densepose") and (boxes_xywh is not None):
dpout = instances.pred_densepose
if select is not None:
dpout = dpout[select]
boxes_xywh = boxes_xywh[select]
return dpout.to_result(boxes_xywh)
else:
return None
def generate_instance(self, instance: Instances, class_names: List[str], total_classes: List[str]):
instance = instance.to('cpu')
boxes = instance.pred_boxes.tensor.numpy()
masks = None
scores = instance.scores.numpy()
if instance.has("pred_masks"):
masks = instance.pred_masks.numpy()
for index, name in enumerate(class_names):
if name not in total_classes:
boxes[index:index+1] = 0
scores[index] = 0
if masks is not None:
masks[index:index+1] = False
instance.pred_boxes = Boxes(torch.from_numpy(boxes))
if masks is not None:
instance.pred_masks = torch.from_numpy(masks)
instance.scores = torch.from_numpy(scores)
return instance
def extract_scores_from_instances(instances: Instances, select=None):
if instances.has("scores"):
return instances.scores if select is None else instances.scores[select]
return None