inference

---

source

YOLOXInferenceWrapper

 YOLOXInferenceWrapper (model:torch.nn.modules.module.Module,
                        normalize_mean:torch.Tensor=tensor([[[[0.]],
                        [[0.]],           [[0.]]]]),
                        normalize_std:torch.Tensor=tensor([[[[1.]],
                        [[1.]],           [[1.]]]]),
                        strides:Optional[List[int]]=[8, 16, 32],
                        scale_inp:bool=False, channels_last:bool=False,
                        run_box_and_prob_calculation:bool=True)

This is a wrapper for the YOLOX https://arxiv.org/abs/2107.08430 object detection model. The class handles preprocessing of the input, postprocessing of the model output, and calculation of bounding boxes and their probabilities.

	Type	Default	Details
model	Module		The YOLOX model.
normalize_mean	Tensor	tensor([[[[0.]], [[0.]], [[0.]]]])	The mean values for normalization.
normalize_std	Tensor	tensor([[[[1.]], [[1.]], [[1.]]]])	The standard deviation values for normalization.
strides	typing.Optional[typing.List[int]]	[8, 16, 32]	The strides for the model.
scale_inp	bool	False	Whether to scale the input by dividing by 255.
channels_last	bool	False	Whether the input tensor has channels first.
run_box_and_prob_calculation	bool	True	Whether to calculate the bounding boxes and their probabilities.

model_type = 'yolox_tiny'

model = build_model(model_type, 19, pretrained=True)

test_inp = torch.randn(1, 3, 256, 256)

with torch.no_grad():
    cls_scores, bbox_preds, objectness = model(test_inp)
    
print(f"cls_scores: {[cls_score.shape for cls_score in cls_scores]}")
print(f"bbox_preds: {[bbox_pred.shape for bbox_pred in bbox_preds]}")
print(f"objectness: {[objectness.shape for objectness in objectness]}")

The file ./pretrained_checkpoints/yolox_tiny.pth already exists and overwrite is set to False.
cls_scores: [torch.Size([1, 19, 32, 32]), torch.Size([1, 19, 16, 16]), torch.Size([1, 19, 8, 8])]
bbox_preds: [torch.Size([1, 4, 32, 32]), torch.Size([1, 4, 16, 16]), torch.Size([1, 4, 8, 8])]
objectness: [torch.Size([1, 1, 32, 32]), torch.Size([1, 1, 16, 16]), torch.Size([1, 1, 8, 8])]

norm_stats = [*NORM_STATS[model_type].values()]

# Convert the normalization stats to tensors
mean_tensor = torch.tensor(norm_stats[0]).view(1, 3, 1, 1)
std_tensor = torch.tensor(norm_stats[1]).view(1, 3, 1, 1)

# Set the model to evaluation mode
model.eval();

# Wrap the model with preprocessing and post-processing steps
wrapped_model = YOLOXInferenceWrapper(model, 
                                      mean_tensor, 
                                      std_tensor, 
                                      scale_inp=False, 
                                      channels_last=False)

with torch.no_grad():
    model_output = wrapped_model(test_inp)
model_output.shape

torch.Size([1, 1344, 6])