def save_state(
self,
target_dir: Path,
name: str,
**kwargs,
):
"""
Save case result as pickle file. Identifier of ensembler will
be added to the name.
This version only saves the topk model predictions to speed
up loading.
Args:
target_dir: folder to save result to
name: name of case
Notes:
The device is not saved inside the checkpoint and everything
will be loaded on the CPU.
"""
for model in self.model_results.keys():
boxes = cat(self.model_results[model]["boxes"])
probs = cat(self.model_results[model]["scores"])
labels = cat(self.model_results[model]["labels"])
weights = cat(self.model_results[model]["weights"])
if len(probs) > self.parameters["model_topk"]:
_, idx_sorted = probs.sort(descending=True)
idx_sorted = idx_sorted[:self.parameters["model_topk"]]
self.model_results[model]["boxes"] = boxes[idx_sorted]
self.model_results[model]["scores"] = probs[idx_sorted]
self.model_results[model]["labels"] = labels[idx_sorted]
self.model_results[model]["weights"] = weights[idx_sorted]
return super().save_state(target_dir=target_dir, name=name, **kwargs)
def process_model(self, name: Hashable) ->\
Tuple[Tensor, Tensor, Tensor, Tensor]:
"""
Process the output of a single model on the whole scan
topk candidates -> nms
Args:
name: name of model to process
Returns:
Tensor: processed boxes
Tensor: processed probs
Tensor: processed labels
Tensor: processed weights
"""
# collect predictions on whole case and apply postprocessing
boxes = cat(self.model_results[name]["boxes"]).to(self.device)
probs = cat(self.model_results[name]["scores"]).to(self.device)
labels = cat(self.model_results[name]["labels"]).to(self.device)
weights = cat(self.model_results[name]["weights"]).to(self.device)
return self.postprocess_image(
boxes=boxes,
probs=probs,
labels=labels,
weights=weights,
shape=tuple(self.properties["shape"]),
)
def process_ensemble(
self,
boxes: List[Tensor],
probs: List[Tensor],
labels: List[Tensor],
weights: List[Tensor],
) -> Tuple[Tensor, Tensor, Tensor]:
"""
Ensemble predictions from multiple models
Args:
boxes: predicted boxes List[[N, dims * 2]]
(x1, y1, x2, y2, (z1, z2))
probs: predicted probabilities List[[N]]
labels: predicted label List[[N]]
weights: additional weight List[[N]]
Returns:
Tensor: ensembled box predictions
Tensor: ensembled probabilities
Tensor: ensembled labels
"""
num_models = len(boxes)
boxes = cat(boxes, dim=0)
probs = cat(probs, dim=0)
labels = cat(labels, dim=0)
weights = cat(weights, dim=0)
_, idx = probs.sort(descending=True)
idx = idx[:self.parameters["ensemble_topk"]]
boxes = boxes[idx]
probs = probs[idx]
labels = labels[idx]
weights = weights[idx]
n_exp_preds = torch.tensor([num_models] * len(boxes)).to(boxes)
boxes, probs, labels = self.parameters["ensemble_nms_fn"](
boxes,
probs,
labels,
weights=weights,
iou_thresh=self.parameters["ensemble_iou"],
n_exp_preds=n_exp_preds,
score_thresh=self.parameters["ensemble_score_thresh"],
)
return boxes.cpu(), probs.cpu(), labels.cpu()
def process_model(
self,
name: Hashable,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
"""
Process the output of a single model on the whole scan
topk candidates -> nms
Args:
name: name of model to process
Returns:
Tensor: processed boxes
Tensor: processed probs
Tensor: processed labels
Tensor: processed weights
"""
boxes = to_device(self.model_results[name]["boxes"],
device=self.device)
probs = to_device(self.model_results[name]["scores"],
device=self.device)
labels = to_device(self.model_results[name]["labels"],
device=self.device)
weights = to_device(self.model_results[name]["weights"],
device=self.device)
model_boxes = []
model_probs = []
model_labels = []
model_weights = []
for b, p, l, w in zip(boxes, probs, labels, weights):
if b.numel() > 0:
_b, _p, _l, _w = self.postprocess_image(
boxes=b.float(),
probs=p.float(),
labels=l.float(),
weights=w.float(),
shape=tuple(self.properties["shape"]),
)
model_boxes.append(_b)
model_probs.append(_p)
model_labels.append(_l)
model_weights.append(_w)
return cat(model_boxes), cat(model_probs), cat(model_labels), cat(
model_weights)
def process_model(self, name: Hashable) ->\
Tuple[Tensor, Tensor, Tensor, Tensor]:
"""
Process the output of a single model on the whole scan
topk candidates -> nms
Args:
name: name of model to process
Returns:
Tensor: filtered boxes
Tensor: filtered probs
Tensor: filtered labels
idx: indices kept from original ordered data
"""
# concatenate batches
boxes = cat(self.model_results[name]["boxes"], dim=0)
probs = cat(self.model_results[name]["scores"], dim=0)
labels = cat(self.model_results[name]["labels"], dim=0)
weights = cat(self.model_results[name]["weights"], dim=0)
return boxes, probs, labels, weights
def reduce_cache(self):
"""
Only save a subset of all boxes for further evaluations
"""
if not self.reduced_cache:
self.reduced_cache = True
# we use the mean here to save time ...
self.overlap_map_mean = self.overlap_map.avg()
for model in self.model_results.keys():
batch_idx = self.build_batch_indices(
self.model_results[model]["scores"])
boxes = cat(self.model_results[model]["boxes"])
probs = cat(self.model_results[model]["scores"])
labels = cat(self.model_results[model]["labels"])
weights = cat(self.model_results[model]["weights"])
if len(probs) > self.num_reduced_cache:
_, idx_sorted = probs.sort(descending=True)
idx_sorted = idx_sorted[:self.num_reduced_cache]
batch_idx_keep = [[b for b in bix if b in idx_sorted]
for bix in batch_idx]
assert len(batch_idx_keep) == len(
self.model_results[model]["scores"])
self.model_results[model]["boxes"] = [
boxes[i] for i in batch_idx_keep
]
self.model_results[model]["scores"] = [
probs[i] for i in batch_idx_keep
]
self.model_results[model]["labels"] = [
labels[i] for i in batch_idx_keep
]
self.model_results[model]["weights"] = [
weights[i] for i in batch_idx_keep
]