init app

.gitignore

@@ -0,0 +1,2 @@
+.idea/
+__pycache__/

app.py

@@ -0,0 +1,128 @@
+# https://github.com/MarcoForte/FBA_Matting
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+
+from networks.models import build_model
+from networks.transforms import trimap_transform, normalise_image
+
+REPO_ID = "leonelhs/FBA-Matting"
+
+weights = hf_hub_download(repo_id=REPO_ID, filename="FBA.pth")
+model = build_model(weights)
+model.eval().cpu()
+
+
+def np_to_torch(x, permute=True):
+    if permute:
+        return torch.from_numpy(x).permute(2, 0, 1)[None, :, :, :].float().cpu()
+    else:
+        return torch.from_numpy(x)[None, :, :, :].float().cpu()
+
+
+def scale_input(x: np.ndarray, scale: float, scale_type) -> np.ndarray:
+    ''' Scales inputs to multiple of 8. '''
+    h, w = x.shape[:2]
+    h1 = int(np.ceil(scale * h / 8) * 8)
+    w1 = int(np.ceil(scale * w / 8) * 8)
+    x_scale = cv2.resize(x, (w1, h1), interpolation=scale_type)
+    return x_scale
+
+
+def inference(image_np: np.ndarray, trimap_np: np.ndarray) -> [np.ndarray]:
+    ''' Predict alpha, foreground and background.
+        Parameters:
+        image_np -- the image in rgb format between 0 and 1. Dimensions: (h, w, 3)
+        trimap_np -- two channel trimap, first background then foreground. Dimensions: (h, w, 2)
+        Returns:
+        fg: foreground image in rgb format between 0 and 1. Dimensions: (h, w, 3)
+        bg: background image in rgb format between 0 and 1. Dimensions: (h, w, 3)
+        alpha: alpha matte image between 0 and 1. Dimensions: (h, w)
+    '''
+    h, w = trimap_np.shape[:2]
+    image_scale_np = scale_input(image_np, 1.0, cv2.INTER_LANCZOS4)
+    trimap_scale_np = scale_input(trimap_np, 1.0, cv2.INTER_LANCZOS4)
+
+    with torch.no_grad():
+        image_torch = np_to_torch(image_scale_np)
+        trimap_torch = np_to_torch(trimap_scale_np)
+
+        trimap_transformed_torch = np_to_torch(
+            trimap_transform(trimap_scale_np), permute=False)
+        image_transformed_torch = normalise_image(
+            image_torch.clone())
+
+        output = model(
+            image_torch,
+            trimap_torch,
+            image_transformed_torch,
+            trimap_transformed_torch)
+        output = cv2.resize(
+            output[0].cpu().numpy().transpose(
+                (1, 2, 0)), (w, h), cv2.INTER_LANCZOS4)
+
+    alpha = output[:, :, 0]
+    fg = output[:, :, 1:4]
+    bg = output[:, :, 4:7]
+
+    alpha[trimap_np[:, :, 0] == 1] = 0
+    alpha[trimap_np[:, :, 1] == 1] = 1
+    fg[alpha == 1] = image_np[alpha == 1]
+    bg[alpha == 0] = image_np[alpha == 0]
+
+    return fg, bg, alpha
+
+
+def read_image(name):
+    return (cv2.imread(name) / 255.0)[:, :, ::-1]
+
+
+def read_trimap(name):
+    trimap_im = cv2.imread(name, 0) / 255.0
+    h, w = trimap_im.shape
+    trimap_np = np.zeros((h, w, 2))
+    trimap_np[trimap_im == 1, 1] = 1
+    trimap_np[trimap_im == 0, 0] = 1
+    return trimap_np
+
+
+def predict(image, trimap):
+    image_np = read_image(image)
+    trimap_np = read_trimap(trimap)
+    return inference(image_np, trimap_np)
+
+
+footer = r"""
+<center>
+<b>
+Demo for <a href='https://github.com/MarcoForte/FBA_Matting'>FBA Matting</a>
+</b>
+</center>
+"""
+
+with gr.Blocks(title="FBA Matting") as app:
+    gr.HTML("<center><h1>FBA Matting</h1></center>")
+    gr.HTML("<center><h3>Foreground, Background, Alpha Matting Generator.</h3></center>")
+    with gr.Row().style(equal_height=False):
+        with gr.Column():
+            input_img = gr.Image(type="filepath", label="Input image")
+            input_trimap = gr.Image(type="filepath", label="Trimap image")
+            run_btn = gr.Button(variant="primary")
+        with gr.Column():
+            fg = gr.Image(type="numpy", label="Foreground")
+            bg = gr.Image(type="numpy", label="Background")
+            alpha = gr.Image(type="numpy", label="Alpha")
+
+    run_btn.click(predict, [input_img, input_trimap], [fg, bg, alpha])
+
+    with gr.Row():
+        examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
+        examples = gr.Dataset(components=[input_img], samples=examples_data)
+        examples.click(lambda x: x[0], [examples], [input_img])
+
+    with gr.Row():
+        gr.HTML(footer)
+
+app.launch(share=False, debug=True, enable_queue=True, show_error=True)

