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	Optional	[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])]

/mnt/980_1TB_1/Projects/GitHub/cjm-yolox-pytorch/cjm_yolox_pytorch/model.py:792: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  state_dict = torch.load(checkpoint_path, map_location='cpu')

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])