Skip to content

part_refinement

Nearest-neighbor refinement of part labels on top of another inferer's output.

Takes the per-point argmax of a base inferer and re-assigns the labels that are implausible for the shape (rare parts, parts the shape's category does not own) by a majority vote of each point's nearest neighbors.

Classes:

Functions:

PartRefinementInferer [source] 

PartRefinementInferer(
    base: Inferer,
    part_ids: Optional[Sequence[Sequence[int]]] = None,
    min_count: int = 10,
    num_neighbors: int = 11,
    category_key: str = CATEGORY,
    pos_key: str = POS,
    batch_key: str = BATCH,
)

Bases: Inferer

Nearest-neighbor refinement of part labels on top of another inferer.

Takes the base inferer's per-point argmax and re-assigns the implausible part labels (rare parts, parts the shape's category does not own) by a nearest-neighbor majority vote, returning one-hot scores.

All parameters are forwarded verbatim to part_refinement_inference.

Example 
import torch

import torch_pointcloud as tp
from torch_pointcloud.inferers import PartRefinementInferer, SimpleInferer

model = tp.create_model("dgcnn.shapenetpart.an-tao").eval()
shapes = {
    "pos": torch.rand(2048, 3),
    "batch": torch.zeros(2048, dtype=torch.long),
    "category": torch.nn.functional.one_hot(torch.tensor([0]), num_classes=16).float(),
}

inferer = PartRefinementInferer(SimpleInferer())
scores = inferer(shapes, predictor=lambda d: model(None, d["pos"], d["batch"], d["category"]))
labels = scores.argmax(dim=1)

part_refinement_inference [source] 

part_refinement_inference(
    data: Dict[str, Any],
    *,
    predictor: Callable[[Dict[str, Any]], Tensor],
    base: Optional[Inferer] = None,
    part_ids: Optional[Sequence[Sequence[int]]] = None,
    min_count: int = 10,
    num_neighbors: int = 11,
    category_key: str = CATEGORY,
    pos_key: str = POS,
    batch_key: str = BATCH,
) -> Tensor

Nearest-neighbor refinement of part labels on top of another inferer.

Runs base, takes the per-point argmax, and for every shape re-assigns the labels that are implausible: a predicted part with fewer than min_count points, or a part the shape's category does not own. Each such label is refined in turn (ascending label order): its points take the majority label of their num_neighbors nearest points of the same shape, the label under refinement excluded from the vote, with the already-refined labels feeding the next votes. This is the post-processing of the PointNeXt ShapeNetPart protocol.

Parameters:

  • data (Dict[str, Any]) –

    Dict of per-point tensors. Must contain pos (shape \((N, D)\)), batch (shape \((N,)\)) and the per-shape category under category_key.

  • predictor (Callable[[Dict[str, Any]], Tensor]) –

    Callable mapping a data dict to per-point part scores of shape \((N, C)\).

  • base (Optional[Inferer], default: None ) –

    Inferer running predictor; defaults to SimpleInferer (one forward on the whole batch).

  • part_ids (Optional[Sequence[Sequence[int]]], default: None ) –

    Part labels owned by each category; defaults to the 16-category / 50-part ShapeNetPart table (ShapeNetPart.seg_ids).

  • min_count (int, default: 10 ) –

    Predicted parts with fewer points than this are refined.

  • num_neighbors (int, default: 11 ) –

    Number of nearest neighbors (the point itself included) voting on the new label.

  • category_key (str, default: CATEGORY ) –

    Dict key of the per-shape category, one-hot \((B, K)\) or index \((B,)\).

  • pos_key (str, default: POS ) –

    Dict key for the position tensor.

  • batch_key (str, default: BATCH ) –

    Dict key for the per-point batch index.

Returns:

  • Tensor –

    One-hot refined labels of shape \((N, C)\), so a metric's argmax recovers the refined labels. An empty

  • Tensor –

    scene (\(N = 0\)) returns the base output unchanged.