{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nROOT = '../input/vinbigdata-chest-xray-abnormalities-detection/'\nDATASET_ROOT = '../input/vinbigdata-original-image-dataset/vinbigdata'\nos.listdir(ROOT)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ORIGINAL_TRAIN = os.path.join(ROOT, 'train')\nTRAIN_DIR = os.path.join(DATASET_ROOT, 'train')\nTEST_DIR = os.path.join(DATASET_ROOT, 'test')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport pydicom\nimport random\nimport cv2\nimport torch\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\n\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SEED = 0\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n    \nseed_everything(SEED)    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Let's Inspect the Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"#data_df = pd.read_csv(os.path.join(ROOT, 'train.csv'))\ndata_df = pd.read_csv('../input/effdet-latestvinbigdata-wbf-fused/train_wbf_original.csv', index_col='Unnamed: 0')\ndata_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Remove no finding samples\ndata_df = data_df.loc[data_df['class_id'] != 14].reset_index(drop=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_xray(path, voi_lut = True, fix_monochrome = True):\n    dicom = pydicom.read_file(path)\n    \n    # VOI LUT (if available by DICOM device) is used to transform raw DICOM data to \"human-friendly\" view\n    if voi_lut:\n        data = apply_voi_lut(dicom.pixel_array, dicom)\n    else:\n        data = dicom.pixel_array\n               \n    # depending on this value, X-ray may look inverted - fix that:\n    if fix_monochrome and dicom.PhotometricInterpretation == \"MONOCHROME1\":\n        data = np.amax(data) - data\n        \n    data = data - np.min(data)\n    data = data / np.max(data)\n    data = (data * 255).astype(np.uint8)\n        \n    return data\n\ndef plot_img(img, size=(7, 7), is_rgb=True, title=\"\", cmap='gray'):\n    plt.figure(figsize=size)\n    plt.imshow(img, cmap=cmap)\n    plt.suptitle(title)\n    plt.show()\n\n\ndef plot_imgs(imgs, cols=4, size=7, is_rgb=True, title=\"\", cmap='gray', img_size=(500,500)):\n    rows = len(imgs)//cols + 1\n    fig = plt.figure(figsize=(cols*size, rows*size))\n    for i, img in enumerate(imgs):\n        if img_size is not None:\n            img = cv2.resize(img, img_size)\n        fig.add_subplot(rows, cols, i+1)\n        plt.imshow(img, cmap=cmap)\n    plt.suptitle(title)\n    plt.show()\n\ndef show_boxes(finding_df, img_dir):\n    imgs = []\n    img_ids = finding_df['image_id'].values\n    class_ids = finding_df['class_id'].unique()\n\n    # map label_id to specify color\n    label2color = {class_id:[np.random.randint(0,255) for i in range(3)] for class_id in class_ids}\n    thickness = 3\n    scale = 5\n\n\n    for i in range(8):\n        img_id = random.choice(img_ids)\n        img_path = f'{img_dir}/{img_id}.dicom'\n        img = read_xray(path=img_path)\n        img = cv2.resize(img, None, fx=1/scale, fy=1/scale)\n        img = np.stack([img, img, img], axis=-1)\n\n        boxes = finding_df.loc[finding_df['image_id'] == img_id, ['x_min', 'y_min', 'x_max', 'y_max']].values/scale\n        labels = finding_df.loc[finding_df['image_id'] == img_id, ['class_id']].values.squeeze()\n        \n        try:\n            for label_id, box in zip(labels, boxes):\n                color = label2color[label_id]\n                img = cv2.rectangle(\n                    img,\n                    (int(box[0]), int(box[1])),\n                    (int(box[2]), int(box[3])),\n                    color, thickness\n            )\n            img = cv2.resize(img, (500,500))\n            imgs.append(img)\n        except:\n            continue\n\n    plot_imgs(imgs, cmap=None)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_boxes(data_df, ORIGINAL_TRAIN)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Function to calculate IOU and areas of corresponding boxes\ndef calculate_iou(bbox1, bbox2):\n    # Coordinates must be consistent\n    assert(bbox1['x_min'] < bbox1['x_max'])\n    assert(bbox1['y_min'] < bbox1['y_max'])\n    assert(bbox2['x_min'] < bbox2['x_max'])\n    assert(bbox2['y_min'] < bbox2['y_max'])\n    \n    # Calculate coordinates of the top left corner of the intersection area\n    x_top_left = max(bbox1['x_min'], bbox2['x_min'])\n    y_top_left = max(bbox1['y_min'], bbox2['y_min'])\n    \n    # Calculate coordinates of the bottom right corner of the intersection area\n    x_bottom_right = min(bbox1['x_max'], bbox2['x_max'])\n    y_bottom_right = min(bbox1['y_max'], bbox2['y_max'])\n    \n    # Calculate IOU\n    area_bbox1 = ((bbox1['x_max'] - bbox1['x_min']) * (bbox1['y_max'] - bbox1['y_min']))\n    assert area_bbox1 > 0\n    area_bbox2 = ((bbox2['x_max'] - bbox2['x_min']) * (bbox2['y_max'] - bbox2['y_min']))\n    assert area_bbox2 > 0\n    \n    if x_top_left > x_bottom_right or y_top_left > y_bottom_right:\n        return 0.0, area_bbox1, area_bbox2\n    \n    area_intersection = (x_bottom_right - x_top_left) * (y_bottom_right - y_top_left) \n    assert area_intersection >= 0\n    \n    area_union = area_bbox1 + area_bbox2 - area_intersection\n    \n    iou = area_intersection / area_union\n    assert iou >= 0.0 and iou <= 1.0\n    \n    return iou, area_bbox1, area_bbox2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Remove bounding boxes with high IOU and same class\ndef remove_bboxs(df, threshold=0.5):\n    img_ids = df['image_id'].unique()\n    new_records = list()\n\n    for img_id in img_ids:\n        records = df[df['image_id'] == img_id].reset_index(drop=True)\n        to_drop = list()\n        size = records.shape[0]\n        for i in range(size-1):\n            if i in to_drop:\n                continue\n            bbox1 = records.iloc[[i],:]\n            bbox1 = bbox1.to_dict('records')[0]\n            for j in range(i+1, size):\n                bbox2 = records.iloc[[j],:]\n                bbox2 = bbox2.to_dict('records')[0]\n\n                iou, bb1_area, bb2_area = calculate_iou(bbox1, bbox2)\n                if iou >= threshold and bbox1['class_id'] == bbox2['class_id']:\n                    if bb1_area >= bb2_area:\n                        to_drop.append(i)\n                        break\n                    else:\n                        to_drop.append(j)\n        records = records.loc[~records.index.isin(to_drop)]\n        new_records.append(records)\n\n    return pd.concat(new_records)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(data_df.shape)\ndata_df = remove_bboxs(data_df, threshold = 0.1)\nprint(data_df.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_boxes(data_df, ORIGINAL_TRAIN)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_df = data_df.merge(meta_df, on='image_id')\ndata_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Transform bounding box coordinates to new scale\ndata_df['x_min'] = data_df.apply(lambda row: (row.x_min)/row.dim1*512, axis =1)\ndata_df['y_min'] = data_df.apply(lambda row: (row.y_min)/row.dim0*512, axis =1)\n\ndata_df['x_max'] = data_df.apply(lambda row: (row.x_max)/row.dim1*512, axis =1)\ndata_df['y_max'] = data_df.apply(lambda row: (row.y_max)/row.dim0*512, axis =1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_df['x_max'].max(), data_df['y_max'].max()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_df['class_id'] = data_df['class_id'] + 1\nprint(data_df['class_id'].unique())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data Splitting by Applying K-Fold "},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.model_selection import GroupKFold\n\ndef split_data(df, groups_col, n_splits=5, fold=0):\n    group_fold = GroupKFold(n_splits=n_splits)\n    df['fold'] = -1\n    \n    for fold, (train_idx, val_idx) in enumerate(group_fold.split(df, groups=df[groups_col].tolist())):\n        df.loc[val_idx, 'fold'] = fold\n\n    df_val = df[df.fold==fold]\n    df_train = df[df.fold!=fold]\n    \n    return df_train, df_val","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Take into account one radiologist\n# df_split = data_df.loc[data_df['rad_id'] == 'R9'].reset_index(drop=True)\ndf_split = data_df.copy().reset_index(drop=True)\ndf_train, df_val = split_data(df_split, 'image_id')\ndf_train.shape, df_val.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Creating a Custom Dataset Class"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ChestDataset(Dataset):\n    def __init__(self, df, image_dir, transforms=None, is_original=False, is_validation=False, \n                 is_inference=False, padding=False, img_size=512, img_ext='.png'):\n        super().__init__()\n        \n        self.df = df\n        self.image_ids = df['image_id'].unique()\n        self.image_dir = image_dir\n        self.transforms = transforms\n        self.is_original = is_original\n        self.is_validation = is_validation\n        self.is_inference = is_inference\n        self.img_size = img_size\n        self.padding = padding\n        self.img_ext = img_ext\n        \n    def __getitem__(self, index):\n        if not self.is_inference:\n            if self.is_validation:\n                choice = 'as_is'\n            else:\n                choice = np.random.choice(\n                    ['as_is', 'mixup', 'cutmix'],\n                    1,\n                    p=[0.4, 0.3, 0.3]\n                )\n                \n            if choice == 'as_is':\n                image, boxes, labels = self.load_image_and_boxes(index)\n            elif choice == 'mixup':\n                image, boxes, labels = self.load_mixup_image_and_boxes(index)\n            elif choice == 'cutmix':\n                image, boxes, labels = self.load_cutmix_image_and_boxes(index, self.img_size)\n                \n            ## To prevent ValueError: y_max is less than or equal to y_min for bbox from albumentations bbox_utils\n            labels = np.array(labels, dtype=np.int).reshape(len(labels), 1)\n            combined = np.hstack((boxes.astype(np.int), labels))\n            combined = combined[np.logical_and(combined[:,2] > combined[:,0],\n                                               combined[:,3] > combined[:,1])]\n            boxes = combined[:, :4]\n            labels = combined[:, 4].tolist()\n            area = (boxes[:,2] - boxes[:,0]) * (boxes[:,3] - boxes[:,1])\n            \n            if self.transforms:\n                sample = self.transforms(image=image, bboxes=boxes, labels=labels)\n                image = sample['image']\n                boxes = sample['bboxes']\n                labels = sample['labels']\n                \n                if len(boxes) == 0:\n                    boxes = np.array([[0.0, 0.0, 1.0, 1.0]])\n                    area = [1.0]\n                    labels = [0]\n            \n            target = dict()\n\n            area = torch.as_tensor(area, dtype=torch.float32)\n            boxes = torch.as_tensor(boxes, dtype=torch.float32)\n            labels = torch.as_tensor(labels, dtype=torch.int64)\n            iscrowd = torch.zeros(labels.shape[0], dtype=torch.uint8)\n            idx = torch.tensor([index])\n\n            target['area'] = area\n            target['labels'] = labels\n            target['iscrowd'] = iscrowd\n            target['image_id'] = idx\n            target['boxes'] = boxes\n                \n            return image, target, self.image_ids[index]\n        \n        else:\n            if self.transforms:\n                image = self.transforms(image=image)['image']\n            \n            return image, image_id\n    \n    def load_image_and_boxes(self, index):\n        image_id = self.image_ids[index]\n        info = self.df.loc[self.df['image_id'] == image_id]\n        info = info.reset_index(drop=True)\n        \n        if self.is_original:\n            img_path = os.path.join(self.image_dir, image_id + '.dicom')\n            image = read_xray(img_path)\n            image = np.stack([image, image, image]).transpose(1, 2, 0)  \n            image = image.astype(np.float32)\n        else:\n            img_path = os.path.join(self.image_dir, image_id + self.img_ext)\n            image = cv2.imread(img_path, cv2.IMREAD_COLOR)\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n            image /= 255.0\n            \n            if self.padding:\n                # padding image\n                if image.shape[0] < self.img_size:\n                    pad_count = self.img_size - image.shape[0]\n                    image = np.pad(image, ((0, pad_count), (0, 0), (0, 0)), \n                                            mode='constant', constant_values=(0))\n                elif image.shape[1] < self.img_size:\n                    pad_count = self.img_size - image.shape[1]\n                    image = np.pad(image, ((0, 0), (0, pad_count), (0, 0)), \n                                   mode='constant', constant_values=(0))\n                    \n        if not self.is_inference:\n            boxes = info.loc[:, ['x_min', 'y_min', 'x_max', 'y_max']].values    \n            labels = info['class_id'].values  \n            image, boxes, labels = self.resize_img_and_boxes(image, boxes, labels)\n            \n            return image, boxes, labels\n        else:\n            return image\n        \n    def resize_img_and_boxes(self, image, boxes, labels):\n        resize_transform = A.Compose([A.Resize(height=self.img_size, width=self.img_size, p=1.0)], \n                                      p=1.0, \n                                      bbox_params=A.BboxParams(\n                                          format='pascal_voc',\n                                          min_area=0.1, \n                                          min_visibility=0.1,\n                                          label_fields=['labels'])\n                                     )\n\n        resized = resize_transform(image=image, bboxes=boxes, labels=labels)\n\n        resized_bboxes = np.vstack((list(bx) for bx in resized['bboxes']))\n        \n        return resized['image'], resized_bboxes, resized['labels']\n    \n    def load_mixup_image_and_boxes(self, index):\n        image, boxes, labels = self.load_image_and_boxes(index)\n        r_image, r_boxes, r_labels = self.load_image_and_boxes(random.randint(0, self.image_ids.shape[0] - 1))\n        return (image+r_image)/2, np.vstack((boxes, r_boxes)).astype(np.int32), np.concatenate((labels, r_labels))\n\n    def load_cutmix_image_and_boxes(self, index, imsize=512):\n        \"\"\" \n        This implementation of cutmix author:  https://www.kaggle.com/nvnnghia \n        Refactoring and adaptation: https://www.kaggle.com/shonenkov\n        \"\"\"\n        w, h = imsize, imsize\n        s = imsize // 2\n    \n        xc, yc = [int(random.uniform(imsize * 0.25, imsize * 0.75)) for _ in range(2)]  # center x, y\n        indexes = [index] + [random.randint(0, self.image_ids.shape[0] - 1) for _ in range(3)]\n\n        result_image = np.full((imsize, imsize, 3), 1, dtype=np.float32)\n        result_boxes = []\n        result_labels = np.array([], dtype=np.int)\n\n        for i, index in enumerate(indexes):\n            image, boxes, labels = self.load_image_and_boxes(index)\n            if i == 0:\n                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)\n                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)\n            elif i == 1:  # top right\n                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc\n                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h\n            elif i == 2:  # bottom left\n                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)\n                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)\n            elif i == 3:  # bottom right\n                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)\n                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)\n            result_image[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]\n            padw = x1a - x1b\n            padh = y1a - y1b\n\n            boxes[:, 0] += padw\n            boxes[:, 1] += padh\n            boxes[:, 2] += padw\n            boxes[:, 3] += padh\n\n            result_boxes.append(boxes)\n            result_labels = np.concatenate((result_labels, labels))\n\n        result_boxes = np.concatenate(result_boxes, 0)\n        np.clip(result_boxes[:, 0:], 0, 2 * s, out=result_boxes[:, 0:])\n        result_boxes = result_boxes.astype(np.int32)\n        index_to_use = np.where((result_boxes[:,2]-result_boxes[:,0])*(result_boxes[:,3]-result_boxes[:,1]) > 0)\n        result_boxes = result_boxes[index_to_use]\n        result_labels = result_labels[index_to_use]\n        \n        return result_image, result_boxes, result_labels\n    \n    def __len__(self):\n        return self.image_ids.shape[0] ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Defining Transformations "},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transform():\n    return A.Compose([A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),\n                      A.ShiftScaleRotate(scale_limit=0.1, rotate_limit=45, p=0.5),\n                      A.ToGray(p=0.01),\n                      A.HorizontalFlip(p=0.5),\n                      A.VerticalFlip(p=0.5),\n                      A.Cutout(num_holes=8, max_h_size=16, max_w_size=16, fill_value=0, p=0.3),\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                      p=1.0,\n                      bbox_params=A.BboxParams(format='pascal_voc',min_area=0, min_visibility=0,label_fields=['labels'])\n                    )\n\ndef get_valid_transform():\n    return A.Compose([A.Resize(height=512, width=512, p=1.0),\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                      p=1.0, \n                      bbox_params=A.BboxParams(format='pascal_voc',min_area=0, min_visibility=0,label_fields=['labels'])\n                    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Defining the Model "},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.models.detection import fasterrcnn_resnet50_fpn\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n\ndef get_model(n_classes):\n    model = 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, n_classes)\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Creating the DataLoaders "},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return tuple(zip(*batch))\n\ndef collate_fn2(batch):\n    batch = list(filter(lambda x: x is not None, batch))\n    return tupe(zip(*batch))\n\ndef get_data_loader(df, data_dir, transform, batch_size=8, shuffle=True, is_original=False, \n                    is_validation=False, is_inference=False, padding=False, img_ext='.png'):\n    \n    dataset = ChestDataset(df, data_dir, transform, is_original, is_validation, \n                           is_inference, padding=padding, img_ext=img_ext)\n    \n    loader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        shuffle=shuffle,\n        num_workers=4,\n        collate_fn=collate_fn\n    )\n    \n    return loader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Testing DataLoader "},{"metadata":{"trusted":true},"cell_type":"code","source":"test_loader = get_data_loader(df_train, TRAIN_DIR, get_train_transform(), img_ext='.jpg')\nimgs, targets, ids = next(iter(test_loader))\n\nbatch_size = len(imgs)\ncolsize = batch_size // 2\nfig, axes = plt.subplots(2, colsize, figsize=(16, 8))\n\nfor i in range(batch_size):\n    row = i // colsize    \n    col = i % colsize\n    axes[row][col].imshow(imgs[i].numpy().transpose(1, 2, 0))\n        \nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_loader = get_data_loader(df_val, TRAIN_DIR, get_valid_transform(), is_validation=True, img_ext='.jpg')\nimgs, targets, ids = next(iter(test_loader))\n\nbatch_size = len(imgs)\ncolsize = batch_size // 2\nfig, axes = plt.subplots(2, colsize, figsize=(16, 8))\n\nfor i in range(batch_size):\n    row = i // colsize    \n    col = i % colsize\n    axes[row][col].imshow(imgs[i].numpy().transpose(1, 2, 0))\n        \nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training "},{"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":"def train(train_data, valid_data, model, n_epoch=10, \n          learning_rate=1e-3, device='cpu', print_freq=100):\n    # Create the optimizer\n    print('Device:', device)\n    model = model.to(device)\n    model.train()\n    \n    train_loss_hist = Averager()\n    valid_loss_hist = Averager()\n    \n    params = [p for p in model.parameters() if p.requires_grad]\n    optimizer = torch.optim.SGD(params, lr=learning_rate, momentum=0.6, weight_decay=0.0005, nesterov=True)\n    #lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)\n\n    for e in range(n_epoch):        \n        for images, targets, _ in train_data:\n            images = [img.to(device) for img 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            train_losses = sum([l for l in loss_dict.values()])\n            train_loss = train_losses.item()\n            train_loss_hist.send(train_loss)\n            \n            optimizer.zero_grad()\n            train_losses.backward()\n            optimizer.step()\n\n#         if lr_scheduler is not None:\n#              lr_scheduler.step()\n        \n        print('Saving model, epoch:', e)\n        torch.save(model.state_dict(), f'./model_{e}.th')\n        \n        for images, targets, _ in valid_data:\n            with torch.no_grad():\n                images = [img.to(device) for img 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                val_losses = sum([l for l in loss_dict.values()])\n                val_loss = val_losses.item()\n                valid_loss_hist.send(val_loss)\n        \n        print(f\"Epoch:{e}, Train loss:{train_loss_hist.value}, Validation loss:{valid_loss_hist.value}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\nmodel = get_model(15)\ntrain_loader = get_data_loader(df_train, TRAIN_DIR, get_train_transform(), batch_size = 8, img_ext='.jpg')\nvalid_loader = get_data_loader(df_val, TRAIN_DIR, get_valid_transform(), batch_size = 8, img_ext='.jpg')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrain(train_loader, valid_loader, model, device=device, learning_rate=1e-2, n_epoch = 10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_random_prediction(model_path, df, data, device, threshold=0.5):\n    model = get_model(15)\n    weights = torch.load(model_path)\n    model.load_state_dict(weights)\n    \n    idx = np.random.randint(len(data))\n    i = 0\n    for sample in data:\n        if i == idx:\n            break\n        else:\n            i+=1\n            \n    img, targets, image_id = sample[0][0], sample[1][0], sample[2][0]\n\n    model.to(device)\n    model.eval()\n    \n    outputs = model(img.unsqueeze(1).to(device))\n    img = img.cpu().numpy().transpose(1, 2, 0)\n    boxes = targets['boxes']\n    \n    fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n    for box in boxes:\n        cv2.rectangle(img,\n                  (box[0], box[1]),\n                  (box[2]+box[0], box[3]+box[1]),\n                  (220, 0, 0), 1)\n        \n    boxes = outputs[0]['boxes'].detach().cpu().numpy()\n    prob = outputs[0]['scores'].detach().cpu().numpy()\n    boxes = boxes[prob>threshold]\n    for box in boxes:\n        cv2.rectangle(img,\n                  (box[0], box[1]),\n                  (box[2]+box[0], box[3]+box[1]),\n                  (0, 0, 220), 1)\n    \n    ax.set_axis_off()\n    ax.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_random_prediction('./model_9.th', df_train, valid_loader, device, 0.6)","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}