|
|
import albumentations as A |
|
|
import cv2 |
|
|
import torch |
|
|
|
|
|
from albumentations.pytorch import ToTensorV2 |
|
|
from utils.utils import seed_everything |
|
|
|
|
|
DATASET = "PASCAL_VOC" |
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
DEVICE_COUNT = torch.cuda.device_count() |
|
|
NUM_WORKERS = 0 |
|
|
BATCH_SIZE = 128 |
|
|
SHUFFLE = True |
|
|
IMAGE_SIZE = 416 |
|
|
NUM_CLASSES = 20 |
|
|
LEARNING_RATE = 1e-3 |
|
|
WEIGHT_DECAY = 1e-4 |
|
|
NUM_EPOCHS = 40 |
|
|
CONF_THRESHOLD = 0.05 |
|
|
MAP_IOU_THRESH = 0.5 |
|
|
NMS_IOU_THRESH = 0.45 |
|
|
S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8] |
|
|
PIN_MEMORY = True |
|
|
LOAD_MODEL = False |
|
|
SAVE_MODEL = True |
|
|
CHECKPOINT_FILE = "checkpoint.pth.tar" |
|
|
IMG_DIR = DATASET + "/images/" |
|
|
LABEL_DIR = DATASET + "/labels/" |
|
|
P_MOSAIC = 0.5 |
|
|
|
|
|
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)], |
|
|
] |
|
|
|
|
|
means = [0.485, 0.456, 0.406] |
|
|
|
|
|
scale = 1.1 |
|
|
train_transforms = A.Compose( |
|
|
[ |
|
|
A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)), |
|
|
A.PadIfNeeded( |
|
|
min_height=int(IMAGE_SIZE * scale), |
|
|
min_width=int(IMAGE_SIZE * scale), |
|
|
border_mode=cv2.BORDER_CONSTANT, |
|
|
), |
|
|
A.Rotate(limit=10, interpolation=1, border_mode=4), |
|
|
A.RandomCrop(width=IMAGE_SIZE, height=IMAGE_SIZE), |
|
|
A.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4), |
|
|
A.OneOf( |
|
|
[ |
|
|
A.ShiftScaleRotate( |
|
|
rotate_limit=20, p=0.5, border_mode=cv2.BORDER_CONSTANT |
|
|
), |
|
|
|
|
|
], |
|
|
p=1.0, |
|
|
), |
|
|
A.HorizontalFlip(p=0.5), |
|
|
A.Blur(p=0.1), |
|
|
A.CLAHE(p=0.1), |
|
|
A.Posterize(p=0.1), |
|
|
A.ToGray(p=0.1), |
|
|
A.ChannelShuffle(p=0.05), |
|
|
A.Normalize( |
|
|
mean=[0, 0, 0], |
|
|
std=[1, 1, 1], |
|
|
max_pixel_value=255, |
|
|
), |
|
|
ToTensorV2(), |
|
|
], |
|
|
bbox_params=A.BboxParams( |
|
|
format="yolo", |
|
|
min_visibility=0.4, |
|
|
label_fields=[], |
|
|
), |
|
|
) |
|
|
test_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(), |
|
|
], |
|
|
bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[]), |
|
|
) |
|
|
|
|
|
PASCAL_CLASSES = [ |
|
|
"aeroplane", |
|
|
"bicycle", |
|
|
"bird", |
|
|
"boat", |
|
|
"bottle", |
|
|
"bus", |
|
|
"car", |
|
|
"cat", |
|
|
"chair", |
|
|
"cow", |
|
|
"diningtable", |
|
|
"dog", |
|
|
"horse", |
|
|
"motorbike", |
|
|
"person", |
|
|
"pottedplant", |
|
|
"sheep", |
|
|
"sofa", |
|
|
"train", |
|
|
"tvmonitor", |
|
|
] |
|
|
|
|
|
COCO_LABELS = [ |
|
|
"person", |
|
|
"bicycle", |
|
|
"car", |
|
|
"motorcycle", |
|
|
"airplane", |
|
|
"bus", |
|
|
"train", |
|
|
"truck", |
|
|
"boat", |
|
|
"traffic light", |
|
|
"fire hydrant", |
|
|
"stop sign", |
|
|
"parking meter", |
|
|
"bench", |
|
|
"bird", |
|
|
"cat", |
|
|
"dog", |
|
|
"horse", |
|
|
"sheep", |
|
|
"cow", |
|
|
"elephant", |
|
|
"bear", |
|
|
"zebra", |
|
|
"giraffe", |
|
|
"backpack", |
|
|
"umbrella", |
|
|
"handbag", |
|
|
"tie", |
|
|
"suitcase", |
|
|
"frisbee", |
|
|
"skis", |
|
|
"snowboard", |
|
|
"sports ball", |
|
|
"kite", |
|
|
"baseball bat", |
|
|
"baseball glove", |
|
|
"skateboard", |
|
|
"surfboard", |
|
|
"tennis racket", |
|
|
"bottle", |
|
|
"wine glass", |
|
|
"cup", |
|
|
"fork", |
|
|
"knife", |
|
|
"spoon", |
|
|
"bowl", |
|
|
"banana", |
|
|
"apple", |
|
|
"sandwich", |
|
|
"orange", |
|
|
"broccoli", |
|
|
"carrot", |
|
|
"hot dog", |
|
|
"pizza", |
|
|
"donut", |
|
|
"cake", |
|
|
"chair", |
|
|
"couch", |
|
|
"potted plant", |
|
|
"bed", |
|
|
"dining table", |
|
|
"toilet", |
|
|
"tv", |
|
|
"laptop", |
|
|
"mouse", |
|
|
"remote", |
|
|
"keyboard", |
|
|
"cell phone", |
|
|
"microwave", |
|
|
"oven", |
|
|
"toaster", |
|
|
"sink", |
|
|
"refrigerator", |
|
|
"book", |
|
|
"clock", |
|
|
"vase", |
|
|
"scissors", |
|
|
"teddy bear", |
|
|
"hair drier", |
|
|
"toothbrush", |
|
|
] |
|
|
|
