senoquant 1.0.0b1__py3-none-any.whl → 1.0.0b3__py3-none-any.whl

ufish/model/loss.py

@@ -0,0 +1,62 @@
+import torch
+from torch import nn
+
+
+class DiceLoss(nn.Module):
+    def __init__(self):
+        super(DiceLoss, self).__init__()
+
+    def _dice_loss(self, score, target):
+        target = target.float()
+        smooth = 1e-5
+        intersect = torch.sum(score * target)
+        y_sum = torch.sum(target * target)
+        z_sum = torch.sum(score * score)
+        loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
+        loss = 1 - loss
+        return loss
+
+    def forward(self, y_hat, y):
+        return self._dice_loss(y_hat, y)
+
+
+class RMSELoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mse = nn.MSELoss()
+
+    def forward(self, y_hat, y):
+        return torch.sqrt(self.mse(y_hat, y))
+
+
+class DiceRMSELoss(nn.Module):
+    def __init__(self, dice_ratio=0.6, rmse_ratio=0.4):
+        super().__init__()
+        self.dice_loss = DiceLoss()
+        self.rmse_loss = RMSELoss()
+        self.dice_ratio = dice_ratio
+        self.rmse_ratio = rmse_ratio
+
+    def forward(self, y_hat, y):
+        _dice = self.dice_loss(y_hat, y)
+        _dice = self.dice_ratio * _dice
+        _rmse = self.rmse_loss(y_hat, y)
+        _rmse = self.rmse_ratio * _rmse
+        return _dice + _rmse
+
+
+class DiceCoefLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def _dice_coef(self, y_true, y_pred):
+        smooth = 1.0
+        y_true_f = torch.flatten(y_true)
+        y_pred_f = torch.flatten(y_pred)
+        intersection = torch.sum(y_true_f * y_pred_f)
+        a = (2. * intersection + smooth)
+        b = (torch.sum(y_true_f) + torch.sum(y_pred_f) + smooth)
+        return a / b
+
+    def forward(self, y_hat, y):
+        return - self._dice_coef(y, y_hat)

ufish/model/network/spot_learn.py

@@ -0,0 +1,50 @@
+import torch.nn as nn
+import torch
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout2d(p=0.2)
+        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1)
+        self.relu2 = nn.ReLU()
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.relu1(out)
+        out = self.dropout1(out)
+        out = self.conv2(out)
+        out = self.relu2(out)
+        return out
+
+
+class SpotLearn(nn.Module):
+    def __init__(self, input_channel: int = 1) -> None:
+        super().__init__()
+        self.conv1 = ConvBlock(input_channel, 64)
+        self.down1 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv2 = ConvBlock(64, 128)
+        self.down2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.bottom = ConvBlock(128, 256)
+        self.up1 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.conv3 = ConvBlock(256, 128)
+        self.up2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.conv4 = ConvBlock(128, 64)
+        self.conv5 = nn.Conv2d(64, 1, kernel_size=1, padding=0)
+
+    def forward(self, x):
+        in1 = self.conv1(x)
+        x = self.down1(in1)
+        in2 = self.conv2(x)
+        x = self.down2(in2)
+        x = self.bottom(x)
+        x = self.up1(x)
+        x = torch.cat([x, in2], dim=1)
+        x = self.conv3(x)
+        x = self.up2(x)
+        x = torch.cat([x, in1], dim=1)
+        x = self.conv4(x)
+        x = self.conv5(x)
+        return x

ufish/model/network/ufish_net.py

