File size: 4,356 Bytes
a3d6c18 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
"""
Source url: https://github.com/OPHoperHPO/image-background-remove-tool
Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
License: Apache License 2.0
"""
import pathlib
from typing import Union, Any, Tuple
import PIL.Image
import numpy as np
import torch
ALLOWED_SUFFIXES = [".jpg", ".jpeg", ".bmp", ".png", ".webp"]
def to_tensor(x: Any) -> torch.Tensor:
"""
Returns a PIL.Image.Image as torch tensor without swap tensor dims.
Args:
x: PIL.Image.Image instance
Returns:
torch.Tensor instance
"""
return torch.tensor(np.array(x, copy=True))
def load_image(file: Union[str, pathlib.Path, PIL.Image.Image]) -> PIL.Image.Image:
"""Returns a PIL.Image.Image class by string path or pathlib path or PIL.Image.Image instance
Args:
file: File path or PIL.Image.Image instance
Returns:
PIL.Image.Image instance
Raises:
ValueError: If file not exists or file is directory or file isn't an image or file is not correct PIL Image
"""
if isinstance(file, str) and is_image_valid(pathlib.Path(file)):
return PIL.Image.open(file)
elif isinstance(file, PIL.Image.Image):
return file
elif isinstance(file, pathlib.Path) and is_image_valid(file):
return PIL.Image.open(str(file))
else:
raise ValueError("Unknown input file type")
def convert_image(image: PIL.Image.Image, mode="RGB") -> PIL.Image.Image:
"""Performs image conversion to correct color mode
Args:
image: PIL.Image.Image instance
mode: Colort Mode to convert
Returns:
PIL.Image.Image instance
Raises:
ValueError: If image hasn't convertable color mode, or it is too small
"""
if is_image_valid(image):
return image.convert(mode)
def is_image_valid(image: Union[pathlib.Path, PIL.Image.Image]) -> bool:
"""This function performs image validation.
Args:
image: Path to the image or PIL.Image.Image instance being checked.
Returns:
True if image is valid
Raises:
ValueError: If file not a valid image path or image hasn't convertable color mode, or it is too small
"""
if isinstance(image, pathlib.Path):
if not image.exists():
raise ValueError("File is not exists")
elif image.is_dir():
raise ValueError("File is a directory")
elif image.suffix.lower() not in ALLOWED_SUFFIXES:
raise ValueError(
f"Unsupported image format. Supported file formats: {', '.join(ALLOWED_SUFFIXES)}"
)
elif isinstance(image, PIL.Image.Image):
if not (image.size[0] > 32 and image.size[1] > 32):
raise ValueError("Image should be bigger then (32x32) pixels.")
elif image.mode not in ["RGB", "RGBA", "L"]:
raise ValueError("Wrong image color mode.")
else:
raise ValueError("Unknown input file type")
return True
def transparency_paste(
bg_img: PIL.Image.Image, fg_img: PIL.Image.Image, box=(0, 0)
) -> PIL.Image.Image:
"""
Inserts an image into another image while maintaining transparency.
Args:
bg_img: background image
fg_img: foreground image
box: place to paste
Returns:
Background image with pasted foreground image at point or in the specified box
"""
fg_img_trans = PIL.Image.new("RGBA", bg_img.size)
fg_img_trans.paste(fg_img, box, mask=fg_img)
new_img = PIL.Image.alpha_composite(bg_img, fg_img_trans)
return new_img
def add_margin(
pil_img: PIL.Image.Image,
top: int,
right: int,
bottom: int,
left: int,
color: Tuple[int, int, int, int],
) -> PIL.Image.Image:
"""
Adds margin to the image.
Args:
pil_img: Image that needed to add margin.
top: pixels count at top side
right: pixels count at right side
bottom: pixels count at bottom side
left: pixels count at left side
color: color of margin
Returns:
Image with margin.
"""
width, height = pil_img.size
new_width = width + right + left
new_height = height + top + bottom
# noinspection PyTypeChecker
result = PIL.Image.new(pil_img.mode, (new_width, new_height), color)
result.paste(pil_img, (left, top))
return result
|