Spaces:
Sleeping
Sleeping
Upload 2 files
Browse files
config.py
ADDED
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import albumentations as A
|
2 |
+
import cv2
|
3 |
+
import torch
|
4 |
+
import os
|
5 |
+
|
6 |
+
from albumentations.pytorch import ToTensorV2
|
7 |
+
#from utils import seed_everything
|
8 |
+
from pytorch_lightning import LightningModule, Trainer, seed_everything
|
9 |
+
DATASET = '/content/drive/MyDrive/sunandini/pascal/PASCAL_VOC'
|
10 |
+
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
11 |
+
seed_everything() # If you want deterministic behavior
|
12 |
+
NUM_WORKERS = os.cpu_count()-1
|
13 |
+
BATCH_SIZE = 32
|
14 |
+
IMAGE_SIZE = 416
|
15 |
+
NUM_CLASSES = 20
|
16 |
+
LEARNING_RATE = 1e-5
|
17 |
+
WEIGHT_DECAY = 1e-4
|
18 |
+
NUM_EPOCHS = 40
|
19 |
+
CONF_THRESHOLD = 0.05
|
20 |
+
MAP_IOU_THRESH = 0.5
|
21 |
+
NMS_IOU_THRESH = 0.45
|
22 |
+
S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]
|
23 |
+
PIN_MEMORY = True
|
24 |
+
LOAD_MODEL = False
|
25 |
+
SAVE_MODEL = True
|
26 |
+
CHECKPOINT_FILE = "checkpoint.pth.tar"
|
27 |
+
IMG_DIR = DATASET + "/images/"
|
28 |
+
LABEL_DIR = DATASET + "/labels/"
|
29 |
+
|
30 |
+
ANCHORS = [
|
31 |
+
[(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],
|
32 |
+
[(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],
|
33 |
+
[(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
|
34 |
+
] # Note these have been rescaled to be between [0, 1]
|
35 |
+
|
36 |
+
means = [0.485, 0.456, 0.406]
|
37 |
+
|
38 |
+
scale = 1.1
|
39 |
+
train_transforms = A.Compose(
|
40 |
+
[
|
41 |
+
A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)),
|
42 |
+
A.PadIfNeeded(
|
43 |
+
min_height=int(IMAGE_SIZE * scale),
|
44 |
+
min_width=int(IMAGE_SIZE * scale),
|
45 |
+
border_mode=cv2.BORDER_CONSTANT,
|
46 |
+
),
|
47 |
+
A.Rotate(limit = 10, interpolation=1, border_mode=4),
|
48 |
+
A.RandomCrop(width=IMAGE_SIZE, height=IMAGE_SIZE),
|
49 |
+
A.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.6, p=0.4),
|
50 |
+
A.OneOf(
|
51 |
+
[
|
52 |
+
A.ShiftScaleRotate(
|
53 |
+
rotate_limit=20, p=0.5, border_mode=cv2.BORDER_CONSTANT
|
54 |
+
),
|
55 |
+
# A.Affine(shear=15, p=0.5, mode="constant"),
|
56 |
+
],
|
57 |
+
p=1.0,
|
58 |
+
),
|
59 |
+
A.HorizontalFlip(p=0.5),
|
60 |
+
A.Blur(p=0.1),
|
61 |
+
A.CLAHE(p=0.1),
|
62 |
+
A.Posterize(p=0.1),
|
63 |
+
A.ToGray(p=0.1),
|
64 |
+
A.ChannelShuffle(p=0.05),
|
65 |
+
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
|
66 |
+
ToTensorV2(),
|
67 |
+
],
|
68 |
+
bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[],),
|
69 |
+
)
|
70 |
+
test_transforms = A.Compose(
|
71 |
+
[
|
72 |
+
A.LongestMaxSize(max_size=IMAGE_SIZE),
|
73 |
+
A.PadIfNeeded(
|
74 |
+
min_height=IMAGE_SIZE, min_width=IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
|
75 |
+
),
|
76 |
+
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
|
77 |
+
ToTensorV2(),
|
78 |
+
],
|
79 |
+
bbox_params=A.BboxParams(format="yolo", min_visibility=0.4, label_fields=[]),
|
80 |
+
)
|
81 |
+
|
82 |
+
PASCAL_CLASSES = [
|
83 |
+
"aeroplane",
|
84 |
+
"bicycle",
|
85 |
+
"bird",
|
86 |
+
"boat",
|
87 |
+
"bottle",
|
88 |
+
"bus",
|
89 |
+
"car",
|
90 |
+
"cat",
|
91 |
+
"chair",
|
92 |
+
"cow",
|
93 |
+
"diningtable",
|
94 |
+
"dog",
|
95 |
+
"horse",
|
96 |
+
"motorbike",
|
97 |
+
"person",
|
98 |
+
"pottedplant",
|
99 |
+
"sheep",
|
100 |
+
"sofa",
|
101 |
+
"train",
|
102 |
+
"tvmonitor"
|
103 |
+
]
|
utils.py
ADDED
@@ -0,0 +1,584 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import config
|
2 |
+
import matplotlib.pyplot as plt
|
3 |
+
import matplotlib.patches as patches
|
4 |
+
import numpy as np
|
5 |
+
import os
|
6 |
+
import random
|
7 |
+
import torch
|
8 |
+
|
9 |
+
from collections import Counter
|
10 |
+
from torch.utils.data import DataLoader
|
11 |
+
from tqdm import tqdm
|
12 |
+
|
13 |
+
|
14 |
+
def iou_width_height(boxes1, boxes2):
|
15 |
+
"""
|
16 |
+
Parameters:
|
17 |
+
boxes1 (tensor): width and height of the first bounding boxes
|
18 |
+
boxes2 (tensor): width and height of the second bounding boxes
|
19 |
+
Returns:
|
20 |
+
tensor: Intersection over union of the corresponding boxes
|
21 |
+
"""
|
22 |
+
intersection = torch.min(boxes1[..., 0], boxes2[..., 0]) * torch.min(
|
23 |
+
boxes1[..., 1], boxes2[..., 1]
|
24 |
+
)
|
25 |
+
union = (
|
26 |
+
boxes1[..., 0] * boxes1[..., 1] + boxes2[..., 0] * boxes2[..., 1] - intersection
|
27 |
+
)
|
28 |
+
return intersection / union
|
29 |
+
|
30 |
+
|
31 |
+
def intersection_over_union(boxes_preds, boxes_labels, box_format="midpoint"):
|
32 |
+
"""
|
33 |
+
Video explanation of this function:
|
34 |
+
https://youtu.be/XXYG5ZWtjj0
|
35 |
+
|
36 |
+
This function calculates intersection over union (iou) given pred boxes
|
37 |
+
and target boxes.
|
38 |
+
|
39 |
+
Parameters:
|
40 |
+
boxes_preds (tensor): Predictions of Bounding Boxes (BATCH_SIZE, 4)
|
41 |
+
boxes_labels (tensor): Correct labels of Bounding Boxes (BATCH_SIZE, 4)
|
42 |
+
box_format (str): midpoint/corners, if boxes (x,y,w,h) or (x1,y1,x2,y2)
|
43 |
+
|
44 |
+
Returns:
|
45 |
+
tensor: Intersection over union for all examples
|
46 |
+
"""
|
47 |
+
|
48 |
+
if box_format == "midpoint":
|
49 |
+
box1_x1 = boxes_preds[..., 0:1] - boxes_preds[..., 2:3] / 2
|
50 |
+
box1_y1 = boxes_preds[..., 1:2] - boxes_preds[..., 3:4] / 2
|
51 |
+
box1_x2 = boxes_preds[..., 0:1] + boxes_preds[..., 2:3] / 2
|
52 |
+
box1_y2 = boxes_preds[..., 1:2] + boxes_preds[..., 3:4] / 2
|
53 |
+
box2_x1 = boxes_labels[..., 0:1] - boxes_labels[..., 2:3] / 2
|
54 |
+
box2_y1 = boxes_labels[..., 1:2] - boxes_labels[..., 3:4] / 2
|
55 |
+
box2_x2 = boxes_labels[..., 0:1] + boxes_labels[..., 2:3] / 2
|
56 |
+
box2_y2 = boxes_labels[..., 1:2] + boxes_labels[..., 3:4] / 2
|
57 |
+
|
58 |
+
if box_format == "corners":
|
59 |
+
box1_x1 = boxes_preds[..., 0:1]
|
60 |
+
box1_y1 = boxes_preds[..., 1:2]
|
61 |
+
box1_x2 = boxes_preds[..., 2:3]
|
62 |
+
box1_y2 = boxes_preds[..., 3:4]
|
63 |
+
box2_x1 = boxes_labels[..., 0:1]
|
64 |
+
box2_y1 = boxes_labels[..., 1:2]
|
65 |
+
box2_x2 = boxes_labels[..., 2:3]
|
66 |
+
box2_y2 = boxes_labels[..., 3:4]
|
67 |
+
|
68 |
+
x1 = torch.max(box1_x1, box2_x1)
|
69 |
+
y1 = torch.max(box1_y1, box2_y1)
|
70 |
+
x2 = torch.min(box1_x2, box2_x2)
|
71 |
+
y2 = torch.min(box1_y2, box2_y2)
|
72 |
+
|
73 |
+
intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)
|
74 |
+
box1_area = abs((box1_x2 - box1_x1) * (box1_y2 - box1_y1))
|
75 |
+
box2_area = abs((box2_x2 - box2_x1) * (box2_y2 - box2_y1))
|
76 |
+
|
77 |
+
return intersection / (box1_area + box2_area - intersection + 1e-6)
|
78 |
+
|
79 |
+
|
80 |
+
def non_max_suppression(bboxes, iou_threshold, threshold, box_format="corners"):
|
81 |
+
"""
|
82 |
+
Video explanation of this function:
|
83 |
+
https://youtu.be/YDkjWEN8jNA
|
84 |
+
|
85 |
+
Does Non Max Suppression given bboxes
|
86 |
+
|
87 |
+
Parameters:
|
88 |
+
bboxes (list): list of lists containing all bboxes with each bboxes
|
89 |
+
specified as [class_pred, prob_score, x1, y1, x2, y2]
|
90 |
+
iou_threshold (float): threshold where predicted bboxes is correct
|
91 |
+
threshold (float): threshold to remove predicted bboxes (independent of IoU)
|
92 |
+
box_format (str): "midpoint" or "corners" used to specify bboxes
|
93 |
+
|
94 |
+
Returns:
|
95 |
+
list: bboxes after performing NMS given a specific IoU threshold
|
96 |
+
"""
|
97 |
+
|
98 |
+
assert type(bboxes) == list
|
99 |
+
|
100 |
+
bboxes = [box for box in bboxes if box[1] > threshold]
|
101 |
+
bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
|
102 |
+
bboxes_after_nms = []
|
103 |
+
|
104 |
+
while bboxes:
|
105 |
+
chosen_box = bboxes.pop(0)
|
106 |
+
|
107 |
+
bboxes = [
|
108 |
+
box
|
109 |
+
for box in bboxes
|
110 |
+
if box[0] != chosen_box[0]
|
111 |
+
or intersection_over_union(
|
112 |
+
torch.tensor(chosen_box[2:]),
|
113 |
+
torch.tensor(box[2:]),
|
114 |
+
box_format=box_format,
|
115 |
+
)
|
116 |
+
< iou_threshold
|
117 |
+
]
|
118 |
+
|
119 |
+
bboxes_after_nms.append(chosen_box)
|
120 |
+
|
121 |
+
return bboxes_after_nms
|
122 |
+
|
123 |
+
|
124 |
+
def mean_average_precision(
|
125 |
+
pred_boxes, true_boxes, iou_threshold=0.5, box_format="midpoint", num_classes=20
|
126 |
+
):
|
127 |
+
"""
|
128 |
+
Video explanation of this function:
|
129 |
+
https://youtu.be/FppOzcDvaDI
|
130 |
+
|
131 |
+
This function calculates mean average precision (mAP)
|
132 |
+
|
133 |
+
Parameters:
|
134 |
+
pred_boxes (list): list of lists containing all bboxes with each bboxes
|
135 |
+
specified as [train_idx, class_prediction, prob_score, x1, y1, x2, y2]
|
136 |
+
true_boxes (list): Similar as pred_boxes except all the correct ones
|
137 |
+
iou_threshold (float): threshold where predicted bboxes is correct
|
138 |
+
box_format (str): "midpoint" or "corners" used to specify bboxes
|
139 |
+
num_classes (int): number of classes
|
140 |
+
|
141 |
+
Returns:
|
142 |
+
float: mAP value across all classes given a specific IoU threshold
|
143 |
+
"""
|
144 |
+
|
145 |
+
# list storing all AP for respective classes
|
146 |
+
average_precisions = []
|
147 |
+
|
148 |
+
# used for numerical stability later on
|
149 |
+
epsilon = 1e-6
|
150 |
+
|
151 |
+
for c in range(num_classes):
|
152 |
+
detections = []
|
153 |
+
ground_truths = []
|
154 |
+
|
155 |
+
# Go through all predictions and targets,
|
156 |
+
# and only add the ones that belong to the
|
157 |
+
# current class c
|
158 |
+
for detection in pred_boxes:
|
159 |
+
if detection[1] == c:
|
160 |
+
detections.append(detection)
|
161 |
+
|
162 |
+
for true_box in true_boxes:
|
163 |
+
if true_box[1] == c:
|
164 |
+
ground_truths.append(true_box)
|
165 |
+
|
166 |
+
# find the amount of bboxes for each training example
|
167 |
+
# Counter here finds how many ground truth bboxes we get
|
168 |
+
# for each training example, so let's say img 0 has 3,
|
169 |
+
# img 1 has 5 then we will obtain a dictionary with:
|
170 |
+
# amount_bboxes = {0:3, 1:5}
|
171 |
+
amount_bboxes = Counter([gt[0] for gt in ground_truths])
|
172 |
+
|
173 |
+
# We then go through each key, val in this dictionary
|
174 |
+
# and convert to the following (w.r.t same example):
|
175 |
+
# ammount_bboxes = {0:torch.tensor[0,0,0], 1:torch.tensor[0,0,0,0,0]}
|
176 |
+
for key, val in amount_bboxes.items():
|
177 |
+
amount_bboxes[key] = torch.zeros(val)
|
178 |
+
|
179 |
+
# sort by box probabilities which is index 2
|
180 |
+
detections.sort(key=lambda x: x[2], reverse=True)
|
181 |
+
TP = torch.zeros((len(detections)))
|
182 |
+
FP = torch.zeros((len(detections)))
|
183 |
+
total_true_bboxes = len(ground_truths)
|
184 |
+
|
185 |
+
# If none exists for this class then we can safely skip
|
186 |
+
if total_true_bboxes == 0:
|
187 |
+
continue
|
188 |
+
|
189 |
+
for detection_idx, detection in enumerate(detections):
|
190 |
+
# Only take out the ground_truths that have the same
|
191 |
+
# training idx as detection
|
192 |
+
ground_truth_img = [
|
193 |
+
bbox for bbox in ground_truths if bbox[0] == detection[0]
|
194 |
+
]
|
195 |
+
|
196 |
+
num_gts = len(ground_truth_img)
|
197 |
+
best_iou = 0
|
198 |
+
|
199 |
+
for idx, gt in enumerate(ground_truth_img):
|
200 |
+
iou = intersection_over_union(
|
201 |
+
torch.tensor(detection[3:]),
|
202 |
+
torch.tensor(gt[3:]),
|
203 |
+
box_format=box_format,
|
204 |
+
)
|
205 |
+
|
206 |
+
if iou > best_iou:
|
207 |
+
best_iou = iou
|
208 |
+
best_gt_idx = idx
|
209 |
+
|
210 |
+
if best_iou > iou_threshold:
|
211 |
+
# only detect ground truth detection once
|
212 |
+
if amount_bboxes[detection[0]][best_gt_idx] == 0:
|
213 |
+
# true positive and add this bounding box to seen
|
214 |
+
TP[detection_idx] = 1
|
215 |
+
amount_bboxes[detection[0]][best_gt_idx] = 1
|
216 |
+
else:
|
217 |
+
FP[detection_idx] = 1
|
218 |
+
|
219 |
+
# if IOU is lower then the detection is a false positive
|
220 |
+
else:
|
221 |
+
FP[detection_idx] = 1
|
222 |
+
|
223 |
+
TP_cumsum = torch.cumsum(TP, dim=0)
|
224 |
+
FP_cumsum = torch.cumsum(FP, dim=0)
|
225 |
+
recalls = TP_cumsum / (total_true_bboxes + epsilon)
|
226 |
+
precisions = TP_cumsum / (TP_cumsum + FP_cumsum + epsilon)
|
227 |
+
precisions = torch.cat((torch.tensor([1]), precisions))
|
228 |
+
recalls = torch.cat((torch.tensor([0]), recalls))
|
229 |
+
# torch.trapz for numerical integration
|
230 |
+
average_precisions.append(torch.trapz(precisions, recalls))
|
231 |
+
|
232 |
+
return sum(average_precisions) / len(average_precisions)
|
233 |
+
|
234 |
+
|
235 |
+
def plot_image(image, boxes):
|
236 |
+
"""Plots predicted bounding boxes on the image"""
|
237 |
+
cmap = plt.get_cmap("tab20b")
|
238 |
+
class_labels = config.COCO_LABELS if config.DATASET=='COCO' else config.PASCAL_CLASSES
|
239 |
+
colors = [cmap(i) for i in np.linspace(0, 1, len(class_labels))]
|
240 |
+
im = np.array(image)
|
241 |
+
height, width, _ = im.shape
|
242 |
+
|
243 |
+
# Create figure and axes
|
244 |
+
fig, ax = plt.subplots(1)
|
245 |
+
# Display the image
|
246 |
+
ax.imshow(im)
|
247 |
+
|
248 |
+
# box[0] is x midpoint, box[2] is width
|
249 |
+
# box[1] is y midpoint, box[3] is height
|
250 |
+
|
251 |
+
# Create a Rectangle patch
|
252 |
+
for box in boxes:
|
253 |
+
assert len(box) == 6, "box should contain class pred, confidence, x, y, width, height"
|
254 |
+
class_pred = box[0]
|
255 |
+
box = box[2:]
|
256 |
+
upper_left_x = box[0] - box[2] / 2
|
257 |
+
upper_left_y = box[1] - box[3] / 2
|
258 |
+
rect = patches.Rectangle(
|
259 |
+
(upper_left_x * width, upper_left_y * height),
|
260 |
+
box[2] * width,
|
261 |
+
box[3] * height,
|
262 |
+
linewidth=2,
|
263 |
+
edgecolor=colors[int(class_pred)],
|
264 |
+
facecolor="none",
|
265 |
+
)
|
266 |
+
# Add the patch to the Axes
|
267 |
+
ax.add_patch(rect)
|
268 |
+
plt.text(
|
269 |
+
upper_left_x * width,
|
270 |
+
upper_left_y * height,
|
271 |
+
s=class_labels[int(class_pred)],
|
272 |
+
color="white",
|
273 |
+
verticalalignment="top",
|
274 |
+
bbox={"color": colors[int(class_pred)], "pad": 0},
|
275 |
+
)
|
276 |
+
|
277 |
+
plt.show()
|
278 |
+
|
279 |
+
|
280 |
+
def get_evaluation_bboxes(
|
281 |
+
loader,
|
282 |
+
model,
|
283 |
+
iou_threshold,
|
284 |
+
anchors,
|
285 |
+
threshold,
|
286 |
+
box_format="midpoint",
|
287 |
+
device="cuda",
|
288 |
+
):
|
289 |
+
# make sure model is in eval before get bboxes
|
290 |
+
model.eval()
|
291 |
+
train_idx = 0
|
292 |
+
all_pred_boxes = []
|
293 |
+
all_true_boxes = []
|
294 |
+
for batch_idx, (x, labels) in enumerate(loader):
|
295 |
+
x = x.to(device)
|
296 |
+
|
297 |
+
with torch.no_grad():
|
298 |
+
predictions = model(x)
|
299 |
+
|
300 |
+
batch_size = x.shape[0]
|
301 |
+
bboxes = [[] for _ in range(batch_size)]
|
302 |
+
for i in range(3):
|
303 |
+
S = predictions[i].shape[2]
|
304 |
+
anchor = torch.tensor([*anchors[i]]).to(device) * S
|
305 |
+
boxes_scale_i = cells_to_bboxes(
|
306 |
+
predictions[i], anchor, S=S, is_preds=True
|
307 |
+
)
|
308 |
+
for idx, (box) in enumerate(boxes_scale_i):
|
309 |
+
bboxes[idx] += box
|
310 |
+
|
311 |
+
# we just want one bbox for each label, not one for each scale
|
312 |
+
true_bboxes = cells_to_bboxes(
|
313 |
+
labels[2], anchor, S=S, is_preds=False
|
314 |
+
)
|
315 |
+
|
316 |
+
for idx in range(batch_size):
|
317 |
+
nms_boxes = non_max_suppression(
|
318 |
+
bboxes[idx],
|
319 |
+
iou_threshold=iou_threshold,
|
320 |
+
threshold=threshold,
|
321 |
+
box_format=box_format,
|
322 |
+
)
|
323 |
+
|
324 |
+
for nms_box in nms_boxes:
|
325 |
+
all_pred_boxes.append([train_idx] + nms_box)
|
326 |
+
|
327 |
+
for box in true_bboxes[idx]:
|
328 |
+
if box[1] > threshold:
|
329 |
+
all_true_boxes.append([train_idx] + box)
|
330 |
+
|
331 |
+
train_idx += 1
|
332 |
+
|
333 |
+
model.train()
|
334 |
+
return all_pred_boxes, all_true_boxes
|
335 |
+
|
336 |
+
|
337 |
+
def cells_to_bboxes(predictions, anchors, S, is_preds=True):
|
338 |
+
"""
|
339 |
+
Scales the predictions coming from the model to
|
340 |
+
be relative to the entire image such that they for example later
|
341 |
+
can be plotted or.
|
342 |
+
INPUT:
|
343 |
+
predictions: tensor of size (N, 3, S, S, num_classes+5)
|
344 |
+
anchors: the anchors used for the predictions
|
345 |
+
S: the number of cells the image is divided in on the width (and height)
|
346 |
+
is_preds: whether the input is predictions or the true bounding boxes
|
347 |
+
OUTPUT:
|
348 |
+
converted_bboxes: the converted boxes of sizes (N, num_anchors, S, S, 1+5) with class index,
|
349 |
+
object score, bounding box coordinates
|
350 |
+
"""
|
351 |
+
BATCH_SIZE = predictions.shape[0]
|
352 |
+
num_anchors = len(anchors)
|
353 |
+
box_predictions = predictions[..., 1:5]
|
354 |
+
if is_preds:
|
355 |
+
anchors = anchors.reshape(1, len(anchors), 1, 1, 2)
|
356 |
+
box_predictions[..., 0:2] = torch.sigmoid(box_predictions[..., 0:2])
|
357 |
+
box_predictions[..., 2:] = torch.exp(box_predictions[..., 2:]) * anchors
|
358 |
+
scores = torch.sigmoid(predictions[..., 0:1])
|
359 |
+
best_class = torch.argmax(predictions[..., 5:], dim=-1).unsqueeze(-1)
|
360 |
+
else:
|
361 |
+
scores = predictions[..., 0:1]
|
362 |
+
best_class = predictions[..., 5:6]
|
363 |
+
|
364 |
+
cell_indices = (
|
365 |
+
torch.arange(S)
|
366 |
+
.repeat(predictions.shape[0], 3, S, 1)
|
367 |
+
.unsqueeze(-1)
|
368 |
+
.to(predictions.device)
|
369 |
+
)
|
370 |
+
x = 1 / S * (box_predictions[..., 0:1] + cell_indices)
|
371 |
+
y = 1 / S * (box_predictions[..., 1:2] + cell_indices.permute(0, 1, 3, 2, 4))
|
372 |
+
w_h = 1 / S * box_predictions[..., 2:4]
|
373 |
+
converted_bboxes = torch.cat((best_class, scores, x, y, w_h), dim=-1).reshape(BATCH_SIZE, num_anchors * S * S, 6)
|
374 |
+
return converted_bboxes.tolist()
|
375 |
+
|
376 |
+
def check_class_accuracy(model, loader, threshold):
|
377 |
+
model.eval()
|
378 |
+
tot_class_preds, correct_class = 0, 0
|
379 |
+
tot_noobj, correct_noobj = 0, 0
|
380 |
+
tot_obj, correct_obj = 0, 0
|
381 |
+
|
382 |
+
for idx, (x, y) in enumerate(loader):
|
383 |
+
x = x.to(config.DEVICE)
|
384 |
+
with torch.no_grad():
|
385 |
+
out = model(x)
|
386 |
+
|
387 |
+
for i in range(3):
|
388 |
+
y[i] = y[i].to(config.DEVICE)
|
389 |
+
obj = y[i][..., 0] == 1 # in paper this is Iobj_i
|
390 |
+
noobj = y[i][..., 0] == 0 # in paper this is Iobj_i
|
391 |
+
|
392 |
+
correct_class += torch.sum(
|
393 |
+
torch.argmax(out[i][..., 5:][obj], dim=-1) == y[i][..., 5][obj]
|
394 |
+
)
|
395 |
+
tot_class_preds += torch.sum(obj)
|
396 |
+
|
397 |
+
obj_preds = torch.sigmoid(out[i][..., 0]) > threshold
|
398 |
+
correct_obj += torch.sum(obj_preds[obj] == y[i][..., 0][obj])
|
399 |
+
tot_obj += torch.sum(obj)
|
400 |
+
correct_noobj += torch.sum(obj_preds[noobj] == y[i][..., 0][noobj])
|
401 |
+
tot_noobj += torch.sum(noobj)
|
402 |
+
|
403 |
+
print(f"Class accuracy is: {(correct_class/(tot_class_preds+1e-16))*100:2f}%")
|
404 |
+
print(f"No obj accuracy is: {(correct_noobj/(tot_noobj+1e-16))*100:2f}%")
|
405 |
+
print(f"Obj accuracy is: {(correct_obj/(tot_obj+1e-16))*100:2f}%")
|
406 |
+
model.train()
|
407 |
+
|
408 |
+
return (correct_class/(tot_class_preds+1e-16))*100, (correct_noobj/(tot_noobj+1e-16))*100, (correct_obj/(tot_obj+1e-16))*100
|
409 |
+
|
410 |
+
|
411 |
+
def get_mean_std(loader):
|
412 |
+
# var[X] = E[X**2] - E[X]**2
|
413 |
+
channels_sum, channels_sqrd_sum, num_batches = 0, 0, 0
|
414 |
+
|
415 |
+
for data, _ in loader:
|
416 |
+
channels_sum += torch.mean(data, dim=[0, 2, 3])
|
417 |
+
channels_sqrd_sum += torch.mean(data ** 2, dim=[0, 2, 3])
|
418 |
+
num_batches += 1
|
419 |
+
|
420 |
+
mean = channels_sum / num_batches
|
421 |
+
std = (channels_sqrd_sum / num_batches - mean ** 2) ** 0.5
|
422 |
+
|
423 |
+
return mean, std
|
424 |
+
|
425 |
+
|
426 |
+
def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
|
427 |
+
print("=> Saving checkpoint")
|
428 |
+
checkpoint = {
|
429 |
+
"state_dict": model.state_dict(),
|
430 |
+
"optimizer": optimizer.state_dict(),
|
431 |
+
}
|
432 |
+
torch.save(checkpoint, filename)
|
433 |
+
|
434 |
+
|
435 |
+
def load_checkpoint(checkpoint_file, model, optimizer, lr):
|
436 |
+
print("=> Loading checkpoint")
|
437 |
+
checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)
|
438 |
+
model.load_state_dict(checkpoint["state_dict"])
|
439 |
+
optimizer.load_state_dict(checkpoint["optimizer"])
|
440 |
+
|
441 |
+
# If we don't do this then it will just have learning rate of old checkpoint
|
442 |
+
# and it will lead to many hours of debugging \:
|
443 |
+
for param_group in optimizer.param_groups:
|
444 |
+
param_group["lr"] = lr
|
445 |
+
|
446 |
+
|
447 |
+
def get_loaders(train_csv_path, test_csv_path):
|
448 |
+
from dataset import YOLODataset
|
449 |
+
|
450 |
+
IMAGE_SIZE = config.IMAGE_SIZE
|
451 |
+
train_dataset = YOLODataset(
|
452 |
+
train_csv_path,
|
453 |
+
transform=config.train_transforms,
|
454 |
+
S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
|
455 |
+
img_dir=config.IMG_DIR,
|
456 |
+
label_dir=config.LABEL_DIR,
|
457 |
+
anchors=config.ANCHORS,
|
458 |
+
)
|
459 |
+
test_dataset = YOLODataset(
|
460 |
+
test_csv_path,
|
461 |
+
transform=config.test_transforms,
|
462 |
+
S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
|
463 |
+
img_dir=config.IMG_DIR,
|
464 |
+
label_dir=config.LABEL_DIR,
|
465 |
+
anchors=config.ANCHORS,
|
466 |
+
)
|
467 |
+
train_loader = DataLoader(
|
468 |
+
dataset=train_dataset,
|
469 |
+
batch_size=config.BATCH_SIZE,
|
470 |
+
num_workers=config.NUM_WORKERS,
|
471 |
+
pin_memory=config.PIN_MEMORY,
|
472 |
+
shuffle=True,
|
473 |
+
drop_last=False,
|
474 |
+
)
|
475 |
+
test_loader = DataLoader(
|
476 |
+
dataset=test_dataset,
|
477 |
+
batch_size=config.BATCH_SIZE,
|
478 |
+
num_workers=config.NUM_WORKERS,
|
479 |
+
pin_memory=config.PIN_MEMORY,
|
480 |
+
shuffle=False,
|
481 |
+
drop_last=False,
|
482 |
+
)
|
483 |
+
|
484 |
+
train_eval_dataset = YOLODataset(
|
485 |
+
train_csv_path,
|
486 |
+
transform=config.test_transforms,
|
487 |
+
S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
|
488 |
+
img_dir=config.IMG_DIR,
|
489 |
+
label_dir=config.LABEL_DIR,
|
490 |
+
anchors=config.ANCHORS,
|
491 |
+
)
|
492 |
+
train_eval_loader = DataLoader(
|
493 |
+
dataset=train_eval_dataset,
|
494 |
+
batch_size=config.BATCH_SIZE,
|
495 |
+
num_workers=config.NUM_WORKERS,
|
496 |
+
pin_memory=config.PIN_MEMORY,
|
497 |
+
shuffle=False,
|
498 |
+
drop_last=False,
|
499 |
+
)
|
500 |
+
|
501 |
+
return train_loader, test_loader, train_eval_loader
|
502 |
+
|
503 |
+
def plot_couple_examples(model, loader, thresh, iou_thresh, anchors):
|
504 |
+
model.eval()
|
505 |
+
x, y = next(iter(loader))
|
506 |
+
x = x.to("cuda")
|
507 |
+
with torch.no_grad():
|
508 |
+
out = model(x)
|
509 |
+
bboxes = [[] for _ in range(x.shape[0])]
|
510 |
+
for i in range(3):
|
511 |
+
batch_size, A, S, _, _ = out[i].shape
|
512 |
+
anchor = anchors[i]
|
513 |
+
boxes_scale_i = cells_to_bboxes(
|
514 |
+
out[i], anchor, S=S, is_preds=True
|
515 |
+
)
|
516 |
+
for idx, (box) in enumerate(boxes_scale_i):
|
517 |
+
bboxes[idx] += box
|
518 |
+
|
519 |
+
model.train()
|
520 |
+
|
521 |
+
for i in range(batch_size//4):
|
522 |
+
nms_boxes = non_max_suppression(
|
523 |
+
bboxes[i], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
|
524 |
+
)
|
525 |
+
plot_image(x[i].permute(1,2,0).detach().cpu(), nms_boxes)
|
526 |
+
|
527 |
+
|
528 |
+
|
529 |
+
def seed_everything(seed=42):
|
530 |
+
os.environ['PYTHONHASHSEED'] = str(seed)
|
531 |
+
random.seed(seed)
|
532 |
+
np.random.seed(seed)
|
533 |
+
torch.manual_seed(seed)
|
534 |
+
torch.cuda.manual_seed(seed)
|
535 |
+
torch.cuda.manual_seed_all(seed)
|
536 |
+
torch.backends.cudnn.deterministic = True
|
537 |
+
torch.backends.cudnn.benchmark = False
|
538 |
+
|
539 |
+
|
540 |
+
def clip_coords(boxes, img_shape):
|
541 |
+
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
542 |
+
boxes[:, 0].clamp_(0, img_shape[1]) # x1
|
543 |
+
boxes[:, 1].clamp_(0, img_shape[0]) # y1
|
544 |
+
boxes[:, 2].clamp_(0, img_shape[1]) # x2
|
545 |
+
boxes[:, 3].clamp_(0, img_shape[0]) # y2
|
546 |
+
|
547 |
+
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
548 |
+
# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
549 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
550 |
+
y[..., 0] = w * (x[..., 0] - x[..., 2] / 2) + padw # top left x
|
551 |
+
y[..., 1] = h * (x[..., 1] - x[..., 3] / 2) + padh # top left y
|
552 |
+
y[..., 2] = w * (x[..., 0] + x[..., 2] / 2) + padw # bottom right x
|
553 |
+
y[..., 3] = h * (x[..., 1] + x[..., 3] / 2) + padh # bottom right y
|
554 |
+
return y
|
555 |
+
|
556 |
+
|
557 |
+
def xyn2xy(x, w=640, h=640, padw=0, padh=0):
|
558 |
+
# Convert normalized segments into pixel segments, shape (n,2)
|
559 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
560 |
+
y[..., 0] = w * x[..., 0] + padw # top left x
|
561 |
+
y[..., 1] = h * x[..., 1] + padh # top left y
|
562 |
+
return y
|
563 |
+
|
564 |
+
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
|
565 |
+
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right
|
566 |
+
if clip:
|
567 |
+
clip_boxes(x, (h - eps, w - eps)) # warning: inplace clip
|
568 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
569 |
+
y[..., 0] = ((x[..., 0] + x[..., 2]) / 2) / w # x center
|
570 |
+
y[..., 1] = ((x[..., 1] + x[..., 3]) / 2) / h # y center
|
571 |
+
y[..., 2] = (x[..., 2] - x[..., 0]) / w # width
|
572 |
+
y[..., 3] = (x[..., 3] - x[..., 1]) / h # height
|
573 |
+
return y
|
574 |
+
|
575 |
+
def clip_boxes(boxes, shape):
|
576 |
+
# Clip boxes (xyxy) to image shape (height, width)
|
577 |
+
if isinstance(boxes, torch.Tensor): # faster individually
|
578 |
+
boxes[..., 0].clamp_(0, shape[1]) # x1
|
579 |
+
boxes[..., 1].clamp_(0, shape[0]) # y1
|
580 |
+
boxes[..., 2].clamp_(0, shape[1]) # x2
|
581 |
+
boxes[..., 3].clamp_(0, shape[0]) # y2
|
582 |
+
else: # np.array (faster grouped)
|
583 |
+
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2
|
584 |
+
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2
|