@@ -0,0 +1,204 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        out = F.relu(self.conv1(x))
+        out = F.relu(self.conv2(out))
+        return out
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            channels, channels, kernel_size=1, padding=0)
+        self.conv2 = nn.Conv2d(
+            channels, channels, kernel_size=3, padding=1)
+        self.BN1 = nn.BatchNorm2d(channels)
+        self.BN2 = nn.BatchNorm2d(channels)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.BN1(out)
+        out = F.relu(out)
+        out = self.conv2(out)
+        out = self.BN2(out)
+        out = F.relu(out)
+        out = out + x
+        return out
+
+
+class DownConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DownConv, self).__init__()
+        self.down_conv = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=2, padding=1)
+        self.conv = ResidualBlock(out_channels)
+
+    def forward(self, x):
+        out = self.down_conv(x)
+        out = self.conv(out)
+        return out
+
+
+class UpConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UpConv, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        up = F.interpolate(x, scale_factor=2, mode='nearest')
+        out = self.conv(up)
+        return out
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, input_nc, ratio=16):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+        self.f1 = nn.Conv2d(input_nc, input_nc // ratio, 1, bias=False)
+        self.relu = nn.ReLU()
+        self.f2 = nn.Conv2d(input_nc // ratio, input_nc, 1, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
+        max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
+        out = self.sigmoid(avg_out + max_out)
+        return out
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
+        padding = 3 if kernel_size == 7 else 1
+        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        # 1*h*w
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        # 2*h*w
+        x = self.conv(x)
+        # 1*h*w
+        return self.sigmoid(x)
+
+
+class CBAM(nn.Module):
+    """Convolutional Block Attention Module"""
+    def __init__(self, input_nc, ratio=16, kernel_size=7):
+        super(CBAM, self).__init__()
+        self.channel_attention = ChannelAttention(input_nc, ratio)
+        self.spatial_attention = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        out = self.channel_attention(x) * x
+        out = self.spatial_attention(out) * out
+        return out
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(EncoderBlock, self).__init__()
+        self.conv = ConvBlock(in_channels, out_channels)
+
+    def forward(self, x):
+        out = self.conv(x)
+        return out
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DecoderBlock, self).__init__()
+        self.conv = ConvBlock(in_channels, out_channels)
+
+    def forward(self, x):
+        out = self.conv(x)
+        return out
+
+
+class BottoleneckBlock(nn.Module):
+    def __init__(self, channels):
+        super(BottoleneckBlock, self).__init__()
+        self.conv1 = ResidualBlock(channels)
+        self.cbam = CBAM(channels)
+        self.conv2 = ResidualBlock(channels)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.cbam(out)
+        out = self.conv2(out)
+        return out
+
+
+class FinalDecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(FinalDecoderBlock, self).__init__()
+        self.cbam = CBAM(in_channels)
+        self.conv = ConvBlock(in_channels, out_channels)
+
+    def forward(self, x):
+        out = self.cbam(x)
+        out = self.conv(out)
+        return out
+
+
+class UFishNet(nn.Module):
+    def __init__(
+            self, in_channels=1, out_channels=1,
+            channel_numbers=[32, 32, 32]):
+        super().__init__()
+        self.encoders = nn.ModuleList()
+        self.downsamples = nn.ModuleList()
+        for i, c in enumerate(channel_numbers[:-1]):
+            _ei = (in_channels if i == 0 else c)
+            self.encoders.append(EncoderBlock(_ei, c))
+            self.downsamples.append(DownConv(c, channel_numbers[i+1]))
+
+        self.bottom = BottoleneckBlock(channel_numbers[-1])
+
+        self.decoders = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        rev_nums = channel_numbers[::-1]
+        for i, c in enumerate(rev_nums[1:]):
+            self.upsamples.append(UpConv(rev_nums[i], c))
+            self.decoders.append(DecoderBlock(2*c, c))
+
+        self.final_decoder = FinalDecoderBlock(
+            channel_numbers[0], out_channels)
+
+    def forward(self, x):
+        encodings = []
+        for i, encoder in enumerate(self.encoders):
+            x = encoder(x)
+            encodings.append(x)
+            x = self.downsamples[i](x)
+
+        x = self.bottom(x)
+
+        for i, decoder in enumerate(self.decoders):
+            x = self.upsamples[i](x)
+            diffY = encodings[-i - 1].size()[2] - x.size()[2]
+            diffX = encodings[-i - 1].size()[3] - x.size()[3]
+            x = F.pad(x, (diffX // 2, diffX - diffX // 2,
+                          diffY // 2, diffY - diffY // 2))
+            x = torch.cat([x, encodings[-i - 1]], dim=1)
+            x = decoder(x)
+
+        x = self.final_decoder(x)
+        return x

ufish/model/train.py

@@ -0,0 +1,175 @@
+import typing as T
+from pathlib import Path
+
+import torch
+import numpy as np
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader
+from torch import Tensor
+
+from . import loss as loss_mod  # noqa: F401
+from ..data import FISHSpotsDataset
+from ..utils.log import logger
+
+
+def training_loop(
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        criterion: T.Callable[[Tensor, Tensor], Tensor],
+        writer: SummaryWriter,
+        device: torch.device,
+        train_loader: DataLoader, valid_loader: DataLoader,
+        model_save_dir: str,
+        save_period: int,
+        num_epochs=50,
+        ):
+    """
+    The training loop.
+
+    Args:
+        model: The model to train.
+        optimizer: The optimizer to use.
+        criterion: The loss function to use.
+        writer: The TensorBoard writer.
+        device: The device to use.
+        train_loader: The training data loader.
+        valid_loader: The valid data loader.
+        model_save_dir: The directory to save the model to.
+        save_period: Save the model every `save_period` epochs.
+        num_epochs: The number of epochs to train for.
+    """
+    best_val_loss = float("inf")
+    model_dir_path = Path(model_save_dir)
+    model_dir_path.mkdir(parents=True, exist_ok=True)
+
+    for epoch in range(num_epochs):
+        model.train()
+        epoch_loss = 0.0
+        for idx, batch in enumerate(train_loader):
+            images = batch["image"].to(device, dtype=torch.float)
+            targets = batch["target"].to(device, dtype=torch.float)
+
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, targets)
+            loss.backward()
+            optimizer.step()
+
+            if idx % 10 == 0:
+                logger.info(
+                    f"Epoch: {epoch + 1}/{num_epochs}, "
+                    f"Batch: {idx + 1}/{len(train_loader)}, "
+                    f"Loss: {loss.item():.4f}"
+                )
+                writer.add_scalar(
+                    "Loss/train_batch", loss.item(),
+                    epoch * len(train_loader) + idx)
+                img = images[0, 0].cpu().numpy()
+                img = np.stack((img,)*3, axis=0)
+                # normalize to 0-255
+                img = (img - img.min()) / (img.max() - img.min()) * 255
+                # record images
+                writer.add_image(
+                    "Image/input",
+                    images[0], epoch * len(train_loader) + idx)
+                writer.add_image(
+                    "Image/target",
+                    targets[0], epoch * len(train_loader) + idx)
+                writer.add_image(
+                    "Image/pred",
+                    outputs[0], epoch * len(train_loader) + idx)
+
+            epoch_loss += loss.item()
+
+        epoch_loss /= len(train_loader)
+        writer.add_scalar("Loss/train", epoch_loss, epoch)
+
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for batch in valid_loader:
+                images = batch["image"].to(device, dtype=torch.float)
+                targets = batch["target"].to(device, dtype=torch.float)
+
+                outputs = model(images)
+                loss = criterion(outputs, targets)
+                val_loss += loss.item()
+
+        val_loss /= len(valid_loader)
+        writer.add_scalar("Loss/val", val_loss, epoch)
+
+        logger.info(
+            f"Epoch {epoch + 1}/{num_epochs}, "
+            f"Train Loss: {epoch_loss:.4f}, Val Loss: {val_loss:.4f}"
+        )
+
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            torch.save(
+                model.state_dict(),
+                f"{model_save_dir}/best_model.pth")
+            logger.info(f"Best model saved with Val Loss: {val_loss:.4f}")
+        if epoch % save_period == 0:
+            torch.save(
+                model.state_dict(),
+                f"{model_save_dir}/model_{epoch}.pth")
+            logger.info(f"Model saved at epoch {epoch}.")
+
+    writer.close()
+
+
+def train_on_dataset(
+        model: torch.nn.Module,
+        train_dataset: FISHSpotsDataset,
+        valid_dataset: FISHSpotsDataset,
+        loss_type: str = "DiceRMSELoss",
+        loader_workers: int = 4,
+        num_epochs: int = 50,
+        batch_size: int = 8,
+        lr: float = 1e-4,
+        summary_dir: str = "runs/unet",
+        model_save_dir: str = "./models",
+        save_period: int = 5,
+        ):
+    """Train the UNet model.
+
+    Args:
+        model: The model to train.
+        train_dataset: The training dataset.
+        valid_dataset: The validation dataset.
+        loss_type: The loss function to use.
+        loader_workers: The number of workers to use for the data loader.
+        num_epochs: The number of epochs to train for.
+        batch_size: The batch size.
+        lr: The learning rate.
+        summary_dir: The directory to save the TensorBoard summary to.
+        model_save_dir: The directory to save the model to.
+        save_period: Save the model every `save_period` epochs.
+    """
+    logger.info(
+        f"Loader workers: {loader_workers}, " +
+        f"TensorBoard summary dir: {summary_dir}"
+    )
+    logger.info(
+        f"Model save dir: {model_save_dir}, " +
+        f"Save period: {save_period}"
+    )
+    train_loader = DataLoader(
+        train_dataset, batch_size=batch_size,
+        shuffle=True, num_workers=loader_workers)
+    valid_loader = DataLoader(
+        valid_dataset, batch_size=batch_size,
+        shuffle=False, num_workers=loader_workers)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Training using device: {device}")
+    model = model.to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+
+    criterion = eval(f"loss_mod.{loss_type}")()
+
+    writer = SummaryWriter(summary_dir)
+    training_loop(
+        model, optimizer, criterion, writer, device,
+        train_loader, valid_loader, model_save_dir,
+        save_period, num_epochs)