mirror of
https://github.com/aladdinpersson/Machine-Learning-Collection.git
synced 2026-02-21 11:18:01 +00:00
182 lines
3.5 KiB
Python
182 lines
3.5 KiB
Python
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 = 4
|
|
BATCH_SIZE = 32
|
|
IMAGE_SIZE = 416
|
|
NUM_CLASSES = 80
|
|
LEARNING_RATE = 3e-5
|
|
WEIGHT_DECAY = 1e-4
|
|
NUM_EPOCHS = 100
|
|
CONF_THRESHOLD = 0.6
|
|
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 = True
|
|
SAVE_MODEL = True
|
|
CHECKPOINT_FILE = "checkpoint.pth.tar"
|
|
IMG_DIR = DATASET + "/images/"
|
|
LABEL_DIR = DATASET + "/labels/"
|
|
|
|
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]
|
|
|
|
|
|
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.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=10, p=0.4, border_mode=cv2.BORDER_CONSTANT
|
|
),
|
|
A.IAAAffine(shear=10, p=0.4, 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=[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'
|
|
] |