Object Detection with LTDETRv2¶
This notebook demonstrates how to use LightlyTrain for object detection with our state-of-the-art LTDETRv2 model.
Important: When running on Google Colab make sure to select a GPU runtime for faster processing. You can do this by going to
Runtime>Change runtime typeand selecting a GPU hardware accelerator.
Installation¶
LightlyTrain can be installed directly via pip:
!pip install lightly-train
Important: LightlyTrain is officially supported on
Linux: CPU or CUDA
MacOS: CPU only
Windows (experimental): CPU or CUDA
We are planning to support MPS for MacOS.
Check the installation instructions for more details on installation.
Train an LTDETR model¶
Training your own LTDETR model is straightforward with LightlyTrain.
Download a dataset¶
First download a dataset. The dataset must be in YOLO format, see the documentation for more details. You can use labelformat to convert any dataset to the YOLO format.
!wget https://github.com/lightly-ai/coco128_yolo/releases/download/v0.0.1/coco128_yolo.zip && unzip -q coco128_yolo.zip
The dataset looks like this after the download completes:
coco128_yolo
├── images
│ ├── train2017
│ │ ├── 000000000009.jpg
│ │ ├── 000000000025.jpg
│ │ ├── ...
│ │ └── 000000000650.jpg
│ └── val2017
│ ├── 000000000139.jpg
│ ├── 000000000285.jpg
│ ├── ...
│ └── 000000013201.jpg
└── labels
├── train2017
│ ├── 000000000009.txt
│ ├── 000000000025.txt
│ ├── ...
│ └── 000000000659.txt
└── val2017
├── 000000000139.txt
├── 000000000285.txt
├── ...
└── 000000013201.txt
Start training¶
Next, start the training with the train_object_detection function. You only have to specify the output directory, model, and input data. LightlyTrain automatically sets the remaining training parameters and applies image augmentations. Of course you can always customize these settings if needed.
import lightly_train
lightly_train.train_object_detection(
out="out/my_experiment",
model="ltdetrv2-s-coco",
steps=100, # Small number of steps for demonstration, default is 266_112.
batch_size=4, # Small batch size for demonstration, default is 32.
data={
"path": "coco128_yolo",
"train": "images/train2017",
"val": "images/val2017",
"names": {
0: "person",
1: "bicycle",
2: "car",
3: "motorcycle",
4: "airplane",
5: "bus",
6: "train",
7: "truck",
8: "boat",
9: "traffic light",
10: "fire hydrant",
11: "stop sign",
12: "parking meter",
13: "bench",
14: "bird",
15: "cat",
16: "dog",
17: "horse",
18: "sheep",
19: "cow",
20: "elephant",
21: "bear",
22: "zebra",
23: "giraffe",
24: "backpack",
25: "umbrella",
26: "handbag",
27: "tie",
28: "suitcase",
29: "frisbee",
30: "skis",
31: "snowboard",
32: "sports ball",
33: "kite",
34: "baseball bat",
35: "baseball glove",
36: "skateboard",
37: "surfboard",
38: "tennis racket",
39: "bottle",
40: "wine glass",
41: "cup",
42: "fork",
43: "knife",
44: "spoon",
45: "bowl",
46: "banana",
47: "apple",
48: "sandwich",
49: "orange",
50: "broccoli",
51: "carrot",
52: "hot dog",
53: "pizza",
54: "donut",
55: "cake",
56: "chair",
57: "couch",
58: "potted plant",
59: "bed",
60: "dining table",
61: "toilet",
62: "tv",
63: "laptop",
64: "mouse",
65: "remote",
66: "keyboard",
67: "cell phone",
68: "microwave",
69: "oven",
70: "toaster",
71: "sink",
72: "refrigerator",
73: "book",
74: "clock",
75: "vase",
76: "scissors",
77: "teddy bear",
78: "hair drier",
79: "toothbrush",
},
},
)
Once the training is complete, the output directory looks like this:
out/my_experiment
├── checkpoints
│ ├── best.ckpt
│ └── last.ckpt
├── events.out.tfevents.1764251158.ef9b159fe4b8.273.0
├── exported_models
│ ├── exported_best.pt
│ └── exported_last.pt
└── train.log
The best model checkpoint is saved to out/my_experiment/exported_models/exported_best.pt.
Predict with your own LTDETR weights¶
Download an example image¶
Download an example image for inference with the following command:
!wget -O image.jpg http://images.cocodataset.org/val2017/000000577932.jpg
Then load the model with LightlyTrain’s load_model function. This will automatically download the model weights and load the model.
Load the model weights¶
model = lightly_train.load_model("out/my_experiment/exported_models/exported_best.pt")
Predict the objects¶
Then run model.predict on the image. The method accepts file paths, URLs,
PIL Images, or tensors as input.
results = model.predict("image.jpg")
results["labels"] # Class labels, tensor of shape (num_boxes,)
results["bboxes"] # Bounding boxes in (xmin, ymin, xmax, ymax) absolute pixel
# coordinates of the original image. Tensor of shape (num_boxes, 4).
results["scores"] # Confidence scores, tensor of shape (num_boxes,)
Visualize the results¶
Finally, we visualize the image and results to check what objects were detected.
import matplotlib.pyplot as plt
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes
image = read_image("image.jpg")
image_with_boxes = draw_bounding_boxes(
image,
boxes=results["bboxes"],
labels=[model.classes[label.item()] for label in results["labels"]],
)
plt.imshow(image_with_boxes.permute(1, 2, 0))
plt.show()
Benchmark your checkpoint¶
You can use the benchmark_object_detection (in beta) command to measure the inference performance of an object detection model on a validation dataset.
result = lightly_train.benchmark_object_detection(
out="out/my_benchmark",
dataset_name="coco128", # Human-readable name shown in the report.
model="out/my_experiment/exported_models/exported_best.pt",
data={
"path": "coco128_yolo",
"train": "images/train2017",
"val": "images/val2017",
"names": {
0: "person",
1: "bicycle",
2: "car",
3: "motorcycle",
4: "airplane",
5: "bus",
6: "train",
7: "truck",
8: "boat",
9: "traffic light",
10: "fire hydrant",
11: "stop sign",
12: "parking meter",
13: "bench",
14: "bird",
15: "cat",
16: "dog",
17: "horse",
18: "sheep",
19: "cow",
20: "elephant",
21: "bear",
22: "zebra",
23: "giraffe",
24: "backpack",
25: "umbrella",
26: "handbag",
27: "tie",
28: "suitcase",
29: "frisbee",
30: "skis",
31: "snowboard",
32: "sports ball",
33: "kite",
34: "baseball bat",
35: "baseball glove",
36: "skateboard",
37: "surfboard",
38: "tennis racket",
39: "bottle",
40: "wine glass",
41: "cup",
42: "fork",
43: "knife",
44: "spoon",
45: "bowl",
46: "banana",
47: "apple",
48: "sandwich",
49: "orange",
50: "broccoli",
51: "carrot",
52: "hot dog",
53: "pizza",
54: "donut",
55: "cake",
56: "chair",
57: "couch",
58: "potted plant",
59: "bed",
60: "dining table",
61: "toilet",
62: "tv",
63: "laptop",
64: "mouse",
65: "remote",
66: "keyboard",
67: "cell phone",
68: "microwave",
69: "oven",
70: "toaster",
71: "sink",
72: "refrigerator",
73: "book",
74: "clock",
75: "vase",
76: "scissors",
77: "teddy bear",
78: "hair drier",
79: "toothbrush",
},
},
)
result.print()
The command returns a BenchmarkResult instance and writes two files to the out directory:
out/my_benchmark
├── benchmark_results.json
└── benchmark_summary.md
benchmark_results.json: the full result as JSON.benchmark_summary.md: a human-readable Markdown report.
The report (also available via result.to_markdown()) contains the run configuration, device info, performance metrics, and a throughput & latency table.
Next Steps¶
Object Detection Documentation: If you want to learn more about object detection with LightlyTrain.
Object Detection Model Export: If you want to learn more about export the model to ONNX and TensorRT format.