{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport torch\nimport cv2\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import Dataset\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nimport pydicom\nimport matplotlib.pyplot as plt\nimport colorsys\nimport warnings\nwarnings.filterwarnings('ignore')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# coding: utf-8\n__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'\n\n# Custom by me\n\nimport numpy as np\n\ndef bb_intersection_over_union(A, B):\n    xA = max(A[0], B[0])\n    yA = max(A[1], B[1])\n    xB = min(A[2], B[2])\n    yB = min(A[3], B[3])\n\n    # compute the area of intersection rectangle\n    interArea = max(0, xB - xA) * max(0, yB - yA)\n\n    if interArea == 0:\n        return 0.0, 0.0\n\n    # compute the area of both the prediction and ground-truth rectangles\n    boxAArea = (A[2] - A[0]) * (A[3] - A[1])\n    boxBArea = (B[2] - B[0]) * (B[3] - B[1])\n\n    iou = interArea / float(boxAArea + boxBArea - interArea)\n    return interArea, iou\n\n\ndef prefilter_boxes(boxes, scores, labels, weights, thr):\n    # Create dict with boxes stored by its label\n    new_boxes = dict()\n    for t in range(len(boxes)):\n        for j in range(len(boxes[t])):\n            score = scores[t][j]\n            if score < thr:\n                continue\n            label = int(labels[t][j])\n            box_part = boxes[t][j]\n            b = [int(label), float(score) * weights[t], float(box_part[0]), float(box_part[1]), float(box_part[2]), float(box_part[3])]\n            if label not in new_boxes:\n                new_boxes[label] = []\n            new_boxes[label].append(b)\n\n    # Sort each list in dict by score and transform it to numpy array\n    for k in new_boxes:\n        current_boxes = np.array(new_boxes[k])\n        new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]]\n\n    return new_boxes\n\n\ndef get_weighted_box(boxes, conf_type='avg'):\n    \"\"\"\n    Create weighted box for set of boxes\n    :param boxes: set of boxes to fuse \n    :param conf_type: type of confidence one of 'avg' or 'max'\n    :return: weighted box\n    \"\"\"\n\n    box = np.zeros(6, dtype=np.float32)\n    conf = 0\n    conf_list = []\n    for b in boxes:\n        box[2:] += (b[1] * b[2:])\n        conf += b[1]\n        conf_list.append(b[1])\n    box[0] = boxes[0][0]\n    if conf_type == 'avg':\n        box[1] = conf / len(boxes)\n    elif conf_type == 'max':\n        box[1] = np.array(conf_list).max()\n    box[2:] /= conf\n    return box\n\n\ndef find_matching_box(boxes_list, new_box, match_iou):\n    best_iou = match_iou\n    best_index = -1\n    for i in range(len(boxes_list)):\n        box = boxes_list[i]\n        if box[0] != new_box[0]:\n            continue\n        interArea, iou = bb_intersection_over_union(box[2:], new_box[2:])\n        boxAArea = (box[2:][2] - box[2:][0]) * (box[2:][3] - box[2:][1])\n        if iou > best_iou or interArea >= 0.9 * boxAArea:\n            best_index = i\n            best_iou = iou\n\n    return best_index, best_iou\n\n\ndef weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0, conf_type='avg', allows_overflow=False):\n    \"\"\"\n    :param boxes_list: list of boxes predictions from each model, each box is 4 numbers.\n    It has 3 dimensions (models_number, model_preds, 4)\n    Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]\n    :param scores_list: list of scores for each model\n    :param labels_list: list of labels for each model\n    :param weights: list of weights for each model. Default: None, which means weight == 1 for each model\n    :param iou_thr: IoU value for boxes to be a match\n    :param skip_box_thr: exclude boxes with score lower than this variable\n    :param conf_type: how to calculate confidence in weighted boxes. 'avg': average value, 'max': maximum value\n    :param allows_overflow: false if we want confidence score not exceed 1.0\n\n    :return: boxes: boxes coordinates (Order of boxes: x1, y1, x2, y2).\n    :return: scores: confidence scores\n    :return: labels: boxes labels\n    \"\"\"\n\n    if weights is None:\n        weights = np.ones(len(boxes_list))\n    if len(weights) != len(boxes_list):\n        print('Warning: incorrect number of weights {}. Must be: {}. Set weights equal to 1.'.format(len(weights), len(boxes_list)))\n        weights = np.ones(len(boxes_list))\n    weights = np.array(weights)\n\n    if conf_type not in ['avg', 'max']:\n        print('Unknown conf_type: {}. Must be \"avg\" or \"max\"'.format(conf_type))\n        exit()\n\n    filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr)\n    if len(filtered_boxes) == 0:\n        return np.zeros((0, 4)), np.zeros((0,)), np.zeros((0,))\n\n    overall_boxes = []\n    for label in filtered_boxes:\n        boxes = filtered_boxes[label]\n        new_boxes = []\n        weighted_boxes = []\n\n        # Clusterize boxes\n        for j in range(0, len(boxes)):\n            index, best_iou = find_matching_box(weighted_boxes, boxes[j], iou_thr)\n            if index != -1:\n                new_boxes[index].append(boxes[j])\n                weighted_boxes[index] = get_weighted_box(new_boxes[index], conf_type)\n            else:\n                new_boxes.append([boxes[j].copy()])\n                weighted_boxes.append(boxes[j].copy())\n\n        # Rescale confidence based on number of models and boxes\n        for i in range(len(new_boxes)):\n            if not allows_overflow:\n                weighted_boxes[i][1] = weighted_boxes[i][1] * min(weights.sum(), len(new_boxes[i])) / weights.sum()\n            else:\n                weighted_boxes[i][1] = weighted_boxes[i][1] * len(new_boxes[i]) / weights.sum()\n        overall_boxes.append(np.array(weighted_boxes))\n\n    overall_boxes = np.concatenate(overall_boxes, axis=0)\n    overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]]\n    boxes = overall_boxes[:, 2:]\n    scores = overall_boxes[:, 1]\n    labels = overall_boxes[:, 0]\n    return boxes, scores, labels","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"image_paths = r'../input/vinbigdata-chest-xray-abnormalities-detection/train'\ndf = pd.read_csv(r'../input/vinbigdata-chest-xray-abnormalities-detection/train.csv')\nSIZE = 1024\n\ndf.drop(columns=[\"class_name\"], inplace=True)\n\ndf = df[df.class_id != 14]\n\nskf = StratifiedKFold(n_splits=5, shuffle=True, random_state=101)\ndf_folds = df[['image_id']].copy()\n\ndf_folds.loc[:, 'bbox_count'] = 1\ndf_folds = df_folds.groupby('image_id').count()\ndf_folds.loc[:, 'object_count'] = df.groupby('image_id')['class_id'].nunique()\n\ndf_folds.loc[:, 'stratify_group'] = np.char.add(\n    df_folds['object_count'].values.astype(str),\n    df_folds['bbox_count'].apply(lambda x: f'_{x // 15}').values.astype(str)\n)\n\ndf_folds.loc[:, 'fold'] = 0\nfor fold_number, (train_index, val_index) in enumerate(skf.split(X=df_folds.index, y=df_folds['stratify_group'])):\n    df_folds.loc[df_folds.iloc[val_index].index, 'fold'] = fold_number\n\n# example with fold 0\ndf_folds.reset_index(inplace=True)\n\nvalid_df = pd.merge(df, df_folds[df_folds['fold'] == 0], on='image_id')\ntrain_df = pd.merge(df, df_folds[df_folds['fold'] != 0], on='image_id')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_valid_transforms():\n    return A.Compose(\n        [\n            A.Resize(height=SIZE, width=SIZE, p=1.0),\n            ToTensorV2(p=1.0),\n        ],\n        p=1.0,\n        bbox_params=A.BboxParams(\n            format='pascal_voc',\n            min_area=0,\n            min_visibility=0,\n            label_fields=['labels']\n        )\n    )\n\n\nclass DatasetRetriever(Dataset):\n\n    def __init__(self, marking, image_dir=None, transforms=None):\n        super().__init__()\n\n        self.image_ids = marking[\"image_id\"].unique()\n        self.marking = marking\n        self.image_dir = image_dir\n        self.transforms = transforms\n\n    def __getitem__(self, index: int):\n        image_id = self.image_ids[index]\n        image, boxes, labels = self.load_image_and_boxes(index)\n        target = {'boxes': torch.from_numpy(boxes),\n                  'labels': torch.from_numpy(labels),\n                  'image_id': torch.tensor([index])}\n\n        if self.transforms:\n            sample = self.transforms(**{\n                'image': image,\n                'bboxes': target['boxes'],\n                'labels': labels\n            })\n            if len(sample['bboxes']) > 0:\n                image = sample['image']\n                target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)\n                target['boxes'][:, [0, 1, 2, 3]] = target['boxes'][:, [1, 0, 3, 2]]\n        return image, target, image_id\n\n    def __len__(self) -> int:\n        return self.image_ids.shape[0]\n\n    def load_image_and_boxes(self, index):\n\n        image_id = self.image_ids[index]\n        dicom = pydicom.dcmread(f\"{self.image_dir}/{image_id}.dicom\")\n        image = dicom.pixel_array\n        if \"PhotometricInterpretation\" in dicom:\n            if dicom.PhotometricInterpretation == \"MONOCHROME1\":\n                image = np.amax(image) - image\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.transpose(1, 2, 0)\n\n        records = self.marking[self.marking['image_id'] == image_id]\n        boxes = records[['x_min', 'y_min', 'x_max', 'y_max']].values\n\n        records = self.marking[(self.marking['image_id'] == image_id)]\n        records = records.reset_index(drop=True)\n        labels = records[\"class_id\"].values\n        return image, boxes, np.array(labels, dtype=np.int)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = DatasetRetriever(marking=train_df, image_dir=image_paths, transforms=get_valid_transforms())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def hsv2rgb(h, s, v):\n    return tuple(round(i * 255) for i in colorsys.hsv_to_rgb(h, s, v))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def visualize():\n    mapping = {0: 'Aortic enlargement', 1: 'Atelectasis', 2: 'Calcification', 3: 'Cardiomegaly', 4: 'Consolidation',\n               5: 'ILD', 6: 'Infiltration', 7: 'Lung Opacity', 8: 'Nodule/Mass', 9: 'Other lesion',\n               10: 'Pleural effusion', 11: 'Pleural thickening', 12: 'Pneumothorax', 13: 'Pulmonary fibrosis'}\n\n    font = cv2.FONT_HERSHEY_SIMPLEX\n    fontScale = 1\n    thickness = 4\n    for i in range(30):\n        image, target, image_ids = train_dataset[i]\n        boxes = target['boxes'].cpu().numpy().astype(np.int64)\n        labels = target['labels'].cpu().numpy().astype(np.int64)\n        numpy_image = image.permute(1, 2, 0).cpu().numpy() * 255\n        fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n        image_before = numpy_image.copy()\n        for box, label in zip(boxes, labels):\n            box = np.array(box, dtype=np.int)\n            label = np.array(label, dtype=np.int)\n            image_before = cv2.rectangle(image_before, (int(box[1]), int(box[0])), (int(box[3]), int(box[2])),\n                                         hsv2rgb(int(label) / 14, 1, 1), 2)\n            image_before = cv2.putText(image_before, mapping[int(label)], (box[1], box[0]), font, fontScale,\n                                       hsv2rgb(int(label) / 14, 1, 1), thickness, cv2.LINE_AA)\n\n        image_before = cv2.putText(image_before, 'BOXES_BEFORE', (100, 100), font, fontScale,\n                                   hsv2rgb(1, 1, 1), thickness, cv2.LINE_AA)\n\n        scores = [list(np.ones(labels.shape[0]))]\n        boxes = [[box for box in boxes]]\n        labels = [[label for label in labels]]\n        boxes, scores, labels = weighted_boxes_fusion(boxes, scores, labels, iou_thr=0.5)\n\n        image_after = numpy_image.copy()\n        for box, label in zip(boxes, labels):\n            box = np.array(box, dtype=np.int)\n            label = np.array(label, dtype=np.int)\n            image_after = cv2.rectangle(image_after, (int(box[1]), int(box[0])), (int(box[3]), int(box[2])),\n                                        hsv2rgb(int(label) / 14, 1, 1), 2)\n            image_after = cv2.putText(image_after, mapping[int(label)], (box[1], box[0]), font, fontScale,\n                                      hsv2rgb(int(label) / 14, 1, 1), thickness, cv2.LINE_AA)\n\n        image_after = cv2.putText(image_after, 'BOXES_OPTIMIZED', (100, 100), font, fontScale,\n                                  hsv2rgb(1, 1, 1), thickness, cv2.LINE_AA)\n\n        image = cv2.hconcat([image_before, image_after])\n        ax.imshow(image.astype(np.uint8))\n        plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if __name__ == '__main__':\n    visualize()","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}