Upload model

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,15 @@
+{
+ "architectures": [
+ "BASNetModel"
+ ],
+ "auto_map": {
+ "AutoConfig": "configuration_basnet.BASNetConfig",
+ "AutoModel": "modeling_basnet.BASNetModel"
+ },
+ "kernel_size": 3,
+ "model_type": "basnet",
+ "n_channels": 3,
+ "resnet_model": "microsoft/resnet-34",
+ "torch_dtype": "float32",
+ "transformers_version": "4.42.4"
+}

configuration_basnet.py

@@ -0,0 +1,18 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class BASNetConfig(PretrainedConfig):
+ model_type = "basnet"
+
+ def __init__(
+ self,
+ resnet_model: str = "microsoft/resnet-34",
+ n_channels: int = 3,
+ kernel_size: int = 3,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.resnet_model = resnet_model
+ self.n_channels = n_channels
+
+ self.kernel_size = 3

model.safetensors

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

modeling_basnet.py

@@ -0,0 +1,481 @@
+import logging
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torchvision
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_basnet import BASNetConfig
+
+logger = logging.getLogger(__name__)
+
+
+class RefUnet(nn.Module):
+ def __init__(self, in_ch: int, inc_ch: int) -> None:
+ super().__init__()
+
+ self.conv0 = nn.Conv2d(in_ch, inc_ch, kernel_size=3, padding=1)
+
+ self.conv1 = nn.Conv2d(inc_ch, 64, kernel_size=3, padding=1)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.relu1 = nn.ReLU(inplace=True)
+
+ self.pool1 = nn.MaxPool2d(2, 2, ceil_mode=True)
+
+ self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+ self.bn2 = nn.BatchNorm2d(64)
+ self.relu2 = nn.ReLU(inplace=True)
+
+ self.pool2 = nn.MaxPool2d(2, 2, ceil_mode=True)
+
+ self.conv3 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+ self.bn3 = nn.BatchNorm2d(64)
+ self.relu3 = nn.ReLU(inplace=True)
+
+ self.pool3 = nn.MaxPool2d(2, 2, ceil_mode=True)
+
+ self.conv4 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+ self.bn4 = nn.BatchNorm2d(64)
+ self.relu4 = nn.ReLU(inplace=True)
+
+ self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)
+
+ #####
+
+ self.conv5 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+ self.bn5 = nn.BatchNorm2d(64)
+ self.relu5 = nn.ReLU(inplace=True)
+
+ #####
+
+ self.conv_d4 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+ self.bn_d4 = nn.BatchNorm2d(64)
+ self.relu_d4 = nn.ReLU(inplace=True)
+
+ self.conv_d3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+ self.bn_d3 = nn.BatchNorm2d(64)
+ self.relu_d3 = nn.ReLU(inplace=True)
+
+ self.conv_d2 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+ self.bn_d2 = nn.BatchNorm2d(64)
+ self.relu_d2 = nn.ReLU(inplace=True)
+
+ self.conv_d1 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+ self.bn_d1 = nn.BatchNorm2d(64)
+ self.relu_d1 = nn.ReLU(inplace=True)
+
+ self.conv_d0 = nn.Conv2d(64, 1, kernel_size=3, padding=1)
+
+ self.upscore2 = nn.Upsample(
+ scale_factor=2, mode="bilinear", align_corners=False
+ )
+ # self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear')
+
+ def forward(self, x):
+ hx = x
+ hx = self.conv0(hx)
+
+ hx1 = self.relu1(self.bn1(self.conv1(hx)))
+ hx = self.pool1(hx1)
+
+ hx2 = self.relu2(self.bn2(self.conv2(hx)))
+ hx = self.pool2(hx2)
+
+ hx3 = self.relu3(self.bn3(self.conv3(hx)))
+ hx = self.pool3(hx3)
+
+ hx4 = self.relu4(self.bn4(self.conv4(hx)))
+ hx = self.pool4(hx4)
+
+ hx5 = self.relu5(self.bn5(self.conv5(hx)))
+
+ hx = self.upscore2(hx5)
+
+ d4 = self.relu_d4(self.bn_d4(self.conv_d4(torch.cat((hx, hx4), 1))))
+ hx = self.upscore2(d4)
+
+ d3 = self.relu_d3(self.bn_d3(self.conv_d3(torch.cat((hx, hx3), 1))))
+ hx = self.upscore2(d3)
+
+ d2 = self.relu_d2(self.bn_d2(self.conv_d2(torch.cat((hx, hx2), 1))))
+ hx = self.upscore2(d2)
+
+ d1 = self.relu_d1(self.bn_d1(self.conv_d1(torch.cat((hx, hx1), 1))))
+
+ residual = self.conv_d0(d1)
+
+ return x + residual
+
+
+def conv3x3(in_planes, out_planes, stride=1) -> nn.Conv2d:
+ "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: int = 1
+
+ def __init__(self, inplanes: int, planes: int, stride: int = 1, downsample=None):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(inplanes, planes, stride)
+ self.bn1 = nn.BatchNorm2d(planes)
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = conv3x3(planes, planes)
+ self.bn2 = nn.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 BASNetModel(PreTrainedModel):
+ def __init__(self, config: BASNetConfig) -> None:
+ super().__init__(config)
+
+ resnet = torchvision.models.resnet34(
+ weights=torchvision.models.ResNet34_Weights.IMAGENET1K_V1
+ )
+
+ ## -------------Encoder--------------
+
+ self.inconv = nn.Conv2d(
+ config.n_channels, 64, kernel_size=config.kernel_size, padding=1
+ )
+ self.inbn = nn.BatchNorm2d(64)
+ self.inrelu = nn.ReLU(inplace=True)
+
+ # stage 1
+ self.encoder1 = resnet.layer1 # 256
+ # stage 2
+ self.encoder2 = resnet.layer2 # 128
+ # stage 3
+ self.encoder3 = resnet.layer3 # 64
+ # stage 4
+ self.encoder4 = resnet.layer4 # 32
+
+ self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)
+
+ # stage 5
+ self.resb5_1 = BasicBlock(512, 512)
+ self.resb5_2 = BasicBlock(512, 512)
+ self.resb5_3 = BasicBlock(512, 512) # 16
+
+ self.pool5 = nn.MaxPool2d(2, 2, ceil_mode=True)
+
+ # stage 6
+ self.resb6_1 = BasicBlock(512, 512)
+ self.resb6_2 = BasicBlock(512, 512)
+ self.resb6_3 = BasicBlock(512, 512) # 8
+
+ ## -------------Bridge--------------
+
+ # stage Bridge
+ self.convbg_1 = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, dilation=2, padding=2
+ ) # 8
+ self.bnbg_1 = nn.BatchNorm2d(512)
+ self.relubg_1 = nn.ReLU(inplace=True)
+ self.convbg_m = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, dilation=2, padding=2
+ )
+ self.bnbg_m = nn.BatchNorm2d(512)
+ self.relubg_m = nn.ReLU(inplace=True)
+ self.convbg_2 = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, dilation=2, padding=2
+ )
+ self.bnbg_2 = nn.BatchNorm2d(512)
+ self.relubg_2 = nn.ReLU(inplace=True)
+
+ ## -------------Decoder--------------
+
+ # stage 6d
+ self.conv6d_1 = nn.Conv2d(
+ 1024, 512, kernel_size=config.kernel_size, padding=1
+ ) # 16
+ self.bn6d_1 = nn.BatchNorm2d(512)
+ self.relu6d_1 = nn.ReLU(inplace=True)
+
+ self.conv6d_m = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, dilation=2, padding=2
+ ) ###
+ self.bn6d_m = nn.BatchNorm2d(512)
+ self.relu6d_m = nn.ReLU(inplace=True)
+
+ self.conv6d_2 = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, dilation=2, padding=2
+ )
+ self.bn6d_2 = nn.BatchNorm2d(512)
+ self.relu6d_2 = nn.ReLU(inplace=True)
+
+ # stage 5d
+ self.conv5d_1 = nn.Conv2d(
+ 1024, 512, kernel_size=config.kernel_size, padding=1
+ ) # 16
+ self.bn5d_1 = nn.BatchNorm2d(512)
+ self.relu5d_1 = nn.ReLU(inplace=True)
+
+ self.conv5d_m = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, padding=1
+ ) ###
+ self.bn5d_m = nn.BatchNorm2d(512)
+ self.relu5d_m = nn.ReLU(inplace=True)
+
+ self.conv5d_2 = nn.Conv2d(512, 512, kernel_size=config.kernel_size, padding=1)
+ self.bn5d_2 = nn.BatchNorm2d(512)
+ self.relu5d_2 = nn.ReLU(inplace=True)
+
+ # stage 4d
+ self.conv4d_1 = nn.Conv2d(
+ 1024, 512, kernel_size=config.kernel_size, padding=1
+ ) # 32
+ self.bn4d_1 = nn.BatchNorm2d(512)
+ self.relu4d_1 = nn.ReLU(inplace=True)
+
+ self.conv4d_m = nn.Conv2d(
+ 512, 512, kernel_size=config.kernel_size, padding=1
+ ) ###
+ self.bn4d_m = nn.BatchNorm2d(512)
+ self.relu4d_m = nn.ReLU(inplace=True)
+
+ self.conv4d_2 = nn.Conv2d(512, 256, kernel_size=config.kernel_size, padding=1)
+ self.bn4d_2 = nn.BatchNorm2d(256)
+ self.relu4d_2 = nn.ReLU(inplace=True)
+
+ # stage 3d
+ self.conv3d_1 = nn.Conv2d(
+ 512, 256, kernel_size=config.kernel_size, padding=1
+ ) # 64
+ self.bn3d_1 = nn.BatchNorm2d(256)
+ self.relu3d_1 = nn.ReLU(inplace=True)
+
+ self.conv3d_m = nn.Conv2d(
+ 256, 256, kernel_size=config.kernel_size, padding=1
+ ) ###
+ self.bn3d_m = nn.BatchNorm2d(256)
+ self.relu3d_m = nn.ReLU(inplace=True)
+
+ self.conv3d_2 = nn.Conv2d(256, 128, kernel_size=config.kernel_size, padding=1)
+ self.bn3d_2 = nn.BatchNorm2d(128)
+ self.relu3d_2 = nn.ReLU(inplace=True)
+
+ # stage 2d
+
+ self.conv2d_1 = nn.Conv2d(
+ 256, 128, kernel_size=config.kernel_size, padding=1
+ ) # 128
+ self.bn2d_1 = nn.BatchNorm2d(128)
+ self.relu2d_1 = nn.ReLU(inplace=True)
+
+ self.conv2d_m = nn.Conv2d(
+ 128, 128, kernel_size=config.kernel_size, padding=1
+ ) ###
+ self.bn2d_m = nn.BatchNorm2d(128)
+ self.relu2d_m = nn.ReLU(inplace=True)
+
+ self.conv2d_2 = nn.Conv2d(128, 64, kernel_size=config.kernel_size, padding=1)
+ self.bn2d_2 = nn.BatchNorm2d(64)
+ self.relu2d_2 = nn.ReLU(inplace=True)
+
+ # stage 1d
+ self.conv1d_1 = nn.Conv2d(
+ 128, 64, kernel_size=config.kernel_size, padding=1
+ ) # 256
+ self.bn1d_1 = nn.BatchNorm2d(64)
+ self.relu1d_1 = nn.ReLU(inplace=True)
+
+ self.conv1d_m = nn.Conv2d(
+ 64, 64, kernel_size=config.kernel_size, padding=1
+ ) ###
+ self.bn1d_m = nn.BatchNorm2d(64)
+ self.relu1d_m = nn.ReLU(inplace=True)
+
+ self.conv1d_2 = nn.Conv2d(64, 64, kernel_size=config.kernel_size, padding=1)
+ self.bn1d_2 = nn.BatchNorm2d(64)
+ self.relu1d_2 = nn.ReLU(inplace=True)
+
+ ## -------------Bilinear Upsampling--------------
+ self.upscore6 = nn.Upsample(
+ scale_factor=32, mode="bilinear", align_corners=False
+ ) ###
+ self.upscore5 = nn.Upsample(
+ scale_factor=16, mode="bilinear", align_corners=False
+ )
+ self.upscore4 = nn.Upsample(
+ scale_factor=8, mode="bilinear", align_corners=False
+ )
+ self.upscore3 = nn.Upsample(
+ scale_factor=4, mode="bilinear", align_corners=False
+ )
+ self.upscore2 = nn.Upsample(
+ scale_factor=2, mode="bilinear", align_corners=False
+ )
+
+ # self.upscore6 = nn.Upsample(scale_factor=32, mode='bilinear') ###
+ # self.upscore5 = nn.Upsample(scale_factor=16, mode='bilinear')
+ # self.upscore4 = nn.Upsample(scale_factor=8, mode='bilinear')
+ # self.upscore3 = nn.Upsample(scale_factor=4, mode='bilinear')
+ # self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear')
+
+ ## -------------Side Output--------------
+ self.outconvb = nn.Conv2d(512, 1, kernel_size=3, padding=1)
+ self.outconv6 = nn.Conv2d(512, 1, kernel_size=3, padding=1)
+ self.outconv5 = nn.Conv2d(512, 1, kernel_size=3, padding=1)
+ self.outconv4 = nn.Conv2d(256, 1, kernel_size=3, padding=1)
+ self.outconv3 = nn.Conv2d(128, 1, kernel_size=3, padding=1)
+ self.outconv2 = nn.Conv2d(64, 1, kernel_size=3, padding=1)
+ self.outconv1 = nn.Conv2d(64, 1, kernel_size=3, padding=1)
+
+ ## -------------Refine Module-------------
+ self.refunet = RefUnet(1, 64)
+
+ self.post_init()
+
+ def forward(
+ self, pixel_values: torch.Tensor
+ ) -> Tuple[
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ ]:
+ hx = pixel_values
+
+ ## -------------Encoder-------------
+ hx = self.inconv(hx)
+ hx = self.inbn(hx)
+ hx = self.inrelu(hx)
+
+ h1 = self.encoder1(hx) # 256
+ h2 = self.encoder2(h1) # 128
+ h3 = self.encoder3(h2) # 64
+ h4 = self.encoder4(h3) # 32
+
+ hx = self.pool4(h4) # 16
+
+ hx = self.resb5_1(hx)
+ hx = self.resb5_2(hx)
+ h5 = self.resb5_3(hx)
+
+ hx = self.pool5(h5) # 8
+
+ hx = self.resb6_1(hx)
+ hx = self.resb6_2(hx)
+ h6 = self.resb6_3(hx)
+
+ ## -------------Bridge-------------
+ hx = self.relubg_1(self.bnbg_1(self.convbg_1(h6))) # 8
+ hx = self.relubg_m(self.bnbg_m(self.convbg_m(hx)))
+ hbg = self.relubg_2(self.bnbg_2(self.convbg_2(hx)))
+
+ ## -------------Decoder-------------
+
+ hx = self.relu6d_1(self.bn6d_1(self.conv6d_1(torch.cat((hbg, h6), 1))))
+ hx = self.relu6d_m(self.bn6d_m(self.conv6d_m(hx)))
+ hd6 = self.relu6d_2(self.bn5d_2(self.conv6d_2(hx)))
+
+ hx = self.upscore2(hd6) # 8 -> 16
+
+ hx = self.relu5d_1(self.bn5d_1(self.conv5d_1(torch.cat((hx, h5), 1))))
+ hx = self.relu5d_m(self.bn5d_m(self.conv5d_m(hx)))
+ hd5 = self.relu5d_2(self.bn5d_2(self.conv5d_2(hx)))
+
+ hx = self.upscore2(hd5) # 16 -> 32
+
+ hx = self.relu4d_1(self.bn4d_1(self.conv4d_1(torch.cat((hx, h4), 1))))
+ hx = self.relu4d_m(self.bn4d_m(self.conv4d_m(hx)))
+ hd4 = self.relu4d_2(self.bn4d_2(self.conv4d_2(hx)))
+
+ hx = self.upscore2(hd4) # 32 -> 64
+
+ hx = self.relu3d_1(self.bn3d_1(self.conv3d_1(torch.cat((hx, h3), 1))))
+ hx = self.relu3d_m(self.bn3d_m(self.conv3d_m(hx)))
+ hd3 = self.relu3d_2(self.bn3d_2(self.conv3d_2(hx)))
+
+ hx = self.upscore2(hd3) # 64 -> 128
+
+ hx = self.relu2d_1(self.bn2d_1(self.conv2d_1(torch.cat((hx, h2), 1))))
+ hx = self.relu2d_m(self.bn2d_m(self.conv2d_m(hx)))
+ hd2 = self.relu2d_2(self.bn2d_2(self.conv2d_2(hx)))
+
+ hx = self.upscore2(hd2) # 128 -> 256
+
+ hx = self.relu1d_1(self.bn1d_1(self.conv1d_1(torch.cat((hx, h1), 1))))
+ hx = self.relu1d_m(self.bn1d_m(self.conv1d_m(hx)))
+ hd1 = self.relu1d_2(self.bn1d_2(self.conv1d_2(hx)))
+
+ ## -------------Side Output-------------
+ db = self.outconvb(hbg)
+ db = self.upscore6(db) # 8->256
+
+ d6 = self.outconv6(hd6)
+ d6 = self.upscore6(d6) # 8->256
+
+ d5 = self.outconv5(hd5)
+ d5 = self.upscore5(d5) # 16->256
+
+ d4 = self.outconv4(hd4)
+ d4 = self.upscore4(d4) # 32->256
+
+ d3 = self.outconv3(hd3)
+ d3 = self.upscore3(d3) # 64->256
+
+ d2 = self.outconv2(hd2)
+ d2 = self.upscore2(d2) # 128->256
+
+ d1 = self.outconv1(hd1) # 256
+
+ ## -------------Refine Module-------------
+ dout = self.refunet(d1) # 256
+
+ return (
+ torch.sigmoid(dout),
+ torch.sigmoid(d1),
+ torch.sigmoid(d2),
+ torch.sigmoid(d3),
+ torch.sigmoid(d4),
+ torch.sigmoid(d5),
+ torch.sigmoid(d6),
+ torch.sigmoid(db),
+ )
+
+
+def convert_from_checkpoint(
+ repo_id: str, filename: str, config: Optional[BASNetConfig] = None
+) -> BASNetModel:
+ from huggingface_hub import hf_hub_download
+
+ checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+ config = config or BASNetConfig()
+ model = BASNetModel(config)
+
+ logger.info(f"Loading checkpoint from {checkpoint_path}")
+ state_dict = torch.load(checkpoint_path)
+
+ model.load_state_dict(state_dict, strict=True)
+ model.eval()
+
+ return model