Spaces:
Runtime error
Runtime error
File size: 3,130 Bytes
a00793d |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import albumentations as A
import cv2
import torch
from albumentations.pytorch import ToTensorV2
from utils import seed_everything
DATASET = 'PASCAL_VOC'
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# seed_everything() # If you want deterministic behavior
NUM_WORKERS = 0
BATCH_SIZE = 2
DIV = 32
IMAGE_SIZES = [416, 416, 416, 608, 608]
S = [[x//DIV, x//DIV*2, x//DIV*4] for x in IMAGE_SIZES]
NUM_CLASSES = 20
LEARNING_RATE = 1e-5
WEIGHT_DECAY = 1e-4
NUM_EPOCHS = 10
CONF_THRESHOLD = 0.05
MAP_IOU_THRESH = 0.5
NMS_IOU_THRESH = 0.45
PIN_MEMORY = True
LOAD_MODEL = False
SAVE_MODEL = True
CHECKPOINT_FILE = "checkpoint.pth.tar"
IMG_DIR = DATASET + "/images/"
LABEL_DIR = DATASET + "/labels/"
MOSAIC_PROB = 0.75
INFERENCE_IMAGE_SIZE = 416
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)],
] # Note these have been rescaled to be between [0, 1]
means = [0.45484068, 0.43406072, 0.40103856]
stds = [0.23936155, 0.23471538, 0.23876129]
scale = 1.1
def train_transform(IMAGE_SIZE):
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
),
# A.Affine(shear=15, p=0.5, mode="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=means, std=stds, max_pixel_value=255,),
ToTensorV2(),
],
bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[],),
)
return(train_transforms)
def test_transform(IMAGE_SIZE=416):
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=means, std=stds, max_pixel_value=255,),
ToTensorV2(),
],
bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[]),
)
return(test_transforms)
PASCAL_CLASSES = [
"aeroplane",
"bicycle",
"bird",
"boat",
"bottle",
"bus",
"car",
"cat",
"chair",
"cow",
"diningtable",
"dog",
"horse",
"motorbike",
"person",
"pottedplant",
"sheep",
"sofa",
"train",
"tvmonitor"
] |