Spaces:
Sleeping
Sleeping
File size: 6,031 Bytes
59e40e1 |
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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
"""
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
import warnings
from typing import List, Union
import PIL.Image
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from carvekit.ml.arch.tracerb7.tracer import TracerDecoder
from carvekit.ml.arch.tracerb7.efficientnet import EfficientEncoderB7
from carvekit.ml.files.models_loc import tracer_b7_pretrained, tracer_hair_pretrained
from carvekit.utils.models_utils import get_precision_autocast, cast_network
from carvekit.utils.image_utils import load_image, convert_image
from carvekit.utils.pool_utils import thread_pool_processing, batch_generator
__all__ = ["TracerUniversalB7"]
class TracerUniversalB7(TracerDecoder):
"""TRACER B7 model interface"""
def __init__(
self,
device="cpu",
input_image_size: Union[List[int], int] = 640,
batch_size: int = 4,
load_pretrained: bool = True,
fp16: bool = False,
model_path: Union[str, pathlib.Path] = None,
):
"""
Initialize the U2NET model
Args:
layers_cfg: neural network layers configuration
device: processing device
input_image_size: input image size
batch_size: the number of images that the neural network processes in one run
load_pretrained: loading pretrained model
fp16: use fp16 precision
"""
if model_path is None:
model_path = tracer_b7_pretrained()
super(TracerUniversalB7, self).__init__(
encoder=EfficientEncoderB7(),
rfb_channel=[32, 64, 128],
features_channels=[48, 80, 224, 640],
)
self.fp16 = fp16
self.device = device
self.batch_size = batch_size
if isinstance(input_image_size, list):
self.input_image_size = input_image_size[:2]
else:
self.input_image_size = (input_image_size, input_image_size)
self.transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Resize(self.input_image_size),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
self.to(device)
if load_pretrained:
# TODO remove edge detector from weights. It doesn't work well with this model!
self.load_state_dict(
torch.load(model_path, map_location=self.device), strict=False
)
self.eval()
def data_preprocessing(self, data: PIL.Image.Image) -> torch.FloatTensor:
"""
Transform input image to suitable data format for neural network
Args:
data: input image
Returns:
input for neural network
"""
return torch.unsqueeze(self.transform(data), 0).type(torch.FloatTensor)
@staticmethod
def data_postprocessing(
data: torch.tensor, original_image: PIL.Image.Image
) -> PIL.Image.Image:
"""
Transforms output data from neural network to suitable data
format for using with other components of this framework.
Args:
data: output data from neural network
original_image: input image which was used for predicted data
Returns:
Segmentation mask as PIL Image instance
"""
output = (data.type(torch.FloatTensor).detach().cpu().numpy() * 255.0).astype(
np.uint8
)
output = output.squeeze(0)
mask = Image.fromarray(output).convert("L")
mask = mask.resize(original_image.size, resample=Image.BILINEAR)
return mask
def __call__(
self, images: List[Union[str, pathlib.Path, PIL.Image.Image]]
) -> List[PIL.Image.Image]:
"""
Passes input images though neural network and returns segmentation masks as PIL.Image.Image instances
Args:
images: input images
Returns:
segmentation masks as for input images, as PIL.Image.Image instances
"""
collect_masks = []
autocast, dtype = get_precision_autocast(device=self.device, fp16=self.fp16)
with autocast:
cast_network(self, dtype)
for image_batch in batch_generator(images, self.batch_size):
images = thread_pool_processing(
lambda x: convert_image(load_image(x)), image_batch
)
batches = torch.vstack(
thread_pool_processing(self.data_preprocessing, images)
)
with torch.no_grad():
batches = batches.to(self.device)
masks = super(TracerDecoder, self).__call__(batches)
masks_cpu = masks.cpu()
del batches, masks
masks = thread_pool_processing(
lambda x: self.data_postprocessing(masks_cpu[x], images[x]),
range(len(images)),
)
collect_masks += masks
return collect_masks
class TracerHair(TracerUniversalB7):
"""TRACER HAIR model interface"""
def __init__(
self,
device="cpu",
input_image_size: Union[List[int], int] = 640,
batch_size: int = 4,
load_pretrained: bool = True,
fp16: bool = False,
model_path: Union[str, pathlib.Path] = None,
):
if model_path is None:
model_path = tracer_hair_pretrained()
warnings.warn("TracerHair has not public model yet. Don't use it!", UserWarning)
super(TracerHair, self).__init__(
device=device,
input_image_size=input_image_size,
batch_size=batch_size,
load_pretrained=load_pretrained,
fp16=fp16,
model_path=model_path,
)
|