ManjunathReddy/Yolo3_from_scratch

1import albumentations as A2import cv23import torch4 5from albumentations.pytorch import ToTensorV26 7 8DEVICE = "cuda" if torch.cuda.is_available() else "cpu"9 10DATASET = 'PASCAL_VOC'11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"12# seed_everything()  # If you want deterministic behavior13NUM_WORKERS = 014BATCH_SIZE = 3215IMAGE_SIZE = 41616NUM_CLASSES = 2017LEARNING_RATE = 1e-518WEIGHT_DECAY = 1e-419NUM_EPOCHS = 10020CONF_THRESHOLD = 0.0521MAP_IOU_THRESH = 0.522NMS_IOU_THRESH = 0.4523S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]24PIN_MEMORY = True25LOAD_MODEL = False26SAVE_MODEL = True27CHECKPOINT_FILE = "checkpoint.pth.tar"28IMG_DIR = DATASET + "/images/"29LABEL_DIR = DATASET + "/labels/"30 31means = [0.485, 0.456, 0.406]32 33scale = 1.134 35 36train_transforms = A.Compose(37    [38        A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)),39        A.PadIfNeeded(40            min_height=int(IMAGE_SIZE * scale),41            min_width=int(IMAGE_SIZE * scale),42            border_mode=cv2.BORDER_CONSTANT,43        ),44        A.Rotate(limit = 10, interpolation=1, border_mode=4),45        A.RandomCrop(width=IMAGE_SIZE, height=IMAGE_SIZE),46        A.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),47        A.OneOf(48            [49                A.ShiftScaleRotate(50                    rotate_limit=20, p=0.5, border_mode=cv2.BORDER_CONSTANT51                ),52                # A.Affine(shear=15, p=0.5, mode="constant"),53            ],54            p=1.0,55        ),56        A.HorizontalFlip(p=0.5),57        A.Blur(p=0.1),58        A.CLAHE(p=0.1),59        A.Posterize(p=0.1),60        A.ToGray(p=0.1),61        A.ChannelShuffle(p=0.05),62        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),63        ToTensorV2(),64    ],65    bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[],),66)67test_transforms = A.Compose(68    [69        A.LongestMaxSize(max_size=IMAGE_SIZE),70        A.PadIfNeeded(71            min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT72        ),73        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),74        ToTensorV2(),75    ],76    bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[]),77)78 79 80 81IMAGE_SIZE = 41682transforms = A.Compose(83    [84        A.LongestMaxSize(max_size=IMAGE_SIZE),85        A.PadIfNeeded(86            min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT87        ),88        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),89        ToTensorV2(),90    ],91)92ANCHORS = [93    [(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],94    [(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],95    [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],96]  # Note these have been rescaled to be between [0, 1]97S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]98 99PASCAL_CLASSES = [100    "aeroplane",101    "bicycle",102    "bird",103    "boat",104    "bottle",105    "bus",106    "car",107    "cat",108    "chair",109    "cow",110    "diningtable",111    "dog",112    "horse",113    "motorbike",114    "person",115    "pottedplant",116    "sheep",117    "sofa",118    "train",119    "tvmonitor"120]121