Gosula commited on
Commit
eae1d6d
1 Parent(s): acb86af

Upload 2 files

Browse files
Files changed (2) hide show
  1. config.py +103 -0
  2. utils.py +584 -0
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