networks/layers_WS.py

@@ -0,0 +1,32 @@
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class Conv2d(nn.Conv2d):
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True, eps=1e-5):
+        super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
+                                     padding, dilation, groups, bias)
+        self.out_channels = out_channels
+        self.eps = eps
+
+    def normalize_weight(self):
+        weight = F.batch_norm(
+            self.weight.view(1, self.out_channels, -1), None, None,
+            training=True, momentum=0., eps=self.eps).reshape_as(self.weight)
+        self.weight.data = weight
+
+    def forward(self, x):
+        if self.training:
+            self.normalize_weight()
+        return F.conv2d(x, self.weight, self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+    def train(self, mode: bool = True):
+        super().train(mode=mode)
+        self.normalize_weight()
+
+
+def norm(dim):
+    return nn.GroupNorm(32, dim)

networks/models.py

@@ -0,0 +1,221 @@
+import torch
+import torch.nn as nn
+from networks.resnet_GN_WS import ResNet
+import networks.layers_WS as L
+
+
+def build_model(weights):
+    net_encoder = fba_encoder()
+
+    net_decoder = fba_decoder()
+
+    model = MattingModule(net_encoder, net_decoder)
+
+    if weights != 'default':
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        sd = torch.load(weights, map_location=device)
+        model.load_state_dict(sd, strict=True)
+
+    return model
+
+
+class MattingModule(nn.Module):
+    def __init__(self, net_enc, net_dec):
+        super(MattingModule, self).__init__()
+        self.encoder = net_enc
+        self.decoder = net_dec
+
+    def forward(self, image, two_chan_trimap, image_n, trimap_transformed):
+        resnet_input = torch.cat((image_n, trimap_transformed, two_chan_trimap), 1)
+        conv_out, indices = self.encoder(resnet_input, return_feature_maps=True)
+        return self.decoder(conv_out, image, indices, two_chan_trimap)
+
+
+def fba_encoder():
+    orig_resnet = ResNet()
+    net_encoder = ResnetDilated(orig_resnet, dilate_scale=8)
+
+    num_channels = 3 + 6 + 2
+
+    print(f'modifying input layer to accept {num_channels} channels')
+    net_encoder_sd = net_encoder.state_dict()
+    conv1_weights = net_encoder_sd['conv1.weight']
+
+    c_out, c_in, h, w = conv1_weights.size()
+    conv1_mod = torch.zeros(c_out, num_channels, h, w)
+    conv1_mod[:, :3, :, :] = conv1_weights
+
+    conv1 = net_encoder.conv1
+    conv1.in_channels = num_channels
+    conv1.weight = torch.nn.Parameter(conv1_mod)
+
+    net_encoder.conv1 = conv1
+
+    net_encoder_sd['conv1.weight'] = conv1_mod
+
+    net_encoder.load_state_dict(net_encoder_sd)
+    return net_encoder
+
+
+class ResnetDilated(nn.Module):
+    def __init__(self, orig_resnet, dilate_scale=8):
+        super(ResnetDilated, self).__init__()
+        from functools import partial
+
+        if dilate_scale == 8:
+            orig_resnet.layer3.apply(
+                partial(self._nostride_dilate, dilate=2))
+            orig_resnet.layer4.apply(
+                partial(self._nostride_dilate, dilate=4))
+        elif dilate_scale == 16:
+            orig_resnet.layer4.apply(
+                partial(self._nostride_dilate, dilate=2))
+
+        # take pretrained resnet, except AvgPool and FC
+        self.conv1 = orig_resnet.conv1
+        self.bn1 = orig_resnet.bn1
+        self.relu = orig_resnet.relu
+        self.maxpool = orig_resnet.maxpool
+        self.layer1 = orig_resnet.layer1
+        self.layer2 = orig_resnet.layer2
+        self.layer3 = orig_resnet.layer3
+        self.layer4 = orig_resnet.layer4
+
+    def _nostride_dilate(self, m, dilate):
+        classname = m.__class__.__name__
+        if classname.find('Conv') != -1:
+            # the convolution with stride
+            if m.stride == (2, 2):
+                m.stride = (1, 1)
+                if m.kernel_size == (3, 3):
+                    m.dilation = (dilate // 2, dilate // 2)
+                    m.padding = (dilate // 2, dilate // 2)
+            # other convoluions
+            else:
+                if m.kernel_size == (3, 3):
+                    m.dilation = (dilate, dilate)
+                    m.padding = (dilate, dilate)
+
+    def forward(self, x, return_feature_maps=False):
+        conv_out = [x]
+        x = self.relu(self.bn1(self.conv1(x)))
+        conv_out.append(x)
+        x, indices = self.maxpool(x)
+        x = self.layer1(x)
+        conv_out.append(x)
+        x = self.layer2(x)
+        conv_out.append(x)
+        x = self.layer3(x)
+        conv_out.append(x)
+        x = self.layer4(x)
+        conv_out.append(x)
+
+        if return_feature_maps:
+            return conv_out, indices
+        return [x]
+
+
+def fba_fusion(alpha, img, F, B):
+    F = (alpha * img + (1 - alpha ** 2) * F - alpha * (1 - alpha) * B)
+    B = ((1 - alpha) * img + (2 * alpha - alpha ** 2) * B - alpha * (1 - alpha) * F)
+
+    F = torch.clamp(F, 0, 1)
+    B = torch.clamp(B, 0, 1)
+    la = 0.1
+    alpha = (alpha * la + torch.sum((img - B) * (F - B), 1, keepdim=True)) / (
+                torch.sum((F - B) * (F - B), 1, keepdim=True) + la)
+    alpha = torch.clamp(alpha, 0, 1)
+    return alpha, F, B
+
+
+class fba_decoder(nn.Module):
+    def __init__(self):
+        super(fba_decoder, self).__init__()
+        pool_scales = (1, 2, 3, 6)
+
+        self.ppm = []
+
+        for scale in pool_scales:
+            self.ppm.append(nn.Sequential(
+                nn.AdaptiveAvgPool2d(scale),
+                L.Conv2d(2048, 256, kernel_size=1, bias=True),
+                L.norm(256),
+                nn.LeakyReLU()
+            ))
+        self.ppm = nn.ModuleList(self.ppm)
+
+        self.conv_up1 = nn.Sequential(
+            L.Conv2d(2048 + len(pool_scales) * 256, 256,
+                     kernel_size=3, padding=1, bias=True),
+
+            L.norm(256),
+            nn.LeakyReLU(),
+            L.Conv2d(256, 256, kernel_size=3, padding=1),
+            L.norm(256),
+            nn.LeakyReLU()
+        )
+
+        self.conv_up2 = nn.Sequential(
+            L.Conv2d(256 + 256, 256,
+                     kernel_size=3, padding=1, bias=True),
+            L.norm(256),
+            nn.LeakyReLU()
+        )
+        self.conv_up3 = nn.Sequential(
+            L.Conv2d(256 + 64, 64,
+                     kernel_size=3, padding=1, bias=True),
+            L.norm(64),
+            nn.LeakyReLU()
+        )
+
+        self.unpool = nn.MaxUnpool2d(2, stride=2)
+
+        self.conv_up4 = nn.Sequential(
+            nn.Conv2d(64 + 3 + 3 + 2, 32,
+                      kernel_size=3, padding=1, bias=True),
+            nn.LeakyReLU(),
+            nn.Conv2d(32, 16,
+                      kernel_size=3, padding=1, bias=True),
+
+            nn.LeakyReLU(),
+            nn.Conv2d(16, 7, kernel_size=1, padding=0, bias=True)
+        )
+
+    def forward(self, conv_out, img, indices, two_chan_trimap):
+        conv5 = conv_out[-1]
+
+        input_size = conv5.size()
+        ppm_out = [conv5]
+        for pool_scale in self.ppm:
+            ppm_out.append(nn.functional.interpolate(
+                pool_scale(conv5),
+                (input_size[2], input_size[3]),
+                mode='bilinear', align_corners=False))
+        ppm_out = torch.cat(ppm_out, 1)
+        x = self.conv_up1(ppm_out)
+
+        x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
+
+        x = torch.cat((x, conv_out[-4]), 1)
+
+        x = self.conv_up2(x)
+        x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
+
+        x = torch.cat((x, conv_out[-5]), 1)
+        x = self.conv_up3(x)
+
+        x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
+        x = torch.cat((x, conv_out[-6][:, :3], img, two_chan_trimap), 1)
+
+        output = self.conv_up4(x)
+
+        alpha = torch.clamp(output[:, 0][:, None], 0, 1)
+        F = torch.sigmoid(output[:, 1:4])
+        B = torch.sigmoid(output[:, 4:7])
+
+        # FBA Fusion
+        alpha, F, B = fba_fusion(alpha, img, F, B)
+
+        output = torch.cat((alpha, F, B), 1)
+
+        return output

networks/resnet_GN_WS.py

@@ -0,0 +1,104 @@
+import torch.nn as nn
+import networks.layers_WS as L
+
+__all__ = ['ResNet', 'l_resnet50']
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return L.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                    padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return L.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = L.norm(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = L.norm(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = L.norm(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block=Bottleneck, layers=[3, 4, 6, 3], num_classes=1000):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        self.conv1 = L.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                              bias=False)
+        self.bn1 = L.norm(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                L.norm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x

networks/resnet_bn.py

@@ -0,0 +1,156 @@
+import torch.nn as nn
+import math
+from torch.nn import BatchNorm2d
+
+__all__ = ['ResNet']
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = BatchNorm2d(planes, momentum=0.01)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000):
+        self.inplanes = 128
+        super(ResNet, self).__init__()
+        self.conv1 = conv3x3(3, 64, stride=2)
+        self.bn1 = BatchNorm2d(64)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(64, 64)
+        self.bn2 = BatchNorm2d(64)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = conv3x3(64, 128)
+        self.bn3 = BatchNorm2d(128)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True)
+
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.relu1(self.bn1(self.conv1(x)))
+        x = self.relu2(self.bn2(self.conv2(x)))
+        x = self.relu3(self.bn3(self.conv3(x)))
+        x, indices = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+
+def l_resnet50():
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3])
+    return model

networks/transforms.py

@@ -0,0 +1,27 @@
+import numpy as np
+import torch
+import cv2
+
+
+def dt(a):
+    return cv2.distanceTransform((a * 255).astype(np.uint8), cv2.DIST_L2, 0)
+
+
+def trimap_transform(trimap, L=320):
+    clicks = []
+    for k in range(2):
+        dt_mask = -dt(1 - trimap[:, :, k]) ** 2
+        clicks.append(np.exp(dt_mask / (2 * ((0.02 * L) ** 2))))
+        clicks.append(np.exp(dt_mask / (2 * ((0.08 * L) ** 2))))
+        clicks.append(np.exp(dt_mask / (2 * ((0.16 * L) ** 2))))
+    clicks = np.array(clicks)
+    return clicks
+
+
+# For RGB !
+imagenet_norm_std = torch.from_numpy(np.array([0.229, 0.224, 0.225])).float().cpu()[None, :, None, None]
+imagenet_norm_mean = torch.from_numpy(np.array([0.485, 0.456, 0.406])).float().cpu()[None, :, None, None]
+
+
+def normalise_image(image, mean=imagenet_norm_mean, std=imagenet_norm_std):
+    return (image - mean) / std

requirements.txt

@@ -0,0 +1,3 @@
+torch>=1.4.0
+numpy
+opencv-python