Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+example.png filter=lfs diff=lfs merge=lfs -text
+results.png filter=lfs diff=lfs merge=lfs -text
+Screenshot[[:space:]]2024-01-21[[:space:]]at[[:space:]]11.56.17.png filter=lfs diff=lfs merge=lfs -text
+T1.png filter=lfs diff=lfs merge=lfs -text
+T2.png filter=lfs diff=lfs merge=lfs -text
+t4.png filter=lfs diff=lfs merge=lfs -text
+example_input.jpg filter=lfs diff=lfs merge=lfs -text

MyConfig.py

@@ -0,0 +1,13 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class RMBGConfig(PretrainedConfig):
+    model_type = "SegformerForSemanticSegmentation"
+    def __init__(
+        self,
+        in_ch=3,
+        out_ch=1,
+        **kwargs):
+      self.in_ch = in_ch
+      self.out_ch = out_ch
+      super().__init__(**kwargs)

MyPipe.py

@@ -0,0 +1,76 @@
+import torch, os
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import numpy as np
+from transformers import Pipeline
+from transformers.image_utils import load_image
+from skimage import io
+from PIL import Image
+
+class RMBGPipe(Pipeline):
+  def __init__(self,**kwargs):
+    Pipeline.__init__(self,**kwargs)
+    self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    self.model.to(self.device)
+    self.model.eval()
+
+  def _sanitize_parameters(self, **kwargs):
+    # parse parameters
+    preprocess_kwargs = {}
+    postprocess_kwargs = {}
+    if "model_input_size" in kwargs : 
+      preprocess_kwargs["model_input_size"] = kwargs["model_input_size"]
+    if "return_mask" in kwargs: 
+      postprocess_kwargs["return_mask"] = kwargs["return_mask"]
+    return preprocess_kwargs, {}, postprocess_kwargs
+
+  def preprocess(self,input_image,model_input_size: list=[1024,1024]):
+      # preprocess the input 
+      orig_im = load_image(input_image)
+      orig_im = np.array(orig_im)
+      orig_im_size = orig_im.shape[0:2]
+      preprocessed_image = self.preprocess_image(orig_im, model_input_size).to(self.device)
+      inputs = {
+          "preprocessed_image":preprocessed_image,
+          "orig_im_size":orig_im_size,
+          "input_image" : input_image
+      }
+      return inputs
+
+  def _forward(self,inputs):
+    result = self.model(inputs.pop("preprocessed_image"))
+    inputs["result"] = result
+    return inputs
+  
+  def postprocess(self,inputs,return_mask:bool=False ):
+    result = inputs.pop("result")
+    orig_im_size = inputs.pop("orig_im_size")
+    input_image = inputs.pop("input_image")
+    result_image = self.postprocess_image(result[0][0], orig_im_size)
+    pil_im = Image.fromarray(result_image)
+    if return_mask ==True : 
+      return pil_im
+    input_image = load_image(input_image)
+    no_bg_image = input_image.copy()
+    no_bg_image.putalpha(pil_im)
+    return no_bg_image
+    
+  # utilities functions
+  def preprocess_image(self,im: np.ndarray, model_input_size: list=[1024,1024]) -> torch.Tensor:
+    # same as utilities.py with minor modification
+    if len(im.shape) < 3:
+        im = im[:, :, np.newaxis]
+    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
+    im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear')
+    image = torch.divide(im_tensor,255.0)
+    image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
+    return image
+  
+  def postprocess_image(self,result: torch.Tensor, im_size: list)-> np.ndarray:
+      result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)
+      ma = torch.max(result)
+      mi = torch.min(result)
+      result = (result-mi)/(ma-mi)
+      im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
+      im_array = np.squeeze(im_array)
+      return im_array

README.md

