cjm-yolox-pytorch

A PyTorch implementation of the YOLOX object detection model based on OpenMMLab’s implementation in the mmdetection library.

Install

pip install cjm_yolox_pytorch

How to use

import torch
from cjm_yolox_pytorch.model import MODEL_TYPES, build_model

Select model type

model_type = MODEL_TYPES[0]
model_type

'yolox_tiny'

Build YOLOX model

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

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

with torch.no_grad():
    cls_scores, bbox_preds, objectness = yolox(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])]