Skip to content

README.md

Latest commit

 

History

History
54 lines (46 loc) · 997 Bytes

README.md

File metadata and controls

54 lines (46 loc) · 997 Bytes

Image Classification

Overview

Image classification using TensorRT for optimized inference.

Export Model

  1. Export your PyTorch/TensorFlow model to ONNX:
import torch
model = torch.load("model.pt")
torch.onnx.export(model, ...)
  1. Convert ONNX to TensorRT engine:
mkdir data
trtexec --onnx=data/model.onnx --saveEngine=data/model.engine --fp16

Configure

In data folder, add your config.json:

{
  "engine": {
    "model_path": "./data/model.engine",
    "batch_size": 1,
    "precision": 16
  },
  "confidence_threshold": 0.5,
  "class_names": ["class1", "class2"]
}

Compile

# in root directory
meson setup build -Dbuild_apps=classifier
meson compile -C build

Run

Display

# in root directory
cd build/app/classifier
./classify -i image.jpg -c data/config.json -d

JQuery Pipeline

# in root directory
cd build/app/classifier
./classify -i image.jpg -c data/config.json | jq .data.class_name