{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nfrom pathlib import Path\nimport ast\nimport numpy as np\nfrom PIL import Image, ImageDraw\nfrom torch.utils.data import Dataset, DataLoader\nimport torch\nimport torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection import FasterRCNN\nfrom matplotlib import pyplot as plt\nfrom typing import Union\nimport numpy as np\n\nimport torchvision.transforms as T","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"data_path = Path('../input/global-wheat-detection/')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv(data_path/'train.csv')\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df['bbox'] = df['bbox'].apply(lambda x: ast.literal_eval(x))\nx = np.array(list(df['bbox']))\n\nfor i,dim in enumerate(['x', 'y', 'w', 'h']):\n    df[dim] = x[:,i]\n    \ndf.drop('bbox', axis = 1, inplace = True)\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = df['image_id'].unique()\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = df[df['image_id'].isin(train_ids)]\nvalid_df = df[df['image_id'].isin(valid_ids)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n    def __init__(self, df, image_dir, transforms = None):\n        super().__init__()\n        \n        self.df = df\n        self.image_ids = self.df['image_id'].unique()\n        self.image_dir = Path(image_dir)\n        self.transforms = transforms\n    \n    def __getitem__(self, idx):\n        image_id = self.image_ids[idx]\n        records = self.df[self.df['image_id'] == image_id]\n            \n        im_name = image_id + '.jpg'\n        img = Image.open(self.image_dir/im_name).convert(\"RGB\")\n        img = T.ToTensor()(img)\n        \n        if self.transforms:\n            img = self.transforms(img)\n        \n        boxes = records[['x', 'y', 'w', 'h']].values\n        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n        boxes = torch.tensor(boxes, dtype=torch.int64)\n        \n        labels = torch.ones((records.shape[0],), dtype=torch.int64)\n        \n        \n        target = {}\n        target['boxes'] = boxes\n        target['labels'] = labels\n        target['image_id'] = torch.tensor([idx])\n        \n        return img, target, image_id       \n    \n    def __len__(self):\n        return self.image_ids.shape[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds = WheatDataset(train_df, data_path/'train')\nval_ds = WheatDataset(valid_df, data_path/'train')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dl = DataLoader(train_ds, batch_size=16, num_workers=4, collate_fn=collate_fn, pin_memory = True)\nval_dl = DataLoader(train_ds, batch_size=8, collate_fn=collate_fn, pin_memory = True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"images, targets, _ = next(iter(train_dl))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = targets[2]['boxes']\nboxes.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im = (images[2].permute(1,2,0).numpy() * 255).astype('uint8')\nsample = Image.fromarray(im)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"draw = ImageDraw.Draw(sample)\nfor box in boxes:\n    draw.rectangle(box.numpy(), fill = None, outline = \"red\")\nsample","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained = True)\nnum_classes = 2 # should be initialized as target_col.nunique + 1\nin_features = model.roi_heads.box_predictor.cls_score.in_features\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"params = [p for p in model.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)\nnum_epochs = 5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nmodel.to(device)\nfor epoch in range(num_epochs):\n    epoch_loss = 0\n    for images, targets, _ in train_dl:\n        optimizer.zero_grad()\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        loss_dict = model(images, targets)\n#         print(loss_dict)\n        losses = sum(loss for loss in loss_dict.values())\n        epoch_loss += losses.item()\n\n        losses.backward()\n        optimizer.step()\n    print(f\"loss for epoch {epoch}: {epoch_loss / len(train_dl)}\")\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model, '/kaggle/working/model.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"results=[]\ndetection_threshold = 0.45\nmodel.eval()\nmodel.to(device)\n\nfor images, targets,idx in val_dl:    \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    outputs = model(images)\n\n    for i, image in enumerate(images):\n\n        boxes = outputs[i]['boxes'].data\n        scores = outputs[i]['scores'].data\n        labels = outputs[i]['labels'].data\n\n        keep = torchvision.ops.nms(boxes, scores, 0.3)\n        boxes = boxes[keep]\n        scores = scores[keep]\n        image_id = idx[i]\n    \n        op = (idx[i], boxes, scores)\n        results.append(op)\n\n        break\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"im_name = image_id + '.jpg'\nimg = Image.open(data_path/'train'/im_name).convert(\"RGB\")\nimg = T.ToTensor()(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# make a function to prepare and annotate images\nim = (img.permute(1,2,0).detach().numpy() * 255).astype('uint8')\nvsample = Image.fromarray(im)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"draw = ImageDraw.Draw(vsample)\nfor box in boxes:\n    draw.rectangle(list(box), fill = None, outline = \"red\")\nvsample","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"features = []\ndef store_features(mod, inp, outp):\n    features.append(outp.data) # this will store [batch-size, channel, height, widht]\n    # if you want to store the feature map per image use torch.unbind","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"to_hook = 'backbone.fpn.layer_blocks.3'\nfor name, layer in model.named_modules():\n    if name == to_hook:\n        layer.register_forward_hook(store_features) # if you want to pass more arguments to store features use partial functions","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Reference:\n\nhttps://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train?sortBy=relevance&group=everyone&search=Pytorch&page=1&pageSize=20&competitionId=19989","execution_count":null}],"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}