{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport cv2\nimport torch\nimport matplotlib.pyplot as plt\nimport torchvision\nfrom torchvision.models.detection import FasterRCNN\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.rpn import AnchorGenerator\nfrom torch.utils.data import DataLoader\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom tqdm.notebook import tqdm\nimport albumentations as A\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\ndevice","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_PATH = '../input/vinbigdata-chest-xray-resized-png-1024x1024/train/'\ndf_train = pd.read_csv('../input/vinbigdata-chest-xray-abnormalities-detection/train.csv')\ndf_train_meta = pd.read_csv('../input/vinbigdata-chest-xray-resized-png-1024x1024/train_meta.csv')\nIMAGE_SIZE = 1024","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train.fillna(0, inplace = True)\ndf_train.class_id = df_train.class_id + 1\ndf_train.loc[df_train.class_id == 15, ['class_id']] = 0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(df_train.shape)\ndf_train.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(df_train_meta.shape)\ndf_train_meta.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def update_bboxes(df_train, df_train_meta):\n    \n    df_temp = pd.merge(df_train, df_train_meta, on = \"image_id\")\n    df_temp['x_min'] = ((df_temp['x_min'] / df_temp['dim1']) * IMAGE_SIZE).astype('int')\n    df_temp['y_min'] = ((df_temp['y_min'] / df_temp['dim0']) * IMAGE_SIZE).astype('int')\n    df_temp['x_max'] = ((df_temp['x_max'] / df_temp['dim1']) * IMAGE_SIZE).astype('int')\n    df_temp['y_max'] = ((df_temp['y_max'] / df_temp['dim0']) * IMAGE_SIZE).astype('int')\n    \n    return df_temp.drop(['dim0', 'dim1'], axis = 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = update_bboxes(df_train, df_train_meta)\ndf_train.loc[df_train.class_id == 0, ['x_max', 'y_max']] = 1\ndf_train.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train['area']= (df_train['x_max'] - df_train['x_min']) * (df_train['y_max'] - df_train['y_min'])\nprint(df_train[df_train.area == 0].shape)\ndf_train = df_train[df_train.area != 0].reset_index(drop = True)\ndf_train.drop(['area'], axis = 1, inplace = True)\ndf_train","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class_dict = dict(zip(df_train[\"class_id\"].unique(), df_train['class_name'].unique()))\nlen(class_dict)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class VinBigData:\n    def __init__(self, root, df, transforms = None):\n        self.root = root\n        self.transforms = transforms\n        self.imgs = list(sorted(os.listdir(root)))\n        self.df = df\n        \n    def __getitem__(self, index):\n        img_path = self.imgs[index]\n        \n        image = cv2.imread(os.path.join(self.root,img_path))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        image = image / 255\n        image = image.astype('float32')\n        \n        records = self.df.loc[self.df.image_id == img_path.split('.')[0],:].reset_index(drop = True)\n        \n        if records.loc[0, \"class_id\"] == 0:\n            records = records.loc[[0], :]\n        boxes = records[['x_min','y_min', 'x_max', 'y_max']].values\n        area = (boxes[:,2] - boxes[:,0]) * (boxes[:,3] - boxes[:,1])\n        labels = records['class_id'].values\n        \n        target = {}\n        target['boxes'] = torch.tensor(boxes)\n        target['labels'] = torch.tensor(labels, dtype=torch.int64)\n        target['image_id'] = torch.tensor([index])\n        target['area'] = torch.tensor(area, dtype=torch.float32)\n        target['iscrowd'] = torch.zeros(labels.shape[0], dtype=torch.int64)\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.tensor(sample['bboxes'])\n        \n        if target[\"boxes\"].shape[0] == 0:\n            # Albumentation cuts the target (class 14, 1x1px in the corner)\n            target[\"boxes\"] = torch.from_numpy(np.array([[0.0, 0.0, 1.0, 1.0]]))\n            target[\"area\"] = torch.tensor([1.0], dtype=torch.float32)\n            target[\"labels\"] = torch.tensor([0], dtype=torch.int64)\n        \n        return image, target\n    \n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_valid_transform():\n    return A.Compose([\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data = VinBigData(TRAIN_PATH, df_train, get_valid_transform())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_image(index):\n    plt.figure(figsize=(10,10))\n    image, target = train_data[index]\n    image = image.permute(1,2,0).numpy()\n    for box,label in zip(target['boxes'], target['labels']):\n        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (225,0,0), 1)\n        class_index = label.numpy().reshape(1)[0]\n        label = class_dict[class_index].upper()\n        cv2.putText(image, \n                    label,\n                    (box[0], box[1]), \n                    fontFace = cv2.FONT_HERSHEY_SIMPLEX, \n                    fontScale = 0.5,\n                    color = (255, 0, 0),\n                    thickness = 1,\n                    lineType = cv2.LINE_AA\n                   )\n    plt.imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_image(717)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ntrain_dataloader = DataLoader(train_data, batch_size = 8, shuffle = False, collate_fn = collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# load a model pre-trained pre-trained on COCO\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained = True)\n# replace the classifier with a new one, that has\n# num_classes which is user-defined\nnum_classes = len(class_dict)\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n# move model to the right device\nmodel.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# construct an optimizer\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# and a learning rate scheduler which decreases the learning rate by 10x every 3 epochs\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# let's train it for 10 epochs\nnum_epochs = 1\nerror_batch = []\nfor epoch in tqdm(range(num_epochs)):\n    #initiating loss and num_iterations\n    total_loss = 0\n    itr = 0\n    for images, targets in train_dataloader:\n        images = [image.to(device) for image in images]\n        targets = [{k: v.to(device) for k,v in t.items()} for t in targets]\n        loss_dict = model(images, targets)\n        losses = sum(loss for loss in loss_dict.values())\n        loss_value = losses.item()\n        #update loss and itr\n        total_loss += loss_value\n        itr +=1\n        #update the learning rate\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n        if itr%500 == 0:\n            print(\"Itr : {} loss: {}\".format(itr, total_loss/itr))\n        lr_scheduler.step()\n    print(\"Epoch : {} loss: {}\".format(epoch, total_loss/itr))\n\n# Specify a path\nPATH = \"./VinBigDataFasterRCNN.pth\"\n# Save\ntorch.save(model.state_dict(), PATH)","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}