Hugoyeah/u-rwkv-busi

u-rwkv-polyp

This model is part of the U-RWKV family of medical image segmentation models. It combines the power of RWKV (Receptance Weighted Key Value) attention mechanism with U-Net architecture for efficient and accurate medical image segmentation.

Model description

U-RWKV model trained on Polyp dataset for polyp segmentation

Architecture

• Base architecture: U-Net with RWKV attention
• Input channels: 3
• Output channels: 1
• Base channels: [16, 32, 128, 160, 256]
• Attention mechanism: RWKV (Receptance Weighted Key Value)
• Feature fusion: SE (Squeeze-and-Excitation)

Performance

• Dice score: 0.7887

Usage

import torch
from models.model import U_RWKV

# Load model
model = U_RWKV()
checkpoint = torch.load('model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Inference
with torch.no_grad():
    output = model(input_image)

Training

The model was trained using:

• Loss functions: Dice Loss + BCE Loss
• Optimizer: AdamW
• Learning rate: 1e-4
• Batch size: 16
• Data augmentation: Random flip, rotation, scaling

License

This model is released under the MIT License.