Update files

data/coco_dataset.py

@@ -7,9 +7,10 @@
 import os
 import cv2
 import numpy as np
-from data.augmentations import Augmentations, CopyPaste, CutOut
-from data.process_box import x1y1_to_xcyc, x1y1wh_to_xyxy, xyxy_to_x1y1wh, normalize_bboxes, resize_bboxes, adjust_bboxes
-
+from augmentations import Augmentations, CopyPaste, CutOut
+from process_box import x1y1_to_xcyc, x1y1wh_to_xyxy, xyxy_to_x1y1wh, normalize_bboxes, resize_bboxes, adjust_bboxes
+import imgaug.augmenters as iaa
+from imgaug.augmentables.bbs import BoundingBoxesOnImage
 
 
 class CustomDataset(Dataset):
@@ -20,7 +21,9 @@ def __init__(self, image_path, annotation_path, image_size=640, normalize=True,
         self.image_root = os.path.join('/', *image_path.split('/')[:-1])
         self.image_path = image_path
         self.annotation_path = annotation_path
-        self.normalize = normalize
+        self.normalize = False
+        import json
+        annot = json.load(open(annotation_path))
         self.coco = COCO(self.annotation_path)
         self.image_paths = sorted(os.listdir(self.image_path))
         self.ids = sorted(list(self.coco.imgs.keys()))
@@ -57,7 +60,10 @@ def __getitem__(self, index):
 
         img, ratio = self.load_image(image_id)
         labels = self.load_labels(image_id=image_id, ratio=ratio)
-
+        bboxes = np.array(img_data['labels'][:,:4])
+        bboxes_iaa = BoundingBoxesOnImage([], img.shape).from_xyxy_array(bboxes, img.shape)
+        image_after = bboxes_iaa.draw_on_image(img, size=2)
+        cv2.imwrite('./labelled/image_' + str(index) + '.jpg', image_after)
         if self.augmentations:
 
             img_data = {'img': img, 'labels':labels}

models/utils.py

@@ -3,6 +3,8 @@
 import torch
 import math
 import numpy as np
+import cv2
+from typing import Union
 
 
 def init_weights(model):
@@ -31,4 +33,174 @@ def variance_scaling_(tensor, gain=1.):
     fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
     std = math.sqrt(gain / float(fan_in))
 
-    return _no_grad_normal_(tensor, 0., std)
+    return _no_grad_normal_(tensor, 0., std)
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+import torch.optim as optim
+import math
+
+def get_optimizer(opt, model):
+    if opt['optimizer'] == 'adam':
+        return optim.Adam(model.parameters(), lr=opt['learning_rate'])
+    elif opt['optimizer'] == 'sgd':
+        return optim.SGD(model.parameters(), lr=opt['learning_rate'], momentum=0.9)
+    elif opt['optimizer'] == 'asgd':
+        return optim.ASGD(model.parameters(), lr=opt['learning_rate'])
+
+
+def get_scheduler(opt, optimizer, len_dataset):
+
+    if opt['lr_scheduler'] == 'multistep_lr':
+        return optim.lr_scheduler.MultiStepLR(optimizer, milestones=[int(opt['epochs']*0.7), int(opt['epochs']*0.9)], gamma=0.1)
+    elif opt['lr_scheduler'] == 'cosine':
+        lf = lambda x: (((1 + math.cos(x * math.pi / opt['epochs'])) / 2) ** 1.0) * 0.95 + 0.05
+        return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
+    elif opt['lr_scheduler'] == 'cosine_annealing':
+        return optim.lr_scheduler.CosineAnnealingLR(optimizer, len_dataset, eta_min=0)
+
+class BBoxTransform(nn.Module):
+    def forward(self, anchors, regression):
+        """
+        decode_box_outputs adapted from https://github.com/google/automl/blob/master/efficientdet/anchors.py
+
+        Args:
+            anchors: [batchsize, boxes, (y1, x1, y2, x2)]
+            regression: [batchsize, boxes, (dy, dx, dh, dw)]
+
+        Returns:
+
+        """
+        y_centers_a = (anchors[..., 0] + anchors[..., 2]) / 2
+        x_centers_a = (anchors[..., 1] + anchors[..., 3]) / 2
+        ha = anchors[..., 2] - anchors[..., 0]
+        wa = anchors[..., 3] - anchors[..., 1]
+
+        w = regression[..., 3].exp() * wa
+        h = regression[..., 2].exp() * ha
+
+        y_centers = regression[..., 0] * ha + y_centers_a
+        x_centers = regression[..., 1] * wa + x_centers_a
+
+        ymin = y_centers - h / 2.
+        xmin = x_centers - w / 2.
+        ymax = y_centers + h / 2.
+        xmax = x_centers + w / 2.
+
+        return torch.stack([xmin, ymin, xmax, ymax], dim=2)
+
+class ClipBoxes(nn.Module):
+
+    def __init__(self):
+        super(ClipBoxes, self).__init__()
+
+    def forward(self, boxes, img):
+        batch_size, num_channels, height, width = img.shape
+
+        boxes[:, :, 0] = torch.clamp(boxes[:, :, 0], min=0)
+        boxes[:, :, 1] = torch.clamp(boxes[:, :, 1], min=0)
+
+        boxes[:, :, 2] = torch.clamp(boxes[:, :, 2], max=width - 1)
+        boxes[:, :, 3] = torch.clamp(boxes[:, :, 3], max=height - 1)
+
+        return boxes
+
+def invert_affine(metas: Union[float, list, tuple], preds):
+    for i in range(len(preds)):
+        if len(preds[i]['rois']) == 0:
+            continue
+        else:
+            if metas is float:
+                preds[i]['rois'][:, [0, 2]] = preds[i]['rois'][:, [0, 2]] / metas
+                preds[i]['rois'][:, [1, 3]] = preds[i]['rois'][:, [1, 3]] / metas
+            else:
+                new_w, new_h, old_w, old_h, padding_w, padding_h = metas[i]
+                preds[i]['rois'][:, [0, 2]] = preds[i]['rois'][:, [0, 2]] / (new_w / old_w)
+                preds[i]['rois'][:, [1, 3]] = preds[i]['rois'][:, [1, 3]] / (new_h / old_h)
+    return preds
+
+
+def preprocess(*image_path, max_size=512, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+    ori_imgs = [cv2.imread(img_path) for img_path in image_path]
+    normalized_imgs = [(img[..., ::-1] / 255 - mean) / std for img in ori_imgs]
+    imgs_meta = [aspectaware_resize_padding(img, max_size, max_size,
+                                            means=None) for img in normalized_imgs]
+    framed_imgs = [img_meta[0] for img_meta in imgs_meta]
+    framed_metas = [img_meta[1:] for img_meta in imgs_meta]
+
+    return ori_imgs, framed_imgs, framed_metas
+
+def aspectaware_resize_padding(image, width, height, interpolation=None, means=None):
+    old_h, old_w, c = image.shape
+    if old_w > old_h:
+        new_w = width
+        new_h = int(width / old_w * old_h)
+    else:
+        new_w = int(height / old_h * old_w)
+        new_h = height
+
+    canvas = np.zeros((height, height, c), np.float32)
+    if means is not None:
+        canvas[...] = means
+
+    if new_w != old_w or new_h != old_h:
+        if interpolation is None:
+            image = cv2.resize(image, (new_w, new_h))
+        else:
+            image = cv2.resize(image, (new_w, new_h), interpolation=interpolation)
+
+    padding_h = height - new_h
+    padding_w = width - new_w
+
+    if c > 1:
+        canvas[:new_h, :new_w] = image
+    else:
+        if len(image.shape) == 2:
+            canvas[:new_h, :new_w, 0] = image
+        else:
+            canvas[:new_h, :new_w] = image
+
+    return canvas, new_w, new_h, old_w, old_h, padding_w, padding_h,
+
+def postprocess(x, anchors, regression, classification, regressBoxes, clipBoxes, threshold, iou_threshold):
+    transformed_anchors = regressBoxes(anchors, regression)
+    transformed_anchors = clipBoxes(transformed_anchors, x)
+    scores = torch.max(classification, dim=2, keepdim=True)[0]
+    scores_over_thresh = (scores > threshold)[:, :, 0]
+    out = []
+    for i in range(x.shape[0]):
+        if scores_over_thresh[i].sum() == 0:
+            out.append({
+                'rois': np.array(()),
+                'class_ids': np.array(()),
+                'scores': np.array(()),
+            })
+            continue
+
+        classification_per = classification[i, scores_over_thresh[i, :], ...].permute(1, 0)
+        transformed_anchors_per = transformed_anchors[i, scores_over_thresh[i, :], ...]
+        scores_per = scores[i, scores_over_thresh[i, :], ...]
+        scores_, classes_ = classification_per.max(dim=0)
+        anchors_nms_idx = batched_nms(transformed_anchors_per, scores_per[:, 0], classes_, iou_threshold=iou_threshold)
+
+        if anchors_nms_idx.shape[0] != 0:
+            classes_ = classes_[anchors_nms_idx]
+            scores_ = scores_[anchors_nms_idx]
+            boxes_ = transformed_anchors_per[anchors_nms_idx, :]
+
+            out.append({
+                'rois': boxes_.cpu().numpy(),
+                'class_ids': classes_.cpu().numpy(),
+                'scores': scores_.cpu().numpy(),
+            })
+        else:
+            out.append({
+                'rois': np.array(()),
+                'class_ids': np.array(()),
+                'scores': np.array(()),
+            })
+
+    return out