def __getitem__(self, idx):

        inst = self.lineIndex[idx]
        author = inst[0]
        words = inst[1]
        batch = []
        for word in words:
            if word >= len(self.w_authors[author]):
                word = (word + 37) % len(self.w_authors[author])
            img_path, lb, gt, id = self.w_authors[author][word]
            img = cv2.imread(img_path,
                             0)[lb[0]:lb[1],
                                lb[2]:lb[3]]  #read as grayscale, crop word
            if img.shape[0] == 0 or img.shape[1] == 0:
                return self.__getitem__((idx + 1) % self.__len__())

            if img is None:
                return None

            if img.shape[0] != self.img_height:
                if img.shape[0] < self.img_height and not self.warning:
                    self.warning = True
                    print("WARNING: upsampling image to fit size")
                percent = float(self.img_height) / img.shape[0]
                img = cv2.resize(img, (0, 0),
                                 fx=percent,
                                 fy=percent,
                                 interpolation=cv2.INTER_CUBIC)

            if img is None:
                return None

            if len(img.shape) == 2:
                img = img[..., None]
            if self.augmentation is not None:
                #img = augmentation.apply_random_color_rotation(img)
                img = augmentation.apply_tensmeyer_brightness(img)
                img = grid_distortion.warp_image(img)

            img = img.astype(np.float32)
            img = 1.0 - img / 128.0

            if len(gt) == 0:
                return None
            gt_label = string_utils.str2label_single(gt, self.char_to_idx)

            if self.styles:
                style_i = self.npr.choice(len(self.styles[author][id]))
                style = self.styles[author][id][style_i]
            else:
                style = None
            batch.append({
                "image": img,
                "gt": gt,
                "style": style,
                "gt_label": gt_label,
                "name": '{}_{}'.format(author, word),
                "author": author
            })
        #batch = [b for b in batch if b is not None]
        #These all should be the same size or error
        assert len(set([b['image'].shape[0] for b in batch])) == 1
        assert len(set([b['image'].shape[2] for b in batch])) == 1

        dim0 = batch[0]['image'].shape[0]
        dim1 = max([b['image'].shape[1] for b in batch])
        dim2 = batch[0]['image'].shape[2]

        all_labels = []
        label_lengths = []

        input_batch = np.full((len(batch), dim0, dim1, dim2),
                              PADDING_CONSTANT).astype(np.float32)
        for i in range(len(batch)):
            b_img = batch[i]['image']
            #toPad = (dim1-b_img.shape[1])
            input_batch[i, :, 0:b_img.shape[1], :] = b_img

            l = batch[i]['gt_label']
            all_labels.append(l)
            label_lengths.append(len(l))

        #all_labels = np.concatenate(all_labels)
        label_lengths = torch.IntTensor(label_lengths)
        max_len = label_lengths.max()
        all_labels = [
            np.pad(l, ((0, max_len - l.shape[0]), ), 'constant')
            for l in all_labels
        ]
        all_labels = np.stack(all_labels, axis=1)

        images = input_batch.transpose([0, 3, 1, 2])
        images = torch.from_numpy(images)
        labels = torch.from_numpy(all_labels.astype(np.int32))
        #label_lengths = torch.from_numpy(label_lengths.astype(np.int32))
        if batch[0]['style'] is not None:
            styles = np.stack([b['style'] for b in batch], axis=0)
            styles = torch.from_numpy(styles).float()
        else:
            styles = None

        return {
            "image": images,
            "mask": makeMask(images, self.mask_post),
            "label": labels,
            "style": styles,
            "label_lengths": label_lengths,
            "gt": [b['gt'] for b in batch],
            "name": [b['name'] for b in batch],
            "author": [b['author'] for b in batch]
        }
