/
coco_detr.py
70 lines (64 loc) · 1.99 KB
/
coco_detr.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
from omegaconf import OmegaConf
import detectron2.data.transforms as T
from detectron2.config import LazyCall as L
from detectron2.data import (
build_detection_test_loader,
build_detection_train_loader,
get_detection_dataset_dicts,
)
from detectron2.evaluation import COCOEvaluator
from detrex.data import DetrDatasetMapper
dataloader = OmegaConf.create()
dataloader.train = L(build_detection_train_loader)(
dataset=L(get_detection_dataset_dicts)(names="coco_2017_train"),
mapper=L(DetrDatasetMapper)(
augmentation=[
L(T.RandomFlip)(),
L(T.ResizeShortestEdge)(
short_edge_length=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800),
max_size=1333,
sample_style="choice",
),
],
augmentation_with_crop=[
L(T.RandomFlip)(),
L(T.ResizeShortestEdge)(
short_edge_length=(400, 500, 600),
sample_style="choice",
),
L(T.RandomCrop)(
crop_type="absolute_range",
crop_size=(384, 600),
),
L(T.ResizeShortestEdge)(
short_edge_length=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800),
max_size=1333,
sample_style="choice",
),
],
is_train=True,
mask_on=False,
img_format="RGB",
),
total_batch_size=16,
num_workers=4,
)
dataloader.test = L(build_detection_test_loader)(
dataset=L(get_detection_dataset_dicts)(names="coco_2017_val", filter_empty=False),
mapper=L(DetrDatasetMapper)(
augmentation=[
L(T.ResizeShortestEdge)(
short_edge_length=800,
max_size=1333,
),
],
augmentation_with_crop=None,
is_train=False,
mask_on=False,
img_format="RGB",
),
num_workers=4,
)
dataloader.evaluator = L(COCOEvaluator)(
dataset_name="${..test.dataset.names}",
)