Jia-Liu/U2Net-CoreML

U-2-Net (U2NET) - Core ML Format

U-2-Net model converted to Apple Core ML format for macOS GPU-accelerated inference.

📖 Original Model

• Source Repository: https://github.com/xuebinqin/U-2-Net
• License: Apache 2.0
• Paper: U²-Net: Going Deeper with Nested U-Structure for Salient Object Detection
• Authors: Xuebin Qin, Zichen Zhang, Chenyang Huang, Masood Dehghan, Osmar R. Zaïane, Martin Jagersand

🔄 Conversion Details

This model was converted from PyTorch to Core ML format using the following method:

• Conversion Method: torch.jit.trace → coremltools.convert()
• Input Shape: (1, 3, 320, 320) RGB image
• Precision: FP32
• Compute Units: CPU + GPU
• Output Format: Core ML Package (.mlpackage)

Conversion Script

import torch
import coremltools as ct
from u2net import U2NET

# Load model
u2net = U2NET(3, 1)
u2net.load_state_dict(torch.load("u2net.pth", map_location="cpu"))
u2net.eval()

# Trace and convert
dummy = torch.randn(1, 3, 320, 320)
traced = torch.jit.trace(u2net, dummy)
coreml_model = ct.convert(traced, inputs=[ct.ImageType(shape=(1, 3, 320, 320))])
coreml_model.save("U2NET.mlpackage")

📥 Usage

Load the model in a macOS application using Core ML:

import CoreML
import Vision

let config = MLModelConfiguration()
config.computeUnits = .cpuAndGPU

let model = try U2NET(configuration: config)

📊 Model Inputs/Outputs

Input:

• image: 3-channel RGB image, 320×320 pixels

Outputs (7 saliency maps at different scales):

• 7 output tensors, each 1-channel, 320×320

⚖️ License & Attribution

This repository contains a format conversion of the original U-2-Net model. All rights and credit belong to the original authors.

• Original License: Apache 2.0
• Conversion: Format-only, no weight modifications

📚 Citation

If you use this model, please cite the original paper:

@article{qin2020u2net,
  title={U$^2$-Net: Going Deeper with Nested U-Structure for Salient Object Detection},
  author={Qin, Xuebin and Zhang, Zichen and Huang, Chenyang and Dehghan, Masood and Zaiane, Osmar R and Jagersand, Martin},
  journal={Pattern Recognition},
  volume={106},
  pages={107404},
  year={2020},
  publisher={Elsevier}
}