{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Transformers Are Here To Improve Computer Vision Aswell !\n\nThis kernel is forked from this amazing kernel: https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train\n\n\nPlease upvote that kernel as well!\n\nIn this kernel, I use:\n- Detetion Transformer (https://arxiv.org/abs/2005.12872)\n- Albumentation enabled (simple flip for now)\n","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"Images are taken from the original paper. ","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"![Detection Transformer](https://github.com/facebookresearch/detr/blob/master/.github/DETR.png?raw=true)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"![](https://media.geeksforgeeks.org/wp-content/uploads/20200601074340/detr-arch.jpg)","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Loss Function\n\n   In the paper, DETR is trained on the COCO dataset. Suppose the model outputs N box predictions and class predictions for each image. Then, we need to find the right permutation of predictions which gives the least loss. Hungarian Matching algorithm is used to find the correct permutation of predictions, which gives the least loss when calculated wrt to the original class labels and bounding boxes. The total loss after finding the correct permutation is just the addition of usual cross entropy loss for class labels and 'bounding box loss'. Bounding box loss is in turn the weighted sum of l1 loss and generalised IOU loss.    \n   \n In this implementation, I have given equal weightage to each of the losses. You can play with it and try to improve the results.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Comparison between Faster R-CNN & Detection Transformer Architecture\n\nThis image is taken from https://analyticsindiamag.com/object-detection-with-transformers/","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"![](https://lh4.googleusercontent.com/2Fcm3z2XtYSi0gE5HRbEYWr4o5PgVOvQYWYEkm5clwOTVJB9TF3YfLwSLsZ5vgriompdxhDbtMvZ-yacrI8cEO6UW7Qqur1xOU0gFaUZxKWm5WLElH0swaX0b6MK7qYSQJ5_U7uv)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"DEBUG = True","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\n\nfrom PIL import Image\n\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport torch\nimport torchvision\n\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.rpn import AnchorGenerator\n\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.sampler import SequentialSampler\n\nfrom matplotlib import pyplot as plt\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-10:] if DEBUG else image_ids[-665:]\ntrain_ids = image_ids[:10] if DEBUG else image_ids[:-665]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_df.shape, train_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n\n    def __init__(self, dataframe, image_dir, transforms=None):\n        super().__init__()\n\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n\n        image_id = self.image_ids[index]\n        records = self.df[self.df['image_id'] == image_id]\n\n        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n\n        boxes = records[['x', 'y', 'w', 'h']].values\n#         boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n#         boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        \n#         area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = boxes[:,2]*boxes[:,3]\n        area = torch.as_tensor(area, dtype=torch.float32)\n\n        # there is only one class\n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        \n        # suppose all instances are not crowd\n        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)\n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        # target['masks'] = None\n        target['image_id'] = torch.tensor([index])\n        target['area'] = area\n        target['iscrowd'] = iscrowd\n\n        if self.transforms:\n            sample = {\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            }\n            sample = self.transforms(**sample)\n            image = sample['image']\n            \n            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Albumentations\ndef get_train_transform():\n    return A.Compose([\n        A.Flip(0.5),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'coco', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'coco', 'label_fields': ['labels']})\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create the model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# load Detection Transformer model; pre-trained on COCO\nmodel = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_classes = 2  # 1 class (wheat) + background\n\n# # get number of input features for the classifier\n# in_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# # replace the pre-trained head with a new one\n# model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\n\nmodel.class_embed = nn.Linear(256, num_classes, bias = True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Hungarian Matcher","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"from torchvision.ops.boxes import box_area\n\ndef box_iou(boxes1, boxes2):\n    area1 = box_area(boxes1)\n    area2 = box_area(boxes2)\n\n    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]\n    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]\n\n    wh = (rb - lt).clamp(min=0)  # [N,M,2]\n    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]\n\n    union = area1[:, None] + area2 - inter\n\n    iou = inter / union\n    return iou, union","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def generalized_box_iou(boxes1, boxes2):\n    \"\"\"\n    Generalized IoU from https://giou.stanford.edu/\n    The boxes should be in [x0, y0, x1, y1] format\n    Returns a [N, M] pairwise matrix, where N = len(boxes1)\n    and M = len(boxes2)\n    \"\"\"\n    # degenerate boxes gives inf / nan results\n    # so do an early check\n    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()\n    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()\n    iou, union = box_iou(boxes1, boxes2)\n\n    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])\n    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])\n\n    wh = (rb - lt).clamp(min=0)  # [N,M,2]\n    area = wh[:, :, 0] * wh[:, :, 1]\n\n    return iou - (area - union) / area","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def box_cxcywh_to_xyxy(x):\n    x_c, y_c, w, h = x.unbind(-1)\n    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),\n         (x_c + 0.5 * w), (y_c + 0.5 * h)]\n    return torch.stack(b, dim=-1)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"from scipy.optimize import linear_sum_assignment\n\nclass HungarianMatcher(nn.Module):\n    \"\"\"This class computes an assignment between the targets and the predictions of the network\n    For efficiency reasons, the targets don't include the no_object. Because of this, in general,\n    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,\n    while the others are un-matched (and thus treated as non-objects).\n    \"\"\"\n\n    def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1):\n        \"\"\"Creates the matcher\n        Params:\n            cost_class: This is the relative weight of the classification error in the matching cost\n            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost\n            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost\n        \"\"\"\n        super().__init__()\n        self.cost_class = cost_class\n        self.cost_bbox = cost_bbox\n        self.cost_giou = cost_giou\n        assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, \"all costs cant be 0\"\n\n    @torch.no_grad()\n    def forward(self, outputs, targets):\n        \"\"\" Performs the matching\n        Params:\n            outputs: This is a dict that contains at least these entries:\n                 \"pred_logits\": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits\n                 \"pred_boxes\": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates\n            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:\n                 \"labels\": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth\n                           objects in the target) containing the class labels\n                 \"boxes\": Tensor of dim [num_target_boxes, 4] containing the target box coordinates\n        Returns:\n            A list of size batch_size, containing tuples of (index_i, index_j) where:\n                - index_i is the indices of the selected predictions (in order)\n                - index_j is the indices of the corresponding selected targets (in order)\n            For each batch element, it holds:\n                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)\n        \"\"\"\n        bs, num_queries = outputs[\"pred_logits\"].shape[:2]\n\n        # We flatten to compute the cost matrices in a batch\n        out_prob = outputs[\"pred_logits\"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]\n        out_bbox = outputs[\"pred_boxes\"].flatten(0, 1)  # [batch_size * num_queries, 4]\n\n        # Also concat the target labels and boxes\n        tgt_ids = torch.cat([v[\"labels\"] for v in targets])\n        tgt_bbox = torch.cat([v[\"boxes\"] for v in targets])\n\n        # Compute the classification cost. Contrary to the loss, we don't use the NLL,\n        # but approximate it in 1 - proba[target class].\n        # The 1 is a constant that doesn't change the matching, it can be ommitted.\n        cost_class = -out_prob[:, tgt_ids]\n\n        # Compute the L1 cost between boxes\n        cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)\n\n        # Compute the giou cost betwen boxes\n        cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))\n\n        # Final cost matrix\n        C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou\n        C = C.view(bs, num_queries, -1).cpu()\n\n        sizes = [len(v[\"boxes\"]) for v in targets]\n        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]\n        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Class for Loss","execution_count":null},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef accuracy(output, target, topk=(1,)):\n    \"\"\"Computes the precision@k for the specified values of k\"\"\"\n    if target.numel() == 0:\n        return [torch.zeros([], device=output.device)]\n    maxk = max(topk)\n    batch_size = target.size(0)\n\n    _, pred = output.topk(maxk, 1, True, True)\n    pred = pred.t()\n    correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n    res = []\n    for k in topk:\n        correct_k = correct[:k].view(-1).float().sum(0)\n        res.append(correct_k.mul_(100.0 / batch_size))\n    return res","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"import torch.nn.functional as F\n\n\nclass SetCriterion(nn.Module):\n    \"\"\" This class computes the loss for DETR.\n    The process happens in two steps:\n        1) we compute hungarian assignment between ground truth boxes and the outputs of the model\n        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)\n    \"\"\"\n    def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses):\n        \"\"\" Create the criterion.\n        Parameters:\n            num_classes: number of object categories, omitting the special no-object category\n            matcher: module able to compute a matching between targets and proposals\n            weight_dict: dict containing as key the names of the losses and as values their relative weight.\n            eos_coef: relative classification weight applied to the no-object category\n            losses: list of all the losses to be applied. See get_loss for list of available losses.\n        \"\"\"\n        super().__init__()\n        self.num_classes = num_classes\n        self.matcher = matcher\n        self.weight_dict = weight_dict\n        self.eos_coef = eos_coef\n        self.losses = losses\n        empty_weight = torch.ones(self.num_classes + 1)\n        empty_weight[-1] = self.eos_coef\n        self.register_buffer('empty_weight', empty_weight)\n\n    def loss_labels(self, outputs, targets, indices, num_boxes, log=True):\n        \"\"\"Classification loss (NLL)\n        targets dicts must contain the key \"labels\" containing a tensor of dim [nb_target_boxes]\n        \"\"\"\n        assert 'pred_logits' in outputs\n        src_logits = outputs['pred_logits']\n\n        idx = self._get_src_permutation_idx(indices)\n        target_classes_o = torch.cat([t[\"labels\"][J] for t, (_, J) in zip(targets, indices)])\n        target_classes = torch.full(src_logits.shape[:2], self.num_classes,\n                                    dtype=torch.int64, device=src_logits.device)\n        target_classes[idx] = target_classes_o\n\n        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)\n        losses = {'loss_ce': loss_ce}\n\n        if log:\n            # TODO this should probably be a separate loss, not hacked in this one here\n            losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]\n        return losses\n\n    @torch.no_grad()\n    def loss_cardinality(self, outputs, targets, indices, num_boxes):\n        \"\"\" Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes\n        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients\n        \"\"\"\n        pred_logits = outputs['pred_logits']\n        device = pred_logits.device\n        tgt_lengths = torch.as_tensor([len(v[\"labels\"]) for v in targets], device=device)\n        # Count the number of predictions that are NOT \"no-object\" (which is the last class)\n        card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)\n        card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())\n        losses = {'cardinality_error': card_err}\n        return losses\n\n    def loss_boxes(self, outputs, targets, indices, num_boxes):\n        \"\"\"Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss\n           targets dicts must contain the key \"boxes\" containing a tensor of dim [nb_target_boxes, 4]\n           The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size.\n        \"\"\"\n        assert 'pred_boxes' in outputs\n        idx = self._get_src_permutation_idx(indices)\n        src_boxes = outputs['pred_boxes'][idx]\n        target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)\n\n        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')\n\n        losses = {}\n        losses['loss_bbox'] = loss_bbox.sum() / num_boxes\n\n        loss_giou = 1 - torch.diag(generalized_box_iou(\n            box_cxcywh_to_xyxy(src_boxes),\n            box_cxcywh_to_xyxy(target_boxes)))\n        losses['loss_giou'] = loss_giou.sum() / num_boxes\n        return losses\n\n    def loss_masks(self, outputs, targets, indices, num_boxes):\n        \"\"\"Compute the losses related to the masks: the focal loss and the dice loss.\n           targets dicts must contain the key \"masks\" containing a tensor of dim [nb_target_boxes, h, w]\n        \"\"\"\n        assert \"pred_masks\" in outputs\n\n        src_idx = self._get_src_permutation_idx(indices)\n        tgt_idx = self._get_tgt_permutation_idx(indices)\n\n        src_masks = outputs[\"pred_masks\"]\n\n        # TODO use valid to mask invalid areas due to padding in loss\n        target_masks, valid = nested_tensor_from_tensor_list([t[\"masks\"] for t in targets]).decompose()\n        target_masks = target_masks.to(src_masks)\n\n        src_masks = src_masks[src_idx]\n        # upsample predictions to the target size\n        src_masks = interpolate(src_masks[:, None], size=target_masks.shape[-2:],\n                                mode=\"bilinear\", align_corners=False)\n        src_masks = src_masks[:, 0].flatten(1)\n\n        target_masks = target_masks[tgt_idx].flatten(1)\n\n        losses = {\n            \"loss_mask\": sigmoid_focal_loss(src_masks, target_masks, num_boxes),\n            \"loss_dice\": dice_loss(src_masks, target_masks, num_boxes),\n        }\n        return losses\n\n    def _get_src_permutation_idx(self, indices):\n        # permute predictions following indices\n        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])\n        src_idx = torch.cat([src for (src, _) in indices])\n        return batch_idx, src_idx\n\n    def _get_tgt_permutation_idx(self, indices):\n        # permute targets following indices\n        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])\n        tgt_idx = torch.cat([tgt for (_, tgt) in indices])\n        return batch_idx, tgt_idx\n\n    def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs):\n        loss_map = {\n            'labels': self.loss_labels,\n            'cardinality': self.loss_cardinality,\n            'boxes': self.loss_boxes,\n            'masks': self.loss_masks\n        }\n        assert loss in loss_map, f'do you really want to compute {loss} loss?'\n        return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs)\n\n    def forward(self, outputs, targets):\n        \"\"\" This performs the loss computation.\n        Parameters:\n             outputs: dict of tensors, see the output specification of the model for the format\n             targets: list of dicts, such that len(targets) == batch_size.\n                      The expected keys in each dict depends on the losses applied, see each loss' doc\n        \"\"\"\n        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}\n\n        # Retrieve the matching between the outputs of the last layer and the targets\n        indices = self.matcher(outputs_without_aux, targets)\n\n        # Compute the average number of target boxes accross all nodes, for normalization purposes\n        num_boxes = sum(len(t[\"labels\"]) for t in targets)\n        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)\n#         if is_dist_avail_and_initialized():\n#             torch.distributed.all_reduce(num_boxes)\n        num_boxes = torch.clamp(num_boxes / 1, min=1).item()\n\n        # Compute all the requested losses\n        losses = {}\n        for loss in self.losses:\n            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))\n\n        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.\n        if 'aux_outputs' in outputs:\n            for i, aux_outputs in enumerate(outputs['aux_outputs']):\n                indices = self.matcher(aux_outputs, targets)\n                for loss in self.losses:\n                    if loss == 'masks':\n                        # Intermediate masks losses are too costly to compute, we ignore them.\n                        continue\n                    kwargs = {}\n                    if loss == 'labels':\n                        # Logging is enabled only for the last layer\n                        kwargs = {'log': False}\n                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)\n                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}\n                    losses.update(l_dict)\n\n        return losses","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n\nmodel.train()\nmatcher = HungarianMatcher()\n\n# Play with these weights to improve the results\nweight_dict = weight_dict = {'loss_ce': 1, 'loss_bbox': 1 , 'loss_giou': 1}\n\nlosses = ['labels', 'boxes', 'cardinality']\n\n#eos_coef is the weightage for the background class\ncriterion = SetCriterion(num_classes-1, matcher, weight_dict, eos_coef = 0.1, losses=losses)\n\ncriterion.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Averager:\n    def __init__(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n\n    def send(self, value):\n        self.current_total += value\n        self.iterations += 1\n\n    @property\n    def value(self):\n        if self.iterations == 0:\n            return 0\n        else:\n            return 1.0 * self.current_total / self.iterations\n\n    def reset(self):\n        self.current_total = 0.0\n        self.iterations = 0.0\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())\nvalid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())\n\n\n# split the dataset in train and test set\nindices = torch.randperm(len(train_dataset)).tolist()\n\ntrain_data_loader = DataLoader(\n    train_dataset,\n    batch_size=2,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=2,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Sample","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"images, targets, image_ids = next(iter(train_data_loader))\nimages = list(image.to(device) for image in images)\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)\nsample = images[0].permute(1,2,0).cpu().numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2]+box[0], box[3]+box[1]),\n                  (220, 0, 0), 3)\n    \nax.set_axis_off()\nax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"model.to(device)\nparams = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\nlr_scheduler = None\n\nnum_epochs = 1 if DEBUG else 10","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loss_hist = Averager()\nitr = 1\n\nfor epoch in range(num_epochs):\n    loss_hist.reset()\n    \n    for images, targets, image_ids in train_data_loader:\n\n        images = list(image.to(device) for image in images)\n        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n\n        output = model(images)\n        \n        loss_dict = criterion(output, targets)\n        weight_dict = criterion.weight_dict\n        \n        losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)\n        loss_value = losses.item()\n\n        loss_hist.send(loss_value)\n\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n\n        if itr % 50 == 0:\n            print(f\"Iteration #{itr} loss: {loss_value}\")\n\n        itr += 1\n    \n    # update the learning rate\n    if lr_scheduler is not None:\n        lr_scheduler.step()\n\n    print(f\"Epoch #{epoch} loss: {loss_hist.value}\")   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images, targets, image_ids = next(iter(valid_data_loader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images = list(img.to(device) for img in images)\ntargets = [{k: v.to(device) for k, v in t.items()} for t in targets]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = targets[1]['boxes'].cpu().numpy().astype(np.int32)\nsample = images[1].permute(1,2,0).cpu().numpy()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\nmodel.to(device)\ncpu_device = torch.device(\"cpu\")\n\noutputs = model(images)\noutputs = [{k: v.to(cpu_device) for k, v in outputs.items()}]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\nfor box in boxes:\n    cv2.rectangle(sample,\n                  (box[0], box[1]),\n                  (box[2]+box[0], box[3]+box[1]),\n                  (220, 0, 0), 3)\n    \nax.set_axis_off()\nax.imshow(sample)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), 'detr_baseline.pth')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}