| from typing import Union | |
| from pathlib import Path | |
| import torch | |
| from torchvision.models.segmentation import deeplabv3_resnet101 | |
| from SoccerNet.Evaluation.utils_calibration import SoccerPitch | |
| class InferenceSegmentationModel: | |
| def __init__(self, checkpoint: Union[str, Path], device) -> None: | |
| self.device = device | |
| self.model = deeplabv3_resnet101( | |
| num_classes=len(SoccerPitch.lines_classes) + 1, aux_loss=True | |
| ) | |
| checkpoint_data = torch.load(checkpoint, map_location=self.device, weights_only=False) | |
| self.model.load_state_dict(checkpoint_data["model"], strict=False) | |
| self.model.to(self.device) | |
| self.model.eval() | |
| def inference(self, img_batch): | |
| return self.model(img_batch)["out"].argmax(1) | |