Ejemplo n.º 2
0
    def __getitem__(self, idx):

        if self.train and self.refresh_self:
            if self.used_instances >= self.set_size:
                self.refresh_data(None, None, logged=False)
                self.used_instances = 0
            self.used_instances += self.batch_size

        if self.augmentation is not None and 'affine' in self.augmentation:
            strech = (self.max_strech *
                      2) * np.random.random() - self.max_strech + 1
            skew = (self.max_rot_rad *
                    2) * np.random.random() - self.max_rot_rad
        if self.include_stroke_aug:
            thickness_change = np.random.randint(-4, 5)
            fg_shade = np.random.random() * 0.25 + 0.75
            bg_shade = np.random.random() * 0.2
            blur_size = np.random.randint(2, 4)
            noise_sigma = np.random.random() * 0.02

        batch = []
        for b in range(self.batch_size):
            img_path = os.path.join(self.directory, '{}.png'.format(idx + b))
            img = cv2.imread(img_path, 0)
            if img is None:
                print('Error, could not read {}'.format(img_path))
                return self[(idx + 1) % len(self)]

            if self.augmentation == 'affine':
                if img.shape[1] * strech > self.max_width:
                    strech = self.max_width / img.shape[1]
            if img.shape[1] > self.max_width:
                percent = float(self.max_width) / img.shape[1]
                img = cv2.resize(img, (0, 0),
                                 fx=percent,
                                 fy=1,
                                 interpolation=cv2.INTER_CUBIC)

            if img.shape[1] > self.clip_width:
                img = img[:, :self.clip_width]

            if self.use_fg_mask:
                th, fg_mask = cv2.threshold(
                    img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
                fg_mask = 255 - fg_mask
                ele = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
                fg_mask = cv2.dilate(fg_mask, ele)
                fg_mask = fg_mask / 255
            else:
                fg_mask = None

            if len(img.shape) == 2:
                img = img[..., None]
            if self.augmentation is not None:
                #img = augmentation.apply_random_color_rotation(img)
                if 'affine' in self.augmentation:
                    img, fg_mask = augmentation.affine_trans(
                        img, fg_mask, skew, strech)
                if 'brightness' in self.augmentation:
                    img = augmentation.apply_tensmeyer_brightness(img)
                    assert (fg_mask is None)
                if 'warp' in self.augmentation and random.random(
                ) < self.warp_freq:
                    try:
                        img = grid_distortion.warp_image(img)
                    except cv2.error as e:
                        print(e)
                        print(img.shape)
                    assert (fg_mask is None)
                if 'invert' in self.augmentation and random.random() < 0.25:
                    img = 1 - img

            if self.include_stroke_aug:
                new_img = augmentation.change_thickness(
                    img, thickness_change, fg_shade, bg_shade, blur_size,
                    noise_sigma)
                new_img = new_img * 2 - 1.0

            if len(img.shape) == 2:
                img = img[..., None]

            img = img.astype(np.float32)
            img = 1.0 - img / 128.0

            if self.train:
                gt = self.labels[idx]
            else:
                with open(self.gt_filename) as f:
                    for i in range(0, idx + 1):
                        gt = f.readline()
                gt = gt.strip()
            if gt is None:
                #metadata = pyexiv2.ImageMetadata(img_path)
                #metadata.read()
                #metadata = piexif.load(img_path)
                #if 'gt' in metadata:
                #    gt = metadata['gt']
                #else:
                print('Error unknown label for image: {}'.format(img_path))
                return self.__getitem__((idx + 7) % self.set_size)

            gt_label = string_utils.str2label_single(gt, self.char_to_idx)

            font_idx = '?'
            toRet = {
                "image": img,
                "gt": gt,
                "gt_label": gt_label,
                "author": font_idx,
                "name": '{}_{}'.format(idx + b, font_idx),
                "style": None,
                "spaced_label": None
            }
            if self.use_fg_mask:
                toRet['fg_mask'] = fg_mask
            if self.include_stroke_aug:
                toRet['changed_image'] = new_img
            batch.append(toRet)

        dim0 = batch[0]['image'].shape[0]
        dim1 = max([b['image'].shape[1] for b in batch])
        dim2 = batch[0]['image'].shape[2]

        all_labels = []
        label_lengths = []
        if self.spaced_by_name is not None:
            spaced_labels = []
        else:
            spaced_labels = None
        max_spaced_len = 0

        input_batch = np.full((len(batch), dim0, dim1, dim2),
                              PADDING_CONSTANT).astype(np.float32)
        if self.use_fg_mask:
            fg_masks = np.full((len(batch), dim0, dim1, 1),
                               0).astype(np.float32)
        if self.include_stroke_aug:
            changed_batch = np.full((len(batch), dim0, dim1, dim2),
                                    PADDING_CONSTANT).astype(np.float32)
        for i in range(len(batch)):
            b_img = batch[i]['image']
            toPad = (dim1 - b_img.shape[1])
            if 'center' in batch[0] and batch[0]['center']:
                toPad //= 2
            else:
                toPad = 0
            input_batch[i, :, toPad:toPad + b_img.shape[1], :] = b_img
            if self.use_fg_mask:
                fg_masks[i, :, toPad:toPad + b_img.shape[1],
                         0] = batch[i]['fg_mask']
            if self.include_stroke_aug:
                changed_batch[i, :, toPad:toPad + b_img.shape[1],
                              0] = batch[i]['changed_image']

            l = batch[i]['gt_label']
            all_labels.append(l)
            label_lengths.append(len(l))

            if spaced_labels is not None:
                sl = batch[i]['spaced_label']
                spaced_labels.append(sl)
                max_spaced_len = max(max_spaced_len, sl.shape[0])

        #all_labels = np.concatenate(all_labels)
        label_lengths = torch.IntTensor(label_lengths)
        max_len = label_lengths.max()
        all_labels = [
            np.pad(l, ((0, max_len - l.shape[0]), ), 'constant')
            for l in all_labels
        ]
        all_labels = np.stack(all_labels, axis=1)
        if self.spaced_by_name is not None:
            spaced_labels = [
                np.pad(l, ((0, max_spaced_len - l.shape[0]), (0, 0)),
                       'constant') for l in spaced_labels
            ]
            ddd = spaced_labels
            spaced_labels = np.concatenate(spaced_labels, axis=1)
            spaced_labels = torch.from_numpy(spaced_labels)
            assert (spaced_labels.size(1) == len(batch))

        images = input_batch.transpose([0, 3, 1, 2])
        images = torch.from_numpy(images)
        labels = torch.from_numpy(all_labels.astype(np.int32))
        #label_lengths = torch.from_numpy(label_lengths.astype(np.int32))
        if self.use_fg_mask:
            fg_masks = fg_masks.transpose([0, 3, 1, 2])
            fg_masks = torch.from_numpy(fg_masks)

        if batch[0]['style'] is not None:
            styles = np.stack([b['style'] for b in batch], axis=0)
            styles = torch.from_numpy(styles).float()
        else:
            styles = None
        mask, top_and_bottom, center_line = makeMask(images, self.mask_post,
                                                     self.mask_random)
        toRet = {
            "image": images,
            "mask": mask,
            "top_and_bottom": top_and_bottom,
            "center_line": center_line,
            "label": labels,
            "style": styles,
            "label_lengths": label_lengths,
            "gt": [b['gt'] for b in batch],
            "spaced_label": spaced_labels,
            "name": [b['name'] for b in batch],
            "author": [b['author'] for b in batch],
        }
        if self.use_fg_mask:
            toRet['fg_mask'] = fg_masks
        if self.include_stroke_aug:
            changed_images = changed_batch.transpose([0, 3, 1, 2])
            changed_images = torch.from_numpy(changed_images)
            toRet['changed_image'] = changed_images
        return toRet
    def __getitem__(self, idx):
        if type( self.augmentation) is str and 'affine' in self.augmentation:
            strech = (self.max_strech*2)*np.random.random() - self.max_strech +1
            #self.max_rot_rad = self.max_rot_deg/180 * np.pi
            skew = (self.max_rot_rad*2)*np.random.random() - self.max_rot_rad
        if self.include_stroke_aug:
            thickness_change= np.random.randint(-4,5)
            fg_shade = np.random.random()*0.25 + 0.75
            bg_shade = np.random.random()*0.2
            blur_size = np.random.randint(2,4)
            noise_sigma = np.random.random()*0.02

        batch=[]

        if self.triplet=='hard':
            authors = random.sample(self.authors.keys(),self.triplet_author_size)
            alines=[]
            for author in authors:
                if len(self.authors[author])>=self.triplet_sample_size*self.batch_size:
                    lines = random.sample(range(len(self.authors[author])),self.triplet_sample_size*self.batch_size)
                else:
                    lines = list(range(len(self.authors[author])))
                    random.shuffle(lines)
                    dif = self.triplet_sample_size*self.batch_size-len(self.authors[author])
                    lines += lines[:dif]
                alines += [(author,l) for l in lines]
        else:


            inst = self.lineIndex[idx]
            author=inst[0]
            lines=inst[1]


            alines = [(author,l) for l in lines]
            used_lines = set(lines)
            if self.triplet:
                if len(self.authors[author])<=2*self.batch_size:
                    for l in range(len(self.authors[author])):
                        if l not in used_lines:
                            alines.append((author,l))
                    if len(alines)<2*self.batch_size:
                        dif = 2*self.batch_size - len(alines)
                        for i in range(dif):
                            alines.append(alines[self.batch_size+i])
                else:
                    unused_lines = set(range(len(self.authors[author])))-used_lines
                    for i in range(self.batch_size):
                        l = random.select(unused_lines)
                        unused_lines.remove(l)
                        alines.append((author,l))
                
                other_authors = set(range(len(self.authors)))
                other_authors.remove(author)
                author = random.select(other_authors)
                unused_lines = set(range(len(self.authors[author])))-used_lines
                for i in range(self.batch_size):
                    l = random.select(unused_lines)
                    unused_lines.remove(l)
                    alines.append((author,l))

            

        images=[]
        for author,line in alines:
            if line>=len(self.authors[author]):
                line = (line+37)%len(self.authors[author])
            img_path, lb, gt = self.authors[author][line]
            img_path = os.path.join(self.dirPath,'images_gray',img_path)

            if self.no_spaces:
                gt = gt.replace(' ','')
            if type(self.augmentation) is str and 'normalization' in  self.augmentation and self.normalized_dir is not None and os.path.exists(os.path.join(self.normalized_dir,'{}_{}.png'.format(author,line))):
                img = cv2.imread(os.path.join(self.normalized_dir,'{}_{}.png'.format(author,line)),0)
                readNorm=True
            else:
                img = cv2.imread(img_path,0)
                if img is None:
                    print('Error, could not read image: {}'.format(img_path))
                    return None
                lb[0] = max(lb[0],0)
                lb[2] = max(lb[2],0)
                lb[1] = min(lb[1],img.shape[0])
                lb[3] = min(lb[3],img.shape[1])
                img = img[lb[0]:lb[1],lb[2]:lb[3]] #read as grayscale, crop line
                readNorm=False


            if img.shape[0] != self.img_height:
                if img.shape[0] < self.img_height and not self.warning:
                    self.warning = True
                    print("WARNING: upsampling image to fit size")
                percent = float(self.img_height) / img.shape[0]
                if img.shape[1]*percent > self.max_width:
                    percent = self.max_width/img.shape[1]
                img = cv2.resize(img, (0,0), fx=percent, fy=percent, interpolation = cv2.INTER_CUBIC)
                if img.shape[0]<self.img_height:
                    diff = self.img_height-img.shape[0]
                    img = np.pad(img,((diff//2,diff//2+diff%2),(0,0)),'constant',constant_values=255)
            elif img.shape[1]> self.max_width:
                percent = self.max_width/img.shape[1]
                img = cv2.resize(img, (0,0), fx=percent, fy=percent, interpolation = cv2.INTER_CUBIC)
                if img.shape[0]<self.img_height:
                    diff = self.img_height-img.shape[0]
                    img = np.pad(img,((diff//2,diff//2+diff%2),(0,0)),'constant',constant_values=255)

            if self.augmentation=='affine':
                if img.shape[1]*strech > self.max_width:
                    strech = self.max_width/img.shape[1]
            images.append( (line,gt,img,author) )
            #we split the processing here so that strech will be adjusted for longest image in author batch


        for line,gt,img,author in images:
            if self.fg_masks_dir is not None:
                fg_path = os.path.join(self.fg_masks_dir,'{}_{}.png'.format(author,line))
                fg_mask = cv2.imread(fg_path,0)
                fg_mask = fg_mask/255
                if fg_mask.shape!=img[:,:].shape:
                    print('Error, fg_mask ({}, {}) not the same size as image ({})'.format(fg_path,fg_mask.shape,img[:,:,0].shape))
                    th,fg_mask = cv2.threshold(img,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
                    fg_mask = 255-fg_mask
                    ele = cv2.getStructuringElement(  cv2.MORPH_ELLIPSE, (9,9) )
                    fg_mask = cv2.dilate(fg_mask,ele)
                    fg_mask = fg_mask/255
            else:
                fg_mask=None

                    
            if type(self.augmentation) is str and 'normalization' in  self.augmentation and not readNorm:
                img = normalize_line.deskew(img)
                img = normalize_line.skeletonize(img)
                if self.normalized_dir is not None:
                    cv2.imwrite(os.path.join(self.normalized_dir,'{}_{}.png'.format(author,line)),img)
            if type(self.augmentation) is str and 'affine' in  self.augmentation:
                img,fg_mask = augmentation.affine_trans(img,fg_mask,skew,strech)
            elif self.augmentation is not None and (type(self.augmentation) is not None or 'warp' in self.augmentation):
                #img = augmentation.apply_random_color_rotation(img)
                img = augmentation.apply_tensmeyer_brightness(img)
                img = grid_distortion.warp_image(img)
                assert(fg_mask is None)

            if self.include_stroke_aug:
                new_img = augmentation.change_thickness(img,thickness_change,fg_shade,bg_shade,blur_size,noise_sigma)
                if len(new_img.shape)==2:
                    new_img = new_img[...,None]
                new_img = new_img*2 -1.0

            if len(img.shape)==2:
                img = img[...,None]

            img = img.astype(np.float32)
            if self.remove_bg:
                img = 1.0 - img / 256.0
                #kernel = torch.FloatTensor(7,7).fill_(1/49)
                #blurred_mask = F.conv2d(fg_mask,kernel,padding=3)
                blurred_mask = cv2.blur(fg_mask,(7,7))
                img *= blurred_mask[...,None]
                img = 2*img -1
            else:
                img = 1.0 - img / 128.0



            if len(gt) == 0:
                return None
            gt_label = string_utils.str2label_single(gt, self.char_to_idx)

            if self.styles:
                style_i = self.npr.choice(len(self.styles[author][id]))
                style = self.styles[author][id][style_i]
            else:
                style=None
            name = '{}_{}'.format(author,line)
            if self.identity_spaced:
                spaced_label = gt_label[:,None].astype(np.long)
            else:
                spaced_label = None if self.spaced_by_name is None else self.spaced_by_name[name]
                if spaced_label is not None:
                    assert(spaced_label.shape[1]==1)
            toAppend= {
                "image": img,
                "gt": gt,
                "style": style,
                "gt_label": gt_label,
                "spaced_label": spaced_label,
                "name": name,
                "center": self.center,
                "author": author,
                "author_idx": self.author_list.index(author)
                
                }
            if self.fg_masks_dir is not None:
                toAppend['fg_mask'] = fg_mask
            if self.include_stroke_aug:
                toAppend['changed_image'] = new_img
            batch.append(toAppend)
            
        #batch = [b for b in batch if b is not None]
        #These all should be the same size or error
        assert len(set([b['image'].shape[0] for b in batch])) == 1
        assert len(set([b['image'].shape[2] for b in batch])) == 1

        dim0 = batch[0]['image'].shape[0]
        dim1 = max([b['image'].shape[1] for b in batch])
        dim2 = batch[0]['image'].shape[2]

        all_labels = []
        label_lengths = []
        if self.spaced_by_name is not None or self.identity_spaced:
            spaced_labels = []
        else:
            spaced_labels = None
        max_spaced_len=0

        input_batch = np.full((len(batch), dim0, dim1, dim2), PADDING_CONSTANT).astype(np.float32)
        if self.fg_masks_dir is not None:
            fg_masks = np.full((len(batch), dim0, dim1, 1), 0).astype(np.float32)
        if self.include_stroke_aug:
            changed_batch = np.full((len(batch), dim0, dim1, dim2), PADDING_CONSTANT).astype(np.float32)
        for i in range(len(batch)):
            b_img = batch[i]['image']
            toPad = (dim1-b_img.shape[1])
            if 'center' in batch[0] and batch[0]['center']:
                toPad //=2
            else:
                toPad = 0
            input_batch[i,:,toPad:toPad+b_img.shape[1],:] = b_img
            if self.fg_masks_dir is not None:
                fg_masks[i,:,toPad:toPad+b_img.shape[1],0] = batch[i]['fg_mask']
            if self.include_stroke_aug:
                changed_batch[i,:,toPad:toPad+b_img.shape[1],:] = batch[i]['changed_image']

            l = batch[i]['gt_label']
            all_labels.append(l)
            label_lengths.append(len(l))

            if spaced_labels is not None:
                sl = batch[i]['spaced_label']
                spaced_labels.append(sl)
                max_spaced_len = max(max_spaced_len,sl.shape[0])

        #all_labels = np.concatenate(all_labels)
        label_lengths = torch.IntTensor(label_lengths)
        max_len = label_lengths.max()
        all_labels = [np.pad(l,((0,max_len-l.shape[0]),),'constant') for l in all_labels]
        all_labels = np.stack(all_labels,axis=1)
        if self.spaced_by_name is not None or self.identity_spaced:
            spaced_labels = [np.pad(l,((0,max_spaced_len-l.shape[0]),(0,0)),'constant') for l in spaced_labels]
            ddd = spaced_labels
            spaced_labels = np.concatenate(spaced_labels,axis=1)
            spaced_labels = torch.from_numpy(spaced_labels)
            assert(spaced_labels.size(1) == len(batch))


        images = input_batch.transpose([0,3,1,2])
        images = torch.from_numpy(images)
        labels = torch.from_numpy(all_labels.astype(np.int32))
        #label_lengths = torch.from_numpy(label_lengths.astype(np.int32))
        if self.fg_masks_dir is not None:
            fg_masks = fg_masks.transpose([0,3,1,2])
            fg_masks = torch.from_numpy(fg_masks)
        
        if batch[0]['style'] is not None:
            styles = np.stack([b['style'] for b in batch], axis=0)
            styles = torch.from_numpy(styles).float()
        else:
            styles=None
        mask, top_and_bottom, center_line = makeMask(images,self.mask_post, self.mask_random)
        ##DEBUG
        #for i in range(5):
        #    mask2, top_and_bottom2 = makeMask(images,self.mask_post, self.mask_random)
        #    #extra_masks.append(mask2)
        #    mask2 = ((mask2[0,0]+1)/2).numpy().astype(np.uint8)*255
        #    cv2.imshow('mask{}'.format(i),mask2)
        #mask = ((mask[0,0]+1)/2).numpy().astype(np.uint8)*255
        #cv2.imshow('mask'.format(i),mask)
        #cv2.waitKey()
        toRet= {
            "image": images,
            "mask": mask,
            "top_and_bottom": top_and_bottom,
            "center_line": center_line,
            "label": labels,
            "style": styles,
            "label_lengths": label_lengths,
            "gt": [b['gt'] for b in batch],
            "spaced_label": spaced_labels,
            "name": [b['name'] for b in batch],
            "author": [b['author'] for b in batch],
            "author_idx": [b['author_idx'] for b in batch],
        }
        if self.fg_masks_dir is not None:
            toRet['fg_mask'] = fg_masks
        if self.include_stroke_aug:
            changed_images = changed_batch.transpose([0,3,1,2])
            changed_images = torch.from_numpy(changed_images)
            toRet['changed_image']=changed_images
        return toRet
    def __getitem__(self, idx):

        inst = self.lineIndex[idx]
        author = inst[0]
        lines = inst[1]
        batch = []
        for line in lines:
            if line >= len(self.authors[author]):
                line = (line + 37) % len(self.authors[author])
            img_path, gt, pad_above, pad_below = self.authors[author][line]
            img = cv2.imread(img_path, 0)  #read as grayscale
            if img is None:
                return None

            if pad_above < 0:
                img = img[-pad_above:, :]
                pad_above = 0
            if pad_below < 0:
                img = img[:pad_below, :]
                pad_below = 0
            #if pad_above>0 or pad_below>0:
            img = img = np.pad(img, ((pad_above, pad_below), (10, 10)),
                               'constant',
                               constant_values=255)
            #we also pad a bit on the sides
            #print('{}, {} {}'.format(img_path,pad_above,pad_below))

            if img.shape[0] != self.img_height:
                if img.shape[0] < self.img_height and not self.warning:
                    self.warning = True
                    print("WARNING: upsampling image to fit size")
                percent = float(self.img_height) / img.shape[0]
                if img.shape[1] * percent > self.max_width:
                    percent = self.max_width / img.shape[1]
                img = cv2.resize(img, (0, 0),
                                 fx=percent,
                                 fy=percent,
                                 interpolation=cv2.INTER_CUBIC)
                if img.shape[0] < self.img_height:
                    diff = self.img_height - img.shape[0]
                    img = np.pad(img,
                                 ((diff // 2, diff // 2 + diff % 2), (0, 0)),
                                 'constant',
                                 constant_values=255)

            if len(img.shape) == 2:
                img = img[..., None]
            if self.fg_masks_dir is not None:
                fg_path = os.path.join(self.fg_masks_dir,
                                       '{}_{}.png'.format(author, line))
                fg_mask = cv2.imread(fg_path, 0)
                fg_mask = fg_mask / 255
                if fg_mask.shape != img[:, :, 0].shape:
                    print(
                        'Error, fg_mask ({}, {}) not the same size as image ({})'
                        .format(fg_path, fg_mask.shape, img[:, :, 0].shape))
                    th, fg_mask = cv2.threshold(
                        img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
                    fg_mask = 255 - fg_mask
                    ele = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
                    fg_mask = cv2.dilate(fg_mask, ele)
                    fg_mask = fg_mask / 255

            if self.augmentation is not None:
                #img = augmentation.apply_random_color_rotation(img)
                img = augmentation.apply_tensmeyer_brightness(img)
                img = grid_distortion.warp_image(img)
                if len(img.shape) == 2:
                    img = img[..., None]

            img = img.astype(np.float32)
            img = 1.0 - img / 128.0

            if len(gt) == 0:
                return None
            gt_label = string_utils.str2label_single(gt, self.char_to_idx)

            if self.styles:
                style_i = self.npr.choice(len(self.styles[author][id]))
                style = self.styles[author][id][style_i]
            else:
                style = None
            name = img_path[img_path.rfind('/') + 1:img_path.rfind('.')]
            spaced_label = None if self.spaced_by_name is None else self.spaced_by_name[
                img_path]
            if spaced_label is not None:
                assert (spaced_label.shape[1] == 1)
            toAppend = {
                "image": img,
                "gt": gt,
                "style": style,
                "gt_label": gt_label,
                "spaced_label": spaced_label,
                "name": name,
                "center": self.center,
                "author": author
            }
            if self.fg_masks_dir is not None:
                toAppend['fg_mask'] = fg_mask
            batch.append(toAppend)
        #batch = [b for b in batch if b is not None]
        #These all should be the same size or error
        assert len(set([b['image'].shape[0] for b in batch])) == 1
        assert len(set([b['image'].shape[2] for b in batch])) == 1

        dim0 = batch[0]['image'].shape[0]
        dim1 = max([b['image'].shape[1] for b in batch])
        dim2 = batch[0]['image'].shape[2]

        all_labels = []
        label_lengths = []
        if self.spaced_by_name is not None:
            spaced_labels = []
        else:
            spaced_labels = None
        max_spaced_len = 0

        input_batch = np.full((len(batch), dim0, dim1, dim2),
                              PADDING_CONSTANT).astype(np.float32)
        if self.fg_masks_dir is not None:
            fg_masks = np.full((len(batch), dim0, dim1, 1),
                               0).astype(np.float32)
        for i in range(len(batch)):
            b_img = batch[i]['image']
            toPad = (dim1 - b_img.shape[1])
            if 'center' in batch[0] and batch[0]['center']:
                toPad //= 2
            else:
                toPad = 0
            input_batch[i, :, toPad:toPad + b_img.shape[1], :] = b_img
            if self.fg_masks_dir is not None:
                fg_masks[i, :, toPad:toPad + b_img.shape[1],
                         0] = batch[i]['fg_mask']

            l = batch[i]['gt_label']
            all_labels.append(l)
            label_lengths.append(len(l))

            if spaced_labels is not None:
                sl = batch[i]['spaced_label']
                spaced_labels.append(sl)
                max_spaced_len = max(max_spaced_len, sl.shape[0])

        #all_labels = np.concatenate(all_labels)
        label_lengths = torch.IntTensor(label_lengths)
        max_len = label_lengths.max()
        all_labels = [
            np.pad(l, ((0, max_len - l.shape[0]), ), 'constant')
            for l in all_labels
        ]
        all_labels = np.stack(all_labels, axis=1)
        if self.spaced_by_name is not None:
            spaced_labels = [
                np.pad(l, ((0, max_spaced_len - l.shape[0]), (0, 0)),
                       'constant') for l in spaced_labels
            ]
            ddd = spaced_labels
            spaced_labels = np.concatenate(spaced_labels, axis=1)
            spaced_labels = torch.from_numpy(spaced_labels)
            assert (spaced_labels.size(1) == len(batch))

        images = input_batch.transpose([0, 3, 1, 2])
        images = torch.from_numpy(images)
        labels = torch.from_numpy(all_labels.astype(np.int32))
        #label_lengths = torch.from_numpy(label_lengths.astype(np.int32))
        if self.fg_masks_dir is not None:
            fg_masks = fg_masks.transpose([0, 3, 1, 2])
            fg_masks = torch.from_numpy(fg_masks)

        if batch[0]['style'] is not None:
            styles = np.stack([b['style'] for b in batch], axis=0)
            styles = torch.from_numpy(styles).float()
        else:
            styles = None
        mask, top_and_bottom, center_line = makeMask(images, self.mask_post,
                                                     self.mask_random)
        ##DEBUG
        #for i in range(5):
        #    mask2, top_and_bottom2 = makeMask(images,self.mask_post, self.mask_random)
        #    #extra_masks.append(mask2)
        #    mask2 = ((mask2[0,0]+1)/2).numpy().astype(np.uint8)*255
        #    cv2.imshow('mask{}'.format(i),mask2)
        #mask = ((mask[0,0]+1)/2).numpy().astype(np.uint8)*255
        #cv2.imshow('mask'.format(i),mask)
        #cv2.waitKey()
        toRet = {
            "image": images,
            "mask": mask,
            "top_and_bottom": top_and_bottom,
            "center_line": center_line,
            "label": labels,
            "style": styles,
            "label_lengths": label_lengths,
            "gt": [b['gt'] for b in batch],
            "spaced_label": spaced_labels,
            "name": [b['name'] for b in batch],
            "author": [b['author'] for b in batch],
        }
        if self.fg_masks_dir is not None:
            toRet['fg_mask'] = fg_masks
        return toRet