Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
@raises(KrakenInputException)
def test_not_binarize_empty(self):
"""
Test that mode '1' images aren't binarized again.
"""
with Image.new('1', (1000,1000)) as im:
nlbin(im)
@raises(KrakenInputException)
def test_rpred_outbounds(self):
"""
Tests correct handling of invalid line coordinates.
"""
nn = load_any(os.path.join(resources, 'toy.clstm'))
pred = rpred(nn, self.im, {'boxes': [[-1, -1, 10000, 10000]], 'text_direction': 'horizontal'}, True)
next(pred)
Args:
filename (str): path to an ALTO file.
Returns:
A dict {'image': impath, lines: [{'boundary': [[x0, y0], ...], 'baseline':
[[x0, y0], ...]}, {...], 'text': 'apdjfqpf'}
"""
with open(filename, 'rb') as fp:
base_dir = dirname(filename)
try:
doc = etree.parse(fp)
except etree.XMLSyntaxError as e:
raise KrakenInputException('Parsing {} failed: {}'.format(filename, e))
image = doc.find('.//{*}fileName')
if image is None or not image.text:
raise KrakenInputException('No valid filename found in ALTO file')
lines = doc.findall('.//{*}TextLine')
data = {'image': os.path.join(base_dir, image.text), 'lines': [], 'type': 'baselines'}
for line in lines:
if line.get('BASELINE') is None:
raise KrakenInputException('ALTO file {} contains no baseline information'.format(filename))
pol = line.find('./{*}Shape/{*}Polygon')
boundary = None
if pol is not None:
points = [int(x) for x in pol.get('POINTS').split(' ')]
boundary = list(zip(points[::2], points[1::2]))
points = [int(x) for x in line.get('BASELINE').split(' ')]
baseline = list(zip(points[::2], points[1::2]))
text = ''
for el in line.xpath(".//*[local-name() = 'String'] | .//*[local-name() = 'SP']"):
text += el.get('CONTENT') if el.get('CONTENT') else ' '
data['lines'].append({'baseline': baseline, 'boundary': boundary, 'text': text})
def add(self, image: str) -> None:
"""
Adds a line-image-text pair to the dataset.
Args:
image (str): Input image path
"""
with open(self.split(image), 'r', encoding='utf-8') as fp:
gt = fp.read().strip('\n\r')
for func in self.text_transforms:
gt = func(gt)
if not gt:
raise KrakenInputException('Text line is empty ({})'.format(fp.name))
if self.preload:
im = Image.open(image)
try:
im = self.transforms(im)
except ValueError as e:
raise KrakenInputException('Image transforms failed on {}'.format(image))
self._images.append(im)
else:
self._images.append(image)
self._gt.append(gt)
self.alphabet.update(gt)
def forward(self, line: torch.Tensor) -> np.array:
"""
Performs a forward pass on a torch tensor of a line with shape (C, H, W)
and returns a numpy array (W, C).
"""
# make CHW -> 1CHW
line = line.to(self.device)
line = line.unsqueeze(0)
o = self.nn.nn(line)
if o.size(2) != 1:
raise KrakenInputException('Expected dimension 3 to be 1, actual {}'.format(o.size()))
self.outputs = o.detach().squeeze().cpu().numpy()
return self.outputs
elif segmentation:
logger.debug(u'Adding segmentations.')
self.text_direction = segmentation['text_direction']
for bbox in segmentation['boxes']:
page['lines'].append({'index': self.line_idx,
'left': 100*int(bbox[0]) / im.size[0],
'top': 100*int(bbox[1]) / im.size[1],
'width': 100*(bbox[2] - bbox[0])/im.size[0],
'height': 100*(int(bbox[3]) - int(bbox[1]))/im.size[1],
'bbox': '{}, {}, {}, {}'.format(int(bbox[0]),
int(bbox[1]),
int(bbox[2]),
int(bbox[3]))})
self.line_idx += 1
else:
raise KrakenInputException('Neither segmentations nor records given')
self.pages.append(page)
respective polygonal boundaries. The last and first point of each
boundary polygon is connected.
Raises:
KrakenInputException if the input image is not binarized or the text
direction is invalid.
"""
im_str = get_im_str(im)
logger.info('Segmenting {}'.format(im_str))
model = vgsl.TorchVGSLModel.load_model(model)
model.eval()
if mask:
if mask.mode != '1' and not is_bitonal(mask):
logger.error('Mask is not bitonal')
raise KrakenInputException('Mask is not bitonal')
mask = mask.convert('1')
if mask.size != im.size:
logger.error('Mask size {} doesn\'t match image size {}'.format(mask.size, im.size))
raise KrakenInputException('Mask size {} doesn\'t match image size {}'.format(mask.size, im.size))
logger.info('Masking enabled in segmenter.')
mask = pil2array(mask)
batch, channels, height, width = model.input
transforms = dataset.generate_input_transforms(batch, height, width, channels, 0, valid_norm=False)
res_tf = tf.Compose(transforms.transforms[:2])
scal_im = res_tf(im).convert('L')
with torch.no_grad():
logger.debug('Running network forward pass')
o = model.nn(transforms(im).unsqueeze(0))
logger.debug('Upsampling network output')
bounds (list): A list of tuples (x1, y1, x2, y2)
Yields:
(PIL.Image) the extracted subimage
"""
if bounds['text_direction'].startswith('vertical'):
angle = 90
else:
angle = 0
for box in bounds['boxes']:
if isinstance(box, tuple):
box = list(box)
if (box < [0, 0, 0, 0] or box[::2] > [im.size[0], im.size[0]] or
box[1::2] > [im.size[1], im.size[1]]):
logger.error('bbox {} is outside of image bounds {}'.format(box, im.size))
raise KrakenInputException('Line outside of image bounds')
yield im.crop(box).rotate(angle, expand=True), box