@@ -0,0 +1,170 @@
+---
+license: other
+license_name: bria-rmbg-1.4
+license_link: https://bria.ai/bria-huggingface-model-license-agreement/
+pipeline_tag: image-segmentation
+tags:
+- remove background
+- background
+- background-removal
+- Pytorch
+- vision
+- legal liability
+- transformers
+- transformers.js
+
+extra_gated_description: RMBG v1.4 is available as a source-available model for non-commercial use
+extra_gated_heading: "Fill in this form to get instant access"
+extra_gated_fields:
+  Name: text
+  Company/Org name: text
+  Org Type (Early/Growth Startup, Enterprise, Academy): text
+  Role: text
+  Country: text
+  Email: text
+  By submitting this form, I agree to BRIA’s Privacy policy and Terms & conditions, see links below: checkbox
+---
+
+# BRIA Background Removal v1.4 Model Card
+
+RMBG v1.4 is our state-of-the-art background removal model, designed to effectively separate foreground from background in a range of
+categories and image types. This model has been trained on a carefully selected dataset, which includes:
+general stock images, e-commerce, gaming, and advertising content, making it suitable for commercial use cases powering enterprise content creation at scale. 
+The accuracy, efficiency, and versatility currently rival leading source-available models. 
+It is ideal where content safety, legally licensed datasets, and bias mitigation are paramount. 
+
+Developed by BRIA AI, RMBG v1.4 is available as a source-available model for non-commercial use. 
+
+
+To purchase a commercial license, simply click [Here](https://go.bria.ai/3D5EGp0).
+
+
+[CLICK HERE FOR A DEMO](https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4)
+
+**NOTE** New RMBG version available! Check out [RMBG-2.0](https://huggingface.co/briaai/RMBG-2.0)
+
+Join our [Discord community](https://discord.gg/Nxe9YW9zHS) for more information, tutorials, tools, and to connect with other users!
+
+
+![examples](t4.png)
+
+
+### Model Description
+
+- **Developed by:** [BRIA AI](https://bria.ai/)
+- **Model type:** Background Removal 
+- **License:** [bria-rmbg-1.4](https://bria.ai/bria-huggingface-model-license-agreement/)
+  - The model is released under a Creative Commons license for non-commercial use.
+  - Commercial use is subject to a commercial agreement with BRIA. To purchase a commercial license simply click [Here](https://go.bria.ai/3B4Asxv).
+
+- **Model Description:** BRIA RMBG 1.4 is a saliency segmentation model trained exclusively on a professional-grade dataset.
+- **BRIA:** Resources for more information: [BRIA AI](https://bria.ai/)
+
+
+
+## Training data
+Bria-RMBG model was trained with over 12,000 high-quality, high-resolution, manually labeled (pixel-wise accuracy), fully licensed images.
+Our benchmark included balanced gender, balanced ethnicity, and people with different types of disabilities.
+For clarity, we provide our data distribution according to different categories, demonstrating our model’s versatility.
+
+### Distribution of images:
+
+| Category | Distribution |
+| -----------------------------------| -----------------------------------:|
+| Objects only | 45.11% |
+| People with objects/animals | 25.24% |
+| People only | 17.35% |
+| people/objects/animals with text | 8.52% |
+| Text only | 2.52% |
+| Animals only | 1.89% |
+
+| Category | Distribution |
+| -----------------------------------| -----------------------------------------:|
+| Photorealistic | 87.70% |
+| Non-Photorealistic | 12.30% |
+
+
+| Category | Distribution |
+| -----------------------------------| -----------------------------------:|
+| Non Solid Background | 52.05% |
+| Solid Background | 47.95% 
+
+
+| Category | Distribution |
+| -----------------------------------| -----------------------------------:|
+| Single main foreground object | 51.42% |
+| Multiple objects in the foreground | 48.58% |
+
+
+## Qualitative Evaluation
+
+![examples](results.png)
+
+
+## Architecture
+
+RMBG v1.4 is developed on the [IS-Net](https://github.com/xuebinqin/DIS) enhanced with our unique training scheme and proprietary dataset. 
+These modifications significantly improve the model’s accuracy and effectiveness in diverse image-processing scenarios.
+
+## Installation
+```bash
+pip install -qr https://huggingface.co/briaai/RMBG-1.4/resolve/main/requirements.txt
+```
+
+## Usage
+
+Either load the pipeline
+```python
+from transformers import pipeline
+image_path = "https://farm5.staticflickr.com/4007/4322154488_997e69e4cf_z.jpg"
+pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+pillow_mask = pipe(image_path, return_mask = True) # outputs a pillow mask
+pillow_image = pipe(image_path) # applies mask on input and returns a pillow image
+```
+
+Or load the model 
+```python
+from transformers import AutoModelForImageSegmentation
+from torchvision.transforms.functional import normalize
+model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4",trust_remote_code=True)
+def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
+    if len(im.shape) < 3:
+        im = im[:, :, np.newaxis]
+    # orig_im_size=im.shape[0:2]
+    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
+    im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear')
+    image = torch.divide(im_tensor,255.0)
+    image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
+    return image
+
+def postprocess_image(result: torch.Tensor, im_size: list)-> np.ndarray:
+    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result-mi)/(ma-mi)
+    im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
+    im_array = np.squeeze(im_array)
+    return im_array
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# prepare input
+image_path = "https://farm5.staticflickr.com/4007/4322154488_997e69e4cf_z.jpg"
+orig_im = io.imread(image_path)
+orig_im_size = orig_im.shape[0:2]
+image = preprocess_image(orig_im, model_input_size).to(device)
+
+# inference 
+result=model(image)
+
+# post process
+result_image = postprocess_image(result[0][0], orig_im_size)
+
+# save result
+pil_mask_im = Image.fromarray(result_image)
+orig_image = Image.open(image_path)
+no_bg_image = orig_image.copy()
+no_bg_image.putalpha(pil_mask_im)
+```
+

briarmbg.py

@@ -0,0 +1,458 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from .MyConfig import RMBGConfig
+
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1,stride=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate,stride=stride)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.interpolate(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
+        super(RSU7,self).__init__()
+
+        self.in_ch = in_ch
+        self.mid_ch = mid_ch
+        self.out_ch = out_ch
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) ## 1 -> 1/2
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+        b, c, h, w = x.shape
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+
+class myrebnconv(nn.Module):
+    def __init__(self, in_ch=3,
+                       out_ch=1,
+                       kernel_size=3,
+                       stride=1,
+                       padding=1,
+                       dilation=1,
+                       groups=1):
+        super(myrebnconv,self).__init__()
+
+        self.conv = nn.Conv2d(in_ch,
+                              out_ch,
+                              kernel_size=kernel_size,
+                              stride=stride,
+                              padding=padding,
+                              dilation=dilation,
+                              groups=groups)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.rl = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+        return self.rl(self.bn(self.conv(x)))
+
+
+class BriaRMBG(PreTrainedModel):
+    config_class = RMBGConfig    
+    def __init__(self,config:RMBGConfig = RMBGConfig()):
+        super().__init__(config)
+        in_ch = config.in_ch # 3
+        out_ch = config.out_ch # 1
+        self.conv_in = nn.Conv2d(in_ch,64,3,stride=2,padding=1)
+        self.pool_in = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage1 = RSU7(64,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.conv_in(hx)
+        #hx = self.pool_in(hxin)
+
+        #stage 1
+        hx1 = self.stage1(hxin)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #-------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+        d1 = _upsample_like(d1,x)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,x)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,x)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,x)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,x)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,x)
+
+        return [F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)],[hx1d,hx2d,hx3d,hx4d,hx5d,hx6]
+

config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "briaai/RMBG-1.4",
+  "architectures": [
+    "BriaRMBG"
+  ],
+  "auto_map": {
+    "AutoConfig": "MyConfig.RMBGConfig",
+    "AutoModelForImageSegmentation": "briarmbg.BriaRMBG"
+  },
+  "custom_pipelines": {
+    "image-segmentation": {
+      "impl": "MyPipe.RMBGPipe",
+      "pt": [
+        "AutoModelForImageSegmentation"
+      ],
+      "tf": [],
+      "type": "image"
+    }
+  },
+  "in_ch": 3,
+  "model_type": "SegformerForSemanticSegmentation",
+  "out_ch": 1,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.0.dev0"
+}

example_inference.py

@@ -0,0 +1,39 @@
+from skimage import io
+import torch, os
+from PIL import Image
+from briarmbg import BriaRMBG
+from utilities import preprocess_image, postprocess_image
+from huggingface_hub import hf_hub_download
+
+def example_inference():
+
+    im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
+
+    net = BriaRMBG()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
+    net.to(device)
+    net.eval()    
+
+    # prepare input
+    model_input_size = [1024,1024]
+    orig_im = io.imread(im_path)
+    orig_im_size = orig_im.shape[0:2]
+    image = preprocess_image(orig_im, model_input_size).to(device)
+
+    # inference 
+    result=net(image)
+
+    # post process
+    result_image = postprocess_image(result[0][0], orig_im_size)
+
+    # save result
+    pil_mask_im = Image.fromarray(result_image)
+    orig_image = Image.open(im_path)
+    no_bg_image = orig_image.copy()
+    no_bg_image.putalpha(pil_mask_im)    
+    no_bg_image.save("example_image_no_bg.png")
+
+
+if __name__ == "__main__":
+    example_inference()

handler.py

@@ -0,0 +1,21 @@
+from typing import Dict, List, Any
+from transformers import pipeline
+from PIL import Image
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Initialize the image segmentation pipeline
+        self.pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # Extract the image path from the input data
+        image_path = data.get("image_path", "")
+        
+        # Perform image segmentation
+        pillow_mask = self.pipe(image_path, return_mask=True)  # outputs a pillow mask
+        pillow_image = self.pipe(image_path)  # outputs the segmented image
+        
+        # Save the segmented image at the root folder
+        output_image_path = "segmented_image.png"
+        pillow_image.save(output_image_path)
+        
+        # Return the result as a list of dictionaries
+        return [{"image_path": output_image_path, "mask": pillow_mask}]

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:893c16c340b1ddafc93e78457a4d94190da9b7179149f8574284c83caebf5e8c
+size 176718373

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46ef7fe46f2ae284d8f1aaa24bfa5fca5ef25a34e2c7caa890a0029eb100e87f
+size 176381984

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cafcf770b06757c4eaced21b1a88e57fd2b66de01b8045f35f01535ba742e0f
+size 176153355

onnx/model_fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fdfdb41866d872e0acf4a010c35c1a8547bf0eebe0d1544406bbf1c824cb59d
+size 88217533

onnx/model_quantized.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6648479275dfd0ede0f3a8abc20aa5c437b394681b05e5af6d268250aaf40f3
+size 44403226

onnx/quantize_config.json

@@ -0,0 +1,24 @@
+{
+    "per_channel": false,
+    "reduce_range": false,
+    "per_model_config": {
+        "model": {
+            "op_types": [
+                "Concat",
+                "MaxPool",
+                "Resize",
+                "Conv",
+                "Unsqueeze",
+                "Cast",
+                "Shape",
+                "Relu",
+                "Sigmoid",
+                "Gather",
+                "Constant",
+                "Slice",
+                "Add"
+            ],
+            "weight_type": "QUInt8"
+        }
+    }
+}

preprocessor_config.json

@@ -0,0 +1,23 @@
+{
+  "do_normalize": true,
+  "do_pad": false,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "feature_extractor_type": "ImageFeatureExtractor",
+  "image_std": [
+    1,
+    1,
+    1
+  ],
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "width": 1024,
+    "height": 1024
+  }
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59569acdb281ac9fc9f78f9d33b6f9f17f68e25086b74f9025c35bb5f2848967
+size 176574018

requirements.txt

@@ -0,0 +1,619 @@
+absl-py==1.4.0
+accelerate==1.7.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.11.15
+aiosignal==1.3.2
+alabaster==1.0.0
+albucore==0.0.24
+albumentations==2.0.7
+ale-py==0.11.0
+altair==5.5.0
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.9.0
+argon2-cffi==23.1.0
+argon2-cffi-bindings==21.2.0
+array_record==0.7.2
+arviz==0.21.0
+astropy==7.1.0
+astropy-iers-data==0.2025.5.19.0.38.36
+astunparse==1.6.3
+atpublic==5.1
+attrs==25.3.0
+audioread==3.0.1
+autograd==1.8.0
+babel==2.17.0
+backcall==0.2.0
+backports.tarfile==1.2.0
+beautifulsoup4==4.13.4
+betterproto==2.0.0b6
+bigframes==2.4.0
+bigquery-magics==0.9.0
+bleach==6.2.0
+blinker==1.9.0
+blis==1.3.0
+blobfile==3.0.0
+blosc2==3.3.3
+bokeh==3.7.3
+Bottleneck==1.4.2
+bqplot==0.12.45
+branca==0.8.1
+build==1.2.2.post1
+CacheControl==0.14.3
+cachetools==5.5.2
+catalogue==2.0.10
+certifi==2025.4.26
+cffi==1.17.1
+chardet==5.2.0
+charset-normalizer==3.4.2
+chex==0.1.89
+clarabel==0.10.0
+click==8.2.1
+cloudpathlib==0.21.1
+cloudpickle==3.1.1
+cmake==3.31.6
+cmdstanpy==1.2.5
+colorcet==3.1.0
+colorlover==0.3.0
+colour==0.1.5
+community==1.0.0b1
+confection==0.1.5
+cons==0.4.6
+contourpy==1.3.2
+cramjam==2.10.0
+cryptography==43.0.3
+cuda-python==12.6.2.post1
+cudf-cu12 @ https://pypi.nvidia.com/cudf-cu12/cudf_cu12-25.2.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+cudf-polars-cu12==25.2.2
+cufflinks==0.17.3
+cuml-cu12==25.2.1
+cupy-cuda12x==13.3.0
+curl_cffi==0.11.1
+cuvs-cu12==25.2.1
+cvxopt==1.3.2
+cvxpy==1.6.5
+cycler==0.12.1
+cyipopt==1.5.0
+cymem==2.0.11
+Cython==3.0.12
+dask==2024.12.1
+dask-cuda==25.2.0
+dask-cudf-cu12==25.2.2
+dask-expr==1.1.21
+dataproc-spark-connect==0.7.4
+datascience==0.17.6
+datasets==2.14.4
+db-dtypes==1.4.3
+dbus-python==1.2.18
+debugpy==1.8.0
+decorator==4.4.2
+defusedxml==0.7.1
+diffusers==0.33.1
+dill==0.3.7
+distributed==2024.12.1
+distributed-ucxx-cu12==0.42.0
+distro==1.9.0
+dlib==19.24.6
+dm-tree==0.1.9
+docker-pycreds==0.4.0
+docstring_parser==0.16
+docutils==0.21.2
+dopamine_rl==4.1.2
+duckdb==1.2.2
+earthengine-api==1.5.15
+easydict==1.13
+editdistance==0.8.1
+eerepr==0.1.2
+einops==0.8.1
+en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85
+entrypoints==0.4
+et_xmlfile==2.0.0
+etils==1.12.2
+etuples==0.3.9
+Farama-Notifications==0.0.4
+fastai==2.7.19
+fastcore==1.7.29
+fastdownload==0.0.7
+fastjsonschema==2.21.1
+fastprogress==1.0.3
+fastrlock==0.8.3
+filelock==3.18.0
+firebase-admin==6.8.0
+Flask==3.1.1
+flatbuffers==25.2.10
+flax==0.10.6
+folium==0.19.6
+fonttools==4.58.0
+frozendict==2.4.6
+frozenlist==1.6.0
+fsspec==2025.3.2
+future==1.0.0
+gast==0.6.0
+gcsfs==2025.3.2
+GDAL==3.8.4
+gdown==5.2.0
+geemap==0.35.3
+geocoder==1.38.1
+geographiclib==2.0
+geopandas==1.0.1
+geopy==2.4.1
+gin-config==0.5.0
+gitdb==4.0.12
+GitPython==3.1.44
+glob2==0.7
+google==2.0.3
+google-ai-generativelanguage==0.6.15
+google-api-core==2.24.2
+google-api-python-client==2.169.0
+google-auth==2.38.0
+google-auth-httplib2==0.2.0
+google-auth-oauthlib==1.2.2
+google-cloud-aiplatform==1.93.1
+google-cloud-bigquery==3.33.0
+google-cloud-bigquery-connection==1.18.2
+google-cloud-bigquery-storage==2.31.0
+google-cloud-core==2.4.3
+google-cloud-dataproc==5.18.1
+google-cloud-datastore==2.21.0
+google-cloud-firestore==2.20.2
+google-cloud-functions==1.20.3
+google-cloud-iam==2.19.0
+google-cloud-language==2.17.1
+google-cloud-resource-manager==1.14.2
+google-cloud-spanner==3.54.0
+google-cloud-storage==2.19.0
+google-cloud-translate==3.20.2
+google-colab @ file:///colabtools/dist/google_colab-1.0.0.tar.gz
+google-crc32c==1.7.1
+google-genai==1.16.1
+google-generativeai==0.8.5
+google-pasta==0.2.0
+google-resumable-media==2.7.2
+googleapis-common-protos==1.70.0
+googledrivedownloader==1.1.0
+graphviz==0.20.3
+greenlet==3.2.2
+grpc-google-iam-v1==0.14.2
+grpc-interceptor==0.15.4
+grpcio==1.71.0
+grpcio-status==1.71.0
+grpclib==0.4.8
+gspread==6.2.1
+gspread-dataframe==4.0.0
+gym==0.25.2
+gym-notices==0.0.8
+gymnasium==1.1.1
+h11==0.16.0
+h2==4.2.0
+h5netcdf==1.6.1
+h5py==3.13.0
+hdbscan==0.8.40
+hf_transfer==0.1.9
+highspy==1.10.0
+holidays==0.73
+holoviews==1.20.2
+hpack==4.1.0
+html5lib==1.1
+httpcore==1.0.9
+httpimport==1.4.1
+httplib2==0.22.0
+httpx==0.28.1
+huggingface-hub==0.31.4
+humanize==4.12.3
+hyperframe==6.1.0
+hyperopt==0.2.7
+ibis-framework==9.5.0
+idna==3.10
+imageio==2.37.0
+imageio-ffmpeg==0.6.0
+imagesize==1.4.1
+imbalanced-learn==0.13.0
+immutabledict==4.2.1
+importlib_metadata==8.7.0
+importlib_resources==6.5.2
+imutils==0.5.4
+inflect==7.5.0
+iniconfig==2.1.0
+intel-cmplr-lib-ur==2025.1.1
+intel-openmp==2025.1.1
+ipyevents==2.0.2
+ipyfilechooser==0.6.0
+ipykernel==6.17.1
+ipyleaflet==0.19.2
+ipyparallel==8.8.0
+ipython==7.34.0
+ipython-genutils==0.2.0
+ipython-sql==0.5.0
+ipytree==0.2.2
+ipywidgets==7.7.1
+itsdangerous==2.2.0
+jaraco.classes==3.4.0
+jaraco.context==6.0.1
+jaraco.functools==4.1.0
+jax==0.5.2
+jax-cuda12-pjrt==0.5.1
+jax-cuda12-plugin==0.5.1
+jaxlib==0.5.1
+jeepney==0.9.0
+jieba==0.42.1
+Jinja2==3.1.6
+jiter==0.10.0
+joblib==1.5.0
+jsonpatch==1.33
+jsonpickle==4.1.0
+jsonpointer==3.0.0
+jsonschema==4.23.0
+jsonschema-specifications==2025.4.1
+jupyter-client==6.1.12
+jupyter-console==6.1.0
+jupyter-leaflet==0.19.2
+jupyter-server==1.16.0
+jupyter_core==5.7.2
+jupyter_kernel_gateway @ git+https://github.com/googlecolab/kernel_gateway@b134e9945df25c2dcb98ade9129399be10788671
+jupyterlab_pygments==0.3.0
+jupyterlab_widgets==3.0.15
+kaggle==1.7.4.5
+kagglehub==0.3.12
+keras==3.8.0
+keras-hub==0.18.1
+keras-nlp==0.18.1
+keyring==25.6.0
+keyrings.google-artifactregistry-auth==1.1.2
+kiwisolver==1.4.8
+langchain==0.3.25
+langchain-core==0.3.60
+langchain-text-splitters==0.3.8
+langcodes==3.5.0
+langsmith==0.3.42
+language_data==1.3.0
+launchpadlib==1.10.16
+lazr.restfulclient==0.14.4
+lazr.uri==1.0.6
+lazy_loader==0.4
+libclang==18.1.1
+libcudf-cu12 @ https://pypi.nvidia.com/libcudf-cu12/libcudf_cu12-25.2.1-py3-none-manylinux_2_28_x86_64.whl
+libcugraph-cu12==25.2.0
+libcuml-cu12==25.2.1
+libcuvs-cu12==25.2.1
+libkvikio-cu12==25.2.1
+libpysal==4.13.0
+libraft-cu12==25.2.0
+librosa==0.11.0
+libucx-cu12==1.18.1
+libucxx-cu12==0.42.0
+lightgbm @ file:///tmp/lightgbm/LightGBM/dist/lightgbm-4.5.0-py3-none-linux_x86_64.whl
+linkify-it-py==2.0.3
+llvmlite==0.43.0
+locket==1.0.0
+logical-unification==0.4.6
+lxml==5.4.0
+Mako==1.1.3
+marisa-trie==1.2.1
+Markdown==3.8
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+matplotlib==3.10.0
+matplotlib-inline==0.1.7
+matplotlib-venn==1.1.2
+mdit-py-plugins==0.4.2
+mdurl==0.1.2
+miniKanren==1.0.3
+missingno==0.5.2
+mistune==3.1.3
+mizani==0.13.5
+mkl==2025.0.1
+ml-dtypes==0.4.1
+mlxtend==0.23.4
+more-itertools==10.7.0
+moviepy==1.0.3
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.4.4
+multipledispatch==1.0.0
+multiprocess==0.70.15
+multitasking==0.0.11
+murmurhash==1.0.12
+music21==9.3.0
+namex==0.0.9
+narwhals==1.40.0
+natsort==8.4.0
+nbclassic==1.3.1
+nbclient==0.10.2
+nbconvert==7.16.6
+nbformat==5.10.4
+ndindex==1.10.0
+nest-asyncio==1.6.0
+networkx==3.4.2
+nibabel==5.3.2
+nltk==3.9.1
+notebook==6.5.7
+notebook_shim==0.2.4
+numba==0.60.0
+numba-cuda==0.2.0
+numexpr==2.10.2
+numpy==2.0.2
+nvidia-cublas-cu12==12.5.3.2
+nvidia-cuda-cupti-cu12==12.5.82
+nvidia-cuda-nvcc-cu12==12.5.82
+nvidia-cuda-nvrtc-cu12==12.5.82
+nvidia-cuda-runtime-cu12==12.5.82
+nvidia-cudnn-cu12==9.3.0.75
+nvidia-cufft-cu12==11.2.3.61
+nvidia-curand-cu12==10.3.6.82
+nvidia-cusolver-cu12==11.6.3.83
+nvidia-cusparse-cu12==12.5.1.3
+nvidia-cusparselt-cu12==0.6.2
+nvidia-ml-py==12.575.51
+nvidia-nccl-cu12==2.21.5
+nvidia-nvcomp-cu12==4.2.0.11
+nvidia-nvjitlink-cu12==12.5.82
+nvidia-nvtx-cu12==12.4.127
+nvtx==0.2.11
+nx-cugraph-cu12 @ https://pypi.nvidia.com/nx-cugraph-cu12/nx_cugraph_cu12-25.2.0-py3-none-any.whl
+oauth2client==4.1.3
+oauthlib==3.2.2
+omegaconf==2.3.0
+openai==1.81.0
+opencv-contrib-python==4.11.0.86
+opencv-python==4.11.0.86
+opencv-python-headless==4.11.0.86
+openpyxl==3.1.5
+opt_einsum==3.4.0
+optax==0.2.4
+optree==0.15.0
+orbax-checkpoint==0.11.13
+orjson==3.10.18
+osqp==1.0.4
+packaging==24.2
+pandas==2.2.2
+pandas-datareader==0.10.0
+pandas-gbq==0.29.0
+pandas-stubs==2.2.2.240909
+pandocfilters==1.5.1
+panel==1.7.0
+param==2.2.0
+parso==0.8.4
+parsy==2.1
+partd==1.4.2
+pathlib==1.0.1
+patsy==1.0.1
+peewee==3.18.1
+peft==0.15.2
+pexpect==4.9.0
+pickleshare==0.7.5
+pillow==11.2.1
+platformdirs==4.3.8
+plotly==5.24.1
+plotnine==0.14.5
+pluggy==1.6.0
+ply==3.11
+polars==1.21.0
+pooch==1.8.2
+portpicker==1.5.2
+preshed==3.0.9
+prettytable==3.16.0
+proglog==0.1.12
+progressbar2==4.5.0
+prometheus_client==0.22.0
+promise==2.3
+prompt_toolkit==3.0.51
+propcache==0.3.1
+prophet==1.1.6
+proto-plus==1.26.1
+protobuf==5.29.4
+psutil==5.9.5
+psycopg2==2.9.10
+ptyprocess==0.7.0
+py-cpuinfo==9.0.0
+py4j==0.10.9.7
+pyarrow==18.1.0
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
+pycairo==1.28.0
+pycocotools==2.0.8
+pycparser==2.22
+pycryptodomex==3.23.0
+pydantic==2.11.4
+pydantic_core==2.33.2
+pydata-google-auth==1.9.1
+pydot==3.0.4
+pydotplus==2.0.2
+PyDrive==1.3.1
+PyDrive2==1.21.3
+pyerfa==2.0.1.5
+pygame==2.6.1
+pygit2==1.18.0
+Pygments==2.19.1
+PyGObject==3.42.0
+PyJWT==2.10.1
+pylibcudf-cu12 @ https://pypi.nvidia.com/pylibcudf-cu12/pylibcudf_cu12-25.2.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+pylibcugraph-cu12==25.2.0
+pylibraft-cu12==25.2.0
+pymc==5.22.0
+pymystem3==0.2.0
+pynndescent==0.5.13
+pynvjitlink-cu12==0.6.0
+pynvml==12.0.0
+pyogrio==0.11.0
+pyomo==6.9.2
+PyOpenGL==3.1.9
+pyOpenSSL==24.2.1
+pyparsing==3.2.3
+pyperclip==1.9.0
+pyproj==3.7.1
+pyproject_hooks==1.2.0
+pyshp==2.3.1
+PySocks==1.7.1
+pyspark==3.5.1
+pytensor==2.30.3
+pytest==8.3.5
+python-apt==0.0.0
+python-box==7.3.2
+python-dateutil==2.9.0.post0
+python-louvain==0.16
+python-slugify==8.0.4
+python-snappy==0.7.3
+python-utils==3.9.1
+pytz==2025.2
+pyviz_comms==3.0.4
+PyWavelets==1.8.0
+PyYAML==6.0.2
+pyzmq==24.0.1
+raft-dask-cu12==25.2.0
+rapids-dask-dependency==25.2.0
+ratelim==0.1.6
+referencing==0.36.2
+regex==2024.11.6
+requests==2.32.3
+requests-oauthlib==2.0.0
+requests-toolbelt==1.0.0
+requirements-parser==0.9.0
+rich==13.9.4
+rmm-cu12==25.2.0
+roman-numerals-py==3.1.0
+rpds-py==0.25.1
+rpy2==3.5.17
+rsa==4.9.1
+safetensors==0.5.3
+scikit-image==0.25.2
+scikit-learn==1.6.1
+scipy==1.15.3
+scooby==0.10.1
+scs==3.2.7.post2
+seaborn==0.13.2
+SecretStorage==3.3.3
+Send2Trash==1.8.3
+sentence-transformers==4.1.0
+sentencepiece==0.2.0
+sentry-sdk==2.29.1
+setproctitle==1.3.6
+shap==0.47.2
+shapely==2.1.1
+shellingham==1.5.4
+simple-parsing==0.1.7
+simplejson==3.20.1
+simsimd==6.2.1
+six==1.17.0
+sklearn-compat==0.1.3
+sklearn-pandas==2.2.0
+slicer==0.0.8
+smart-open==7.1.0
+smmap==5.0.2
+sniffio==1.3.1
+snowballstemmer==3.0.1
+sortedcontainers==2.4.0
+soundfile==0.13.1
+soupsieve==2.7
+soxr==0.5.0.post1
+spacy==3.8.6
+spacy-legacy==3.0.12
+spacy-loggers==1.0.5
+spanner-graph-notebook==1.1.6
+Sphinx==8.2.3
+sphinxcontrib-applehelp==2.0.0
+sphinxcontrib-devhelp==2.0.0
+sphinxcontrib-htmlhelp==2.1.0
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==2.0.0
+sphinxcontrib-serializinghtml==2.0.0
+SQLAlchemy==2.0.41
+sqlglot==25.20.2
+sqlparse==0.5.3
+srsly==2.5.1
+stanio==0.5.1
+statsmodels==0.14.4
+stringzilla==3.12.5
+stumpy==1.13.0
+sympy==1.13.1
+tables==3.10.2
+tabulate==0.9.0
+tbb==2022.1.0
+tblib==3.1.0
+tcmlib==1.3.0
+tenacity==9.1.2
+tensorboard==2.18.0
+tensorboard-data-server==0.7.2
+tensorflow==2.18.0
+tensorflow-datasets==4.9.8
+tensorflow-hub==0.16.1
+tensorflow-io-gcs-filesystem==0.37.1
+tensorflow-metadata==1.17.1
+tensorflow-probability==0.25.0
+tensorflow-text==2.18.1
+tensorflow_decision_forests==1.11.0
+tensorstore==0.1.74
+termcolor==3.1.0
+terminado==0.18.1
+text-unidecode==1.3
+textblob==0.19.0
+tf-slim==1.1.0
+tf_keras==2.18.0
+thinc==8.3.6
+threadpoolctl==3.6.0
+tifffile==2025.5.21
+tiktoken==0.9.0
+timm==1.0.15
+tinycss2==1.4.0
+tokenizers==0.21.1
+toml==0.10.2
+toolz==0.12.1
+torch @ https://download.pytorch.org/whl/cu124/torch-2.6.0%2Bcu124-cp311-cp311-linux_x86_64.whl
+torchao==0.10.0
+torchaudio @ https://download.pytorch.org/whl/cu124/torchaudio-2.6.0%2Bcu124-cp311-cp311-linux_x86_64.whl
+torchdata==0.11.0
+torchsummary==1.5.1
+torchtune==0.6.1
+torchvision @ https://download.pytorch.org/whl/cu124/torchvision-0.21.0%2Bcu124-cp311-cp311-linux_x86_64.whl
+tornado==6.4.2
+tqdm==4.67.1
+traitlets==5.7.1
+traittypes==0.2.1
+transformers==4.52.2
+treelite==4.4.1
+treescope==0.1.9
+triton==3.2.0
+tsfresh==0.21.0
+tweepy==4.15.0
+typeguard==4.4.2
+typer==0.15.3
+types-pytz==2025.2.0.20250516
+types-setuptools==80.8.0.20250521
+typing-inspection==0.4.1
+typing_extensions==4.13.2
+tzdata==2025.2
+tzlocal==5.3.1
+uc-micro-py==1.0.3
+ucx-py-cu12==0.42.0
+ucxx-cu12==0.42.0
+umap-learn==0.5.7
+umf==0.10.0
+uritemplate==4.1.1
+urllib3==2.4.0
+vega-datasets==0.9.0
+wadllib==1.3.6
+wandb==0.19.11
+wasabi==1.1.3
+wcwidth==0.2.13
+weasel==0.4.1
+webcolors==24.11.1
+webencodings==0.5.1
+websocket-client==1.8.0
+websockets==15.0.1
+Werkzeug==3.1.3
+widgetsnbextension==3.6.10
+wordcloud==1.9.4
+wrapt==1.17.2
+wurlitzer==3.1.1
+xarray==2025.3.1
+xarray-einstats==0.8.0
+xgboost==2.1.4
+xlrd==2.0.1
+xxhash==3.5.0
+xyzservices==2025.4.0
+yarl==1.20.0
+ydf==0.12.0
+yellowbrick==1.5
+yfinance==0.2.61
+zict==3.0.0
+zipp==3.21.0
+zstandard==0.23.0

utilities.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import numpy as np
+
+def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
+    if len(im.shape) < 3:
+        im = im[:, :, np.newaxis]
+    # orig_im_size=im.shape[0:2]
+    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
+    im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear').type(torch.uint8)
+    image = torch.divide(im_tensor,255.0)
+    image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
+    return image
+
+
+def postprocess_image(result: torch.Tensor, im_size: list)-> np.ndarray:
+    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result-mi)/(ma-mi)
+    im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
+    im_array = np.squeeze(im_array)
+    return im_array
+