GatedConv之一 Dataset

img处理

1 trainset = train_dataset.InpaintDataset(opt)
2 print('The overall number of images equals to %d' % len(trainset))
View Code
 1     def __init__(self, opt):
 2         assert opt.mask_type in ALLMASKTYPES
 3         self.opt = opt
 4         # img_path:'F:\\pycharm\\Dataset\\paris\\paris_eval_gt\\001_im.png'等
 5         self.imglist = utils.get_files(opt.baseroot)
 6 
 7     def __len__(self):
 8         return len(self.imglist)
 9 
10     def __getitem__(self, index):
11         # image
12         global SEED
13         img = cv2.imread(self.imglist[index])               # (227, 227, 3)
14         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)          # (227, 227, 3)
15         # set the different image size for each batch (data augmentation)
16         if index % self.opt.batch_size == 0:
17             SEED += 2
18         img, height, width = self.random_crop(img, SEED)   # 裁剪后的size
19         
20         img = torch.from_numpy(img.astype(np.float32) / 255.0).permute(2, 0, 1).contiguous()
21 
22 #       mask = torch.from_numpy(mask.astype(np.float32)).contiguous()
23 #       mask = self.random_mask()[0]
24         return img, height, width
25 
26     def random_crop(self, img, seed):
27 
28         width = img.shape[1]
29         height = img.shape[0]
30 
31         x = 0
32         y = 0
33         
34         crop = img[y: y + height, x: x + width]
35 
36         return crop, height, width
View Code

生成不规则的mask

 1     # 不规则的mask，画弧线
 2     def random_ff_mask(shape, max_angle = 10, max_len = 40, max_width = 50, times = 15):
 3         """Generate a random free form mask with configuration.
 4         Args:
 5             config: Config should have configuration including IMG_SHAPES,
 6                 VERTICAL_MARGIN, HEIGHT, HORIZONTAL_MARGIN, WIDTH.
 7         Returns:
 8             tuple: (top, left, height, width)
 9         """
10         height = shape[0]
11         width = shape[1]
12         mask = np.zeros((height, width), np.float32)
13         times = np.random.randint(times-5, times)
14         for i in range(times):
15             start_x = np.random.randint(width)
16             start_y = np.random.randint(height)
17             for j in range(1 + np.random.randint(5)):
18                 angle = 0.01 + np.random.randint(max_angle)
19                 if i % 2 == 0:
20                     angle = 2 * 3.1415926 - angle
21                 length = 10 + np.random.randint(max_len-20, max_len)
22                 brush_w = 5 + np.random.randint(max_width-30, max_width)
23                 end_x = (start_x + length * np.sin(angle)).astype(np.int32)
24                 end_y = (start_y + length * np.cos(angle)).astype(np.int32)
25                 cv2.line(mask, (start_y, start_x), (end_y, end_x), 1.0, brush_w)
26                 start_x, start_y = end_x, end_y
27         return mask.reshape((1, ) + mask.shape).astype(np.float32)
View Code

生成方形的mask

 1     def random_bbox(self, shape, margin, bbox_shape):
 2         """Generate a random tlhw with configuration.
 3         Args:
 4             config: Config should have configuration including IMG_SHAPES, VERTICAL_MARGIN, HEIGHT, HORIZONTAL_MARGIN, WIDTH.
 5         Returns:
 6             tuple: (top, left, height, width)
 7         """
 8         img_height = shape
 9         img_width = shape
10         height = bbox_shape
11         width = bbox_shape
12         ver_margin = margin
13         hor_margin = margin
14         maxt = img_height - ver_margin - height
15         maxl = img_width - hor_margin - width
16         t = np.random.randint(low = ver_margin, high = maxt)
17         l = np.random.randint(low = hor_margin, high = maxl)
18         h = height
19         w = width
20         return (t, l, h, w)
21 
22     def bbox2mask(self, shape, margin, bbox_shape, times):
23         """Generate mask tensor from bbox.
24         Args:
25             bbox: configuration tuple, (top, left, height, width)
26             config: Config should have configuration including IMG_SHAPES,
27                 MAX_DELTA_HEIGHT, MAX_DELTA_WIDTH.
28         Returns:
29             tf.Tensor: output with shape [1, H, W, 1]
30         """
31         bboxs = []
32         for i in range(times):
33             bbox = self.random_bbox(shape, margin, bbox_shape)
34             bboxs.append(bbox)
35         height = shape
36         width = shape
37         mask = np.zeros((height, width), np.float32)
38         for bbox in bboxs:
39             h = int(bbox[2] * 0.1) + np.random.randint(int(bbox[2] * 0.2 + 1))
40             w = int(bbox[3] * 0.1) + np.random.randint(int(bbox[3] * 0.2) + 1)
41             mask[(bbox[0] + h) : (bbox[0] + bbox[2] - h), (bbox[1] + w) : (bbox[1] + bbox[3] - w)] = 1.
42         return mask.reshape((1, ) + mask.shape).astype(np.float32)
43
View Code

posted @ 2021-03-29 20:38  临近边缘  阅读(201)  评论(0)    收藏  举报