Create dataset.py

eastmountyxz · web-flow · commit 279b7f8aa566 · 2021-10-30T16:58:16.000+08:00
diff --git a/blog46-FasterRCNN-Wheat/dataset.py b/blog46-FasterRCNN-Wheat/dataset.py
@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Oct 29 13:42:38 2021
+@author: xiuzhang
+"""
+import numpy as np
+import cv2
+import torch
+from torch.utils.data import Dataset
+
+
+class WheatDataset(Dataset):
+
+    def __init__(self, dataframe, image_dir, transforms=None):
+        super().__init__()
+
+        self.image_ids = dataframe['image_id'].unique()
+        self.df = dataframe
+        self.image_dir = image_dir
+        self.transforms = transforms
+
+    def __getitem__(self, index: int):
+        image_id = self.image_ids[index]
+        records = self.df[self.df['image_id'] == image_id]
+
+        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
+        image /= 255.0
+
+        boxes = records[['x', 'y', 'w', 'h']].values
+        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]
+        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]
+
+        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
+        area = torch.as_tensor(area, dtype=torch.float32)
+
+        # there is only one class
+        labels = torch.ones((records.shape[0],), dtype=torch.int64)
+
+        # suppose all instances are not crowd
+        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)
+
+        target = {}
+        target['boxes'] = boxes
+        target['labels'] = labels
+        # target['masks'] = None
+        target['image_id'] = torch.tensor([index])
+        target['area'] = area
+        target['iscrowd'] = iscrowd
+
+        if self.transforms:
+            sample = {
+                'image': image,
+                'bboxes': target['boxes'],
+                'labels': labels
+            }
+            sample = self.transforms(**sample)
+            image = sample['image']
+
+            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)
+
+        return image, target, image_id
+
+    def __len__(self) -> int:
+        return self.image_ids.shape[0]
+
