|
- import os
- import cv2
- import random
- import numpy as np
- from PIL import Image
- import mindspore.dataset.vision.py_transforms as py_vision
- from mindspore import Tensor
-
-
- class ImageFolder(object):
- def __init__(self, root, crop_size=256, mode='train', augmentation_prob=0.5, type=2):
- """Initializes image paths and preprocessing module."""
- self.root = root
- self.gt_paths = self.root.replace("/images/", "/npy_anno_s4/")
- self.image_paths = list(map(lambda x: os.path.join(root, x), os.listdir(root)))
- self.crop_size = crop_size
- self.mode = mode
- self.CV2RotationDegreeList = ["None", cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_180, cv2.ROTATE_90_COUNTERCLOCKWISE]
- self.augmentation_prob = augmentation_prob
- print("image count in {} path :{}".format(self.mode, len(self.image_paths)))
- self.Transform = py_vision.ToTensor()
- self.nor = py_vision.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
- self.type = type
-
- def __getitem__(self, index):
- """Reads an image from a file and preprocesses it and returns."""
- image_path = self.image_paths[index]
- filename = os.path.basename(image_path)
- gt_path_pos = os.path.join(self.gt_paths, 'positive', filename.replace(".png", ".npy"))
- gt_path_neg = os.path.join(self.gt_paths, 'negative', filename.replace(".png", ".npy"))
-
- image = np.array(Image.open(image_path))
- gt_prob_map_pos = np.load(gt_path_pos).astype(np.float32)
- gt_prob_map_neg = np.load(gt_path_neg).astype(np.float32)
- gt_prob_maps = np.stack([gt_prob_map_pos, gt_prob_map_neg], axis=-1)
-
- p_transform = random.random()
- if (self.mode == 'train') and p_transform <= self.augmentation_prob:
- # random rotation
- RotationDegreeIndex = random.randint(0, 3)
- if RotationDegreeIndex != 0:
- CV2RotationDegree = self.CV2RotationDegreeList[RotationDegreeIndex]
- image = cv2.rotate(image, CV2RotationDegree)
- gt_prob_maps = cv2.rotate(gt_prob_maps, CV2RotationDegree)
-
- # random flip
- if random.random() < 0.5:
- image = cv2.flip(image, 1)
- gt_prob_maps = cv2.flip(gt_prob_maps, 1)
- if random.random() < 0.5:
- image = cv2.flip(image, 0)
- gt_prob_maps = cv2.flip(gt_prob_maps, 0)
- # random crop
- x_start = random.randint(0, image.shape[0] - self.crop_size) # width
- y_start = random.randint(0, image.shape[1] - self.crop_size) # height
- image = image[y_start:y_start + self.crop_size, x_start:x_start + self.crop_size, :]
- gt_prob_maps = gt_prob_maps[y_start:y_start + self.crop_size, x_start:x_start + self.crop_size, :]
-
- image = self.Transform(image)
- image = self.nor(image)
- gt_prob_maps = np.transpose(gt_prob_maps, axes=(2, 0, 1))
- if not self.type == 2:
- gt_prob_maps = np.expand_dims(gt_prob_maps[self.type], axis=0)
-
- return [image, gt_prob_maps, index]
-
- def __len__(self):
- """Returns the total number of font files."""
- return len(self.image_paths)
-
-
- if __name__ == '__main__':
- ImageFolder(root='BCData/images/test/',
- crop_size=512,
- mode='test',
- augmentation_prob=1.0).__getitem__(0)
|