|
import albumentations as A |
|
import cv2 |
|
import torch |
|
|
|
from albumentations.pytorch import ToTensorV2 |
|
|
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
IMAGE_SIZE = 416 |
|
transforms = A.Compose( |
|
[ |
|
A.LongestMaxSize(max_size=IMAGE_SIZE), |
|
A.PadIfNeeded( |
|
min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT |
|
), |
|
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,), |
|
ToTensorV2(), |
|
], |
|
) |
|
ANCHORS = [ |
|
[(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)], |
|
[(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)], |
|
[(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)], |
|
] |
|
S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8] |
|
|
|
PASCAL_CLASSES = [ |
|
"aeroplane", |
|
"bicycle", |
|
"bird", |
|
"boat", |
|
"bottle", |
|
"bus", |
|
"car", |
|
"cat", |
|
"chair", |
|
"cow", |
|
"diningtable", |
|
"dog", |
|
"horse", |
|
"motorbike", |
|
"person", |
|
"pottedplant", |
|
"sheep", |
|
"sofa", |
|
"train", |
|
"tvmonitor" |
|
] |
|
|