{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install bbox-visualizer","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport re\nimport pydicom\nimport warnings\n\nfrom sklearn.model_selection import train_test_split\n\nimport seaborn as sns\nimport bbox_visualizer as bbv\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\n\nwarnings.filterwarnings('ignore')\n\nDIR_INPUT = '/kaggle/input/vinbigdata-chest-xray-abnormalities-detection'\nWEIGHTS_FILE = './'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv(os.path.join(DIR_INPUT, \"train.csv\"))\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize = (8,8))\n\nsns.countplot(train.class_name)\nplt.xticks(fontsize = 14, rotation = 90)\nplt.yticks(fontsize = 14)\nplt.title(\"Distribution of labels\", fontsize = 16)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.fillna(0, inplace = True)\ntrain.loc[train[\"class_id\"] == 14, ['x_max', 'y_max']] = 1.0\ntrain[\"class_id\"] = train[\"class_id\"] + 1\ntrain.loc[train[\"class_id\"] == 15, [\"class_id\"]] = 0\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = train[train.class_name!='No finding'].reset_index(drop=True)\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize = (8,8))\n\nsns.countplot(train.class_name)\nplt.xticks(fontsize = 14, rotation = 90)\nplt.yticks(fontsize = 14)\nplt.title(\"Distribution of labels\", fontsize = 16)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df, valid_df = train_test_split(train, test_size=0.33, random_state=42)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class VinBigDataset(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):\n        \n        image_id = self.image_ids[index]\n        records = self.df[(self.df['image_id'] == image_id)]\n        records = records.reset_index(drop=True)\n\n        dicom = pydicom.dcmread(f\"{self.image_dir}/{image_id}.dicom\")\n        \n        image = dicom.pixel_array\n        \n        if \"PhotometricInterpretation\" in dicom:\n            if dicom.PhotometricInterpretation == \"MONOCHROME1\":\n                image = np.amax(image) - image\n        \n        intercept = dicom.RescaleIntercept if \"RescaleIntercept\" in dicom else 0.0\n        slope = dicom.RescaleSlope if \"RescaleSlope\" in dicom else 1.0\n        \n        if slope != 1:\n            image = slope * image.astype(np.float64)\n            image = image.astype(np.int16)\n            \n        image += np.int16(intercept)        \n        \n        image = np.stack([image, image, image])\n        image = image.astype('float32')\n        image = image - image.min()\n        image = image / image.max()\n        image = image * 255.0\n        image = image.transpose(1,2,0)\n       \n        if records.loc[0, \"class_id\"] == 0:\n            records = records.loc[[0], :]\n        \n        boxes = records[['x_min', 'y_min', 'x_max', 'y_max']].values\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n        area = torch.as_tensor(area, dtype=torch.float32)\n        labels = torch.tensor(records[\"class_id\"].values, 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['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.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 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        A.ShiftScaleRotate(scale_limit=0.1, rotate_limit=45, p=0.25),\n        A.LongestMaxSize(max_size=800, p=1.0),\n\n        # FasterRCNN will normalize.\n        A.Normalize(mean=(0, 0, 0), std=(1, 1, 1), max_pixel_value=255.0, p=1.0),\n        ToTensorV2(p=1.0)\n    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})\n\ndef get_valid_transform():\n    return A.Compose([\n        A.Normalize(mean=(0, 0, 0), std=(1, 1, 1), max_pixel_value=255.0, p=1.0),\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":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\nDIR_TRAIN = os.path.join(DIR_INPUT, \"train\")\ntrain_dataset = VinBigDataset(train_df, DIR_TRAIN,get_train_transform())\nvalid_dataset = VinBigDataset(valid_df, DIR_TRAIN,get_valid_transform())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data_loader = DataLoader(\n    train_dataset,\n    batch_size=8,\n    shuffle=True,\n    num_workers=4,\n    collate_fn=collate_fn\n)\n\nvalid_data_loader = DataLoader(\n    valid_dataset,\n    batch_size=8,\n    shuffle=False,\n    num_workers=4,\n    collate_fn=collate_fn\n)","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":{},"cell_type":"markdown","source":"# Sample"},{"metadata":{"trusted":true},"cell_type":"code","source":"images, targets = next(iter(train_data_loader))\nimages = list(image.to(device) for image in images)\n\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":"def visualize_plot(idx,images,targets):\n    class_types = {\n        0: 'Aortic enlargement',\n        1: 'Atelectasis',\n        2: 'Calcification',\n        3: 'Cardiomegaly',\n        4: 'Consolidation',\n        5: 'ILD',\n        6: 'Infiltration',\n        7: 'Lung Opacity',\n        8: 'Nodule/Mass',\n        9: 'Other lesion',\n        10: 'Pleural effusion',\n        11: 'Pleural thickening',\n        12: 'Pneumothorax',\n        13: 'Pulmonary fibrosis'\n        }\n    boxes = targets[idx]['boxes'].cpu().numpy().astype(np.int32)\n    labels = targets[idx]['labels']-1\n    sample = images[idx].permute(1,2,0).cpu().numpy()\n    \n    img = sample.copy()\n    plt.figure(figsize=(16, 16))\n    for box,label in zip(boxes,labels):\n        bbv.add_label(img, class_types[label.item()], box, \n                      draw_bg=True,\n                      text_bg_color=(255,255,0),\n                      text_color=(0,0,0),\n                      )\n        cv2.rectangle(img ,\n                      (box[0], box[1]),\n                      (box[2], box[3]),\n                      (255,255,0), 3)\n\n\n    plt.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"visualize_plot(4,images,targets)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# replace the classifier with a new one, that has\n# num_classes which is user-defined\nnum_classes = 15  # 14 class (anomalies) + background\n# load a model; pre-trained on COCO\nmodel = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n# get number of input features for the classifier\nin_features = model.roi_heads.box_predictor.cls_score.in_features\n\n# replace the pre-trained head with a new one\nmodel.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\nmodel.to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Optimizer"},{"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.0005, momentum=0.9, weight_decay=0.0005)\nlr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)\n\nnum_epochs = 12","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"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","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 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        loss_dict = model(images, targets)\n\n        losses = sum(loss for loss in loss_dict.values())\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 % 100 == 0:\n            print(f\"Iteration #{itr} loss: {loss_hist.value}\")\n\n        itr += 1\n        \n        # !!!REMOVE THIS!!!\n        break\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":"# Specify a path\nPATH = \"./model_state.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}