"""
Helpers for cropping image depending on resolution.
TODO: replace recursive code with something more adaptive, right now we only do
- no crops / 2x2 / 4x4 / 16x16 / etc.
and the stitching code is complicated, but would be nice to have e.g. 3x3.
"""
import math
import os
import itertools
import torch
from blueprints.multiscale_blueprint import MultiscaleLoss
import functools
import operator
def prod(it):
return functools.reduce(operator.mul, it, 1)
# Images with H * W > prod(_NEEDS_CROP_DIM) will be split into crops
# We set this empirically such that crops fit into our TITAN X (Pascal) with 12GB VRAM.
# You can set this from the console using AC_NEEDS_CROP_DIM, e.g.,
#
# AC_NEEDS_CROP_DIM=2000,2000 python test.py ...
#
# But expect OOM errors for big values.
_NEEDS_CROP_DIM_DEFAULT = '2000,1500'
_NEEDS_CROP_DIM = os.environ.get('AC_NEEDS_CROP_DIM', _NEEDS_CROP_DIM_DEFAULT)
if _NEEDS_CROP_DIM != _NEEDS_CROP_DIM_DEFAULT:
print('*** AC_NEEDS_CROP_DIM =', _NEEDS_CROP_DIM)
_NEEDS_CROP_DIM = prod(map(int, _NEEDS_CROP_DIM.split(',')))
print('*** AC_NEEDS_CROP_DIM =', _NEEDS_CROP_DIM)
def _assert_valid_image(i):
if len(i.shape) != 4 or i.shape[1] != 3:
raise ValueError(f'Expected BCHW image, got {i.shape}')
def needs_crop(img, needs_crop_dim=_NEEDS_CROP_DIM):
_assert_valid_image(img)
H, W = img.shape[-2:]
return H * W > needs_crop_dim
def _crop16(im):
for im_cropped in _crop4(im):
yield from _crop4(im_cropped)
def iter_crops(img, needs_crop_dim=_NEEDS_CROP_DIM):
_assert_valid_image(img)
if not needs_crop(img, needs_crop_dim):
yield img
return
for img_crop in _crop4(img):
yield from iter_crops(img_crop, needs_crop_dim)
def _crop4(img):
_assert_valid_image(img)
H, W = img.shape[-2:]
imgs = [img[..., :H//2, :W//2], # Top left
img[..., :H//2, W//2:], # Top right
img[..., H//2:, :W//2], # Bottom left
img[..., H//2:, W//2:]] # Bottom right
# Validate that we got all pixels
assert sum(prod(img.shape[-2:]) for img in imgs) == \
prod(img.shape[-2:])
return imgs
def _get_crop_idx_mapping(side):
"""Helper method to get the order of crops.
:param side: how many crops live on each side.
Example. Say you have an image that gets devided into 16 crops, i.e., the image gets cut into 16 parts:
[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15]],
However, due to our recursive cropping code, this results in crops that are ordered like this:
index of crop:
0 1 2 3 4 ...
corresponds to part in image:
0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15
This method returns the inverse, going from the index of the crop back to the index in the image.
"""
a = torch.arange(side * side).reshape(1, 1, side, side)
a = torch.cat((a, a, a), dim=1)
# Create mapping
# Index of crop in original image -> index of crop in the order it was extracted,
# E.g. 2 -> 4 means it's the 2nd crop, but in the image, it's at position 4 (see above).
crops = {i: crop[0, 0, ...].flatten().item()
for i, crop in enumerate(iter_crops(a, 1))}
return crops
def stitch(parts):
side = int(math.sqrt(len(parts)))
if side * side != len(parts):
raise ValueError(f'Invalid number of parts {len(parts)}')
rows = []
# Sort by original position in image
crops_idx_mapping = _get_crop_idx_mapping(side)
parts_sorted = (
part for _, part in sorted(
enumerate(parts), key=lambda ip: crops_idx_mapping[ip[0]]))
parts_itr = iter(parts_sorted) # Turn into iterator so we can easily grab elements
for _ in range(side):
parts_row = itertools.islice(parts_itr, side) # Get `side` number of parts
row = torch.cat(list(parts_row), dim=3) # cat on W dimension
rows.append(row)
assert next(parts_itr, None) is None, f'Iterator should be empty, got {len(rows)} rows'
img = torch.cat(rows, dim=2) # cat on H dimension
# Validate.
B, C, H_part, W_part = parts[0].shape
expected_shape = (B, C, H_part * side, W_part * side)
assert img.shape == expected_shape, f'{img.shape} != {expected_shape}'
return img
class CropLossCombinator(object):
"""Used to combine the bpsp of different crops into one. Supports crops of varying dimensions."""
def __init__(self):
self._num_bits_total = 0.
self._num_subpixels_total = 0
def add(self, bpsp, num_subpixels_crop):
bits = bpsp * num_subpixels_crop
self._num_bits_total += bits
self._num_subpixels_total += num_subpixels_crop
def get_bpsp(self):
assert self._num_subpixels_total > 0
return self._num_bits_total / self._num_subpixels_total
def test_auto_crop():
import torch
import pytorch_ext as pe
for H, W, num_crops_expected in [(10000, 6000, 64),
(4928, 3264, 16),
(2048, 2048, 4),
(1024, 1024, 1),
]:
img = (torch.rand(1, 3, H, W) * 255).round().long()
print(img.shape)
if num_crops_expected > 1:
assert needs_crop(img)
crops = list(iter_crops(img, 2048 * 1024))
assert len(crops) == num_crops_expected
pe.assert_equal(stitch(crops), img)
else:
pe.assert_equal(next(iter_crops(img, 2048 * 1024)), img)
