{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"import torch\nfrom engine import train_one_epoch, evaluate\nimport utils\nimport transforms as T\nimport torchvision\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nimport torch.utils.data\nfrom PIL import Image, ImageFile\nimport pandas as pd\nfrom tqdm import tqdm\n\nImageFile.LOAD_TRUNCATED_IMAGES = True","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","collapsed":true,"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":false,"_kg_hide-output":true},"cell_type":"code","source":"def _get_instance_segmentation_model(num_classes):\n    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n    in_features = model.roi_heads.box_predictor.cls_score.in_features\n    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"import collections\nimport os\nimport numpy as np\n\n\nclass OpenDataset(torch.utils.data.Dataset):\n    def __init__(self, image_dir, df_path, height, width, transforms=None):\n        self.transforms = transforms\n        self.image_dir = image_dir\n        self.df = pd.read_csv(df_path)\n        self.height = height\n        self.width = width\n        self.image_info = collections.defaultdict(dict)\n        \n        # Filling up image_info is left as an exercise to the reader\n        \n\n    def __getitem__(self, idx):\n        # load images ad masks\n        img_path = self.image_info[idx][\"image_path\"]\n        img = Image.open(img_path).convert(\"RGB\")\n        img = img.resize((self.width, self.height), resample=Image.BILINEAR)\n        \n        # processing part and extraction of boxes is left as an exercise to the reader\n\n        target = {}\n        target[\"boxes\"] = boxes\n        target[\"labels\"] = labels\n        target[\"image_id\"] = image_id\n        target[\"area\"] = area\n        target[\"iscrowd\"] = iscrowd\n\n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n\n        return img, target\n\n    def __len__(self):\n        return len(self.image_info)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"num_classes = # define number of classes here\ndevice = torch.device('cuda:0')\n\ndataset_train = OpenDataset(\"../input/train/\", \"../input/train.csv\", 128, 128, transforms=None)\n\nmodel_ft = get_instance_segmentation_model(num_classes)\nmodel_ft.to(device)\n\ndata_loader = torch.utils.data.DataLoader(\n    dataset_train, batch_size=4, shuffle=True, num_workers=8,\n    collate_fn=utils.collate_fn)\n\nparams = [p for p in model_ft.parameters() if p.requires_grad]\noptimizer = torch.optim.SGD(params, lr=0.005,\n                            momentum=0.9, weight_decay=0.0005)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,\n                                               step_size=5,\n                                               gamma=0.1)\nnum_epochs = 8\nfor epoch in range(num_epochs):\n    train_one_epoch(model_ft, optimizer, data_loader, device, epoch, print_freq=10)\n    lr_scheduler.step()\n\ntorch.save(model_ft.state_dict(), \"model.bin\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"#### I have left certain parts as exercise to the reader. This model is quite fast and should give good results if used properly. In case of any questions, feel free to ask. Upvotes are appreciated if this helps you."}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}