Update src/run/yolov3/config.py

src/run/yolov3/config.py

@@ -1,8 +1,9 @@
 import albumentations as A
 import cv2
 import torch
-from albumentations.pytorch import ToTensorV2
 
+from albumentations.pytorch import ToTensorV2
+from utils import seed_everything
 
 DATASET = "PASCAL_VOC"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -10,6 +11,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 NUM_WORKERS = 0
 BATCH_SIZE = 32
 IMAGE_SIZE = 416
+MAX_IMAGE_SIZE = 416
 NUM_CLASSES = 20
 LEARNING_RATE = 1e-5
 WEIGHT_DECAY = 1e-4
@@ -18,9 +20,11 @@ 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/"
@@ -31,7 +35,8 @@ ANCHORS = [
     [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
 ]  # Note these have been rescaled to be between [0, 1]
 
-means = [0.485, 0.456, 0.406]
+mean = [0.485, 0.456, 0.406]
+std = [0.229, 0.224, 0.225]
 
 scale = 1.1
 train_transforms = A.Compose(
@@ -61,8 +66,8 @@ train_transforms = A.Compose(
         A.ToGray(p=0.1),
         A.ChannelShuffle(p=0.05),
         A.Normalize(
-            mean=[0, 0, 0],
-            std=[1, 1, 1],
+            mean=mean,
+            std=std,
             max_pixel_value=255,
         ),
         ToTensorV2(),
@@ -80,8 +85,8 @@ test_transforms = A.Compose(
             min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
         ),
         A.Normalize(
-            mean=[0, 0, 0],
-            std=[1, 1, 1],
+            mean=mean,
+            std=std,
             max_pixel_value=255,
         ),
         ToTensorV2(),
@@ -90,6 +95,7 @@ test_transforms = A.Compose(
 )
 
 PASCAL_CLASSES = [
+    # "background",
     "aeroplane",
     "bicycle",
     "bird",
@@ -193,4 +199,4 @@ COCO_LABELS = [
     "teddy bear",
     "hair drier",
     "toothbrush",
-]
+]