dataset.py 文件源码

def filter_image(path, min_size=300, max_size=None, min_aspect=1 / 4, max_aspect=4):
    import numpy
    from PIL import Image
    from skimage.color import rgb2hsv

    try:
        assert os.path.getsize(path) <= 100 * 1000 * 1000, 'file size error'

        image = Image.open(path).convert('RGB')

        w, h = image.size
        if min_size is not None:
            assert w >= min_size and h >= min_size, 'min size error'
        if max_size is not None:
            assert w <= max_size or h <= max_size, 'max size error'

        aspect = w / h
        assert min_aspect is None or aspect >= min_aspect, 'min aspect error'
        assert max_aspect is None or aspect <= max_aspect, 'max aspect error'

        # check data process
        data = data_process(path, test=True)

        # remove low saturation
        rgb = numpy.array(image, numpy.float32) / 255
        hsv = rgb2hsv(rgb)
        sigma = numpy.var(hsv[:, :, 1])
        assert sigma > 0.02, 'saturation error'

    except Exception as e:
        print(path, e)
        return False

    return True