{"cells":[{"metadata":{},"cell_type":"markdown","source":"**Beginner notebook for Dataset class and Dataloader in pytorch. Thanks Peter and Abhishek Thakur for their awesome tutorials and notebooks** ","execution_count":null},{"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 os\nimport re\nimport torch\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nimport cv2\nimport albumentations\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Regex for bounding box in the dataframe","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Read train.csv\n\n- split the bounding box list into x, y, width and height\n- compute x1 and y1 for pascal format [x, y, x1 = x+w, y1 = y+h]\n- compute area = width * height ","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"df = pd.read_csv(\"/kaggle/input/global-wheat-detection/train.csv\")\n\ndf['x'] = -1\ndf['y'] = -1\ndf['w'] = -1\ndf['h'] = -1\n\ndf[['x', 'y', 'w', 'h']] = np.stack(df['bbox'].apply(lambda x: expand_bbox(x)))\n\ndf['x'] = df['x'].astype(float)\ndf['y'] = df['y'].astype(float)\ndf['w'] = df['w'].astype(float)\ndf['h'] = df['h'].astype(float)\n\ndf.drop(columns=['bbox'], inplace=True)\ndf['x1'] = df['x'] + df['w']\ndf['y1'] = df['y'] + df['h']\ndf['area'] = df['w'] * df['h']\n\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Split the dataset\n\n- splitting based on image_id\n- 90% sample images for training \n- 10% sample images for testing\n- make new column(folds) where 0 = validation, 1 = train","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"val_len = int(len(df['image_id'].unique()) * 0.1)\n\ntrain_set = df['image_id'].unique()[val_len:]\nval_set = df['image_id'].unique()[:val_len]\n\ndf['folds'] = -1 \n\ndf.loc[df.image_id.isin(train_set), 'folds'] = 1\ndf.loc[df.image_id.isin(val_set), 'folds'] = 0\n\ndf[\"image_id\"] = df[\"image_id\"] + \".jpg\"\n\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Pytorch dataset class","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class WheatDataset(Dataset):\n    def __init__(self, datafame, is_train, transforms = None):\n        super(WheatDataset, self).__init__()\n\n        self.df = datafame\n        \n        self.transforms = transforms\n        \n        if is_train:\n            self.df = self.df[self.df.folds == 1].reset_index(drop=True)\n            self.img_ids = self.df['image_id'].unique()\n        else:\n            self.df = self.df[self.df.folds == 0].reset_index(drop=True)\n            self.img_ids = self.df['image_id'].unique()\n        \n    def __len__(self):\n        return len(self.img_ids)\n    \n    def __getitem__(self, item):\n        \n        img_idx = self.img_ids[item]\n        \n        images = cv2.imread(os.path.join(\"/kaggle/input/global-wheat-detection/train\", img_idx), cv2.IMREAD_COLOR)\n        images = cv2.cvtColor(images, cv2.COLOR_BGR2RGB).astype(np.float32)\n        \n        data = self.df[self.df['image_id'] == img_idx]\n        bbox = data[['x','y','x1','y1']].values\n        area = data['area'].values\n        \n        labels = np.ones((bbox.shape[0],))\n        is_crowd = np.zeros((bbox.shape[0],))\n        taget = {'image':images , 'bboxes': bbox, 'labels': labels,}\n        \n        if self.transforms:\n            transforms = self.transforms(**taget)\n            images, bboxs, labels = transforms['image'], transforms['bboxes'], transforms['labels']\n            \n        images = np.transpose(images, (2, 0, 1)).astype(np.float32)\n        \n        dict_target = {\"boxes\": torch.tensor(bboxs, dtype = torch.long), \n                    \"area\" : torch.tensor(area, dtype = torch.long),\n                    \"labels\": torch.tensor(labels, dtype = torch.long),\n                    \"iscrowd\": torch.tensor(is_crowd, dtype = torch.long),\n                   }\n        \n        return torch.tensor(images, dtype = torch.float), dict_target","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"IMG_HEIGHT = 512\nIMG_WIDTH = 512\nIMG_MEAN = (0.5, 0.5, 0.5)\nIMG_STD = (0.5, 0.5, 0.5)\nBACTH_SIZE = 4","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augumentations\n\n- Flip = random vertical, random horizontal or random both vertical and horizontal\n- resize to 512x512\n- Normalize","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"train_aug = albumentations.Compose([\n            albumentations.Flip(),\n            albumentations.Resize(IMG_HEIGHT, IMG_HEIGHT, always_apply = True),            \n            albumentations.Normalize(IMG_MEAN, IMG_STD, always_apply = True)\n            ], bbox_params = albumentations.BboxParams(format = 'pascal_voc', \n                                                       label_fields=['labels']))\n\ntrain_data = WheatDataset(datafame= df , is_train = True, transforms = train_aug)\n\nprint(\"Number of images in training dataset\", len(train_data))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataloader\n\n- Since the bounding box size is varies for each image, we need collate_fn","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(batch):\n    return list(zip(*batch))\n\ntrainloader = DataLoader(train_data,\n                           batch_size=BACTH_SIZE,\n                           shuffle=False,\n                           num_workers=0,\n                           collate_fn = collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"idx = 18\ndata = train_data[idx]\nimg, bbx = data\nbbx = bbx['boxes']\nimg = img / 2 + 0.5  \nnpimg = img.numpy()\nbbx = bbx.numpy()\nnpimg = np.transpose(npimg,(1, 2, 0))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(15, 10))\n\nfor b in bbx:\n        ax.add_patch(\n            patches.Rectangle(\n            (b[0], b[1]),\n            b[2]-b[0],\n            b[3]-b[1],\n            linewidth=2,\n            fill=False,\n            color='red'))\n    \nax.set_axis_off()\nax.imshow(npimg)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def imshow(img_data, bbx_data, size):\n    fig, ax = plt.subplots(nrows = size, ncols = 2, figsize=(30, 30))\n    for i, (img, bbx) in enumerate(zip(img_data,bbx_data)):\n        img = img / 2 + 0.5\n        npimg = img.numpy()\n        npimg = np.transpose(npimg,(1, 2, 0))\n        ax[i, 0].set_title(\"Image\")\n        ax[i, 0].set_axis_off()\n        ax[i, 0].imshow(npimg)\n        bboxes = bbx['boxes'].numpy()\n        for box in bboxes:\n            ax[i,1].add_patch(\n                patches.Rectangle(\n                (box[0], box[1]),\n                box[2]-box[0],\n                box[3]-box[1],\n                linewidth=2,\n                fill=False,\n                color='red'))\n        ax[i, 1].set_axis_off()\n        ax[i, 1].set_title(\"Image with bounding box\")\n        ax[i, 1].imshow(npimg)\n\ndataiter = iter(trainloader)\ndata = dataiter.next()\nimages, bbx = data\nimshow(images, bbx, len(images))\nplt.show()\nplt.close()","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":1}