import torch import torch.nn as nn # =================================================== # U-Net の実装 # =================================================== class DoubleConv(nn.Module): “”” 畳み込みブロック: Conv → BN → ReLU → Conv → BN → ReLU U-Netの基本ブロック “”” def __init__(self, in_channels, out_channels): super(DoubleConv, self).__init__() # Sequential: 複数の層を順番に適用 self.double_conv = nn.Sequential( # 1つ目の畳み込み # kernel_size=3: 3×3フィルタ # padding=1: 出力サイズを維持 nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), # Batch Normalization: 学習を安定化 nn.BatchNorm2d(out_channels), # ReLU活性化関数 # inplace=True: メモリ効率のためその場で計算 nn.ReLU(inplace=True), # 2つ目の畳み込み nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x)

class UNet(nn.Module): “”” U-Net アーキテクチャ Args: in_channels: 入力チャンネル数（RGB=3、グレースケール=1） num_classes: 出力クラス数 features: 最初の畳み込み層のフィルタ数（デフォルト64） “”” def __init__(self, in_channels=3, num_classes=2, features=[64, 128, 256, 512]): super(UNet, self).__init__() # ダウンサンプリング用のモジュールリスト self.downs = nn.ModuleList() # アップサンプリング用のモジュールリスト self.ups = nn.ModuleList() # MaxPooling: 解像度を半分に self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # =================================== # Encoder（ダウンサンプリング）部分 # =================================== for feature in features: # 各段階でDoubleConvを追加 self.downs.append(DoubleConv(in_channels, feature)) in_channels = feature # =================================== # Bottleneck（最も深い層） # =================================== # 最後の特徴数の2倍（512→1024） self.bottleneck = DoubleConv(features[-1], features[-1] * 2) # =================================== # Decoder（アップサンプリング）部分 # =================================== # 逆順に処理 for feature in reversed(features): # アップサンプリング: 解像度を2倍に # ConvTranspose2d: 転置畳み込み（逆畳み込み） self.ups.append( nn.ConvTranspose2d( feature * 2, # 入力チャンネル（Skip Connectionで2倍） feature, # 出力チャンネル kernel_size=2, stride=2 ) ) # Skip Connectionで結合後のDoubleConv # feature * 2: アップサンプリング後 + Encoderからの特徴 self.ups.append(DoubleConv(feature * 2, feature)) # =================================== # 最終出力層 # =================================== # 1×1畳み込みでクラス数に変換 self.final_conv = nn.Conv2d(features[0], num_classes, kernel_size=1) def forward(self, x): # Skip Connection用の特徴を保存するリスト skip_connections = [] # =================================== # Encoder（ダウンサンプリング） # =================================== for down in self.downs: x = down(x) # Skip Connection用に保存 skip_connections.append(x) # プーリングでダウンサンプリング x = self.pool(x) # =================================== # Bottleneck # =================================== x = self.bottleneck(x) # Skip Connectionは逆順で使用 skip_connections = skip_connections[::-1] # =================================== # Decoder（アップサンプリング） # =================================== # ups は [UpConv, DoubleConv, UpConv, DoubleConv, …] の順 for idx in range(0, len(self.ups), 2): # アップサンプリング x = self.ups[idx](x) # Skip Connectionから特徴を取得 skip = skip_connections[idx // 2] # サイズが異なる場合は調整（paddingの影響） if x.shape != skip.shape: x = nn.functional.interpolate( x, size=skip.shape[2:], mode=’bilinear’, align_corners=True ) # Concatenate（チャンネル方向に結合） # dim=1: チャンネル次元で結合 x = torch.cat([skip, x], dim=1) # DoubleConvで処理 x = self.ups[idx + 1](x) # =================================== # 最終出力 # =================================== return self.final_conv(x)

# モデルの作成 model = UNet(in_channels=3, num_classes=2) # テスト用の入力（バッチサイズ1、RGB、256×256） x = torch.randn(1, 3, 256, 256) # 順伝播 output = model(x) print(f”入力サイズ: {x.shape}”) print(f”出力サイズ: {output.shape}”)

import torch import torch.nn as nn class DiceLoss(nn.Module): “”” Dice Loss: 1 – Dice係数 クラス不均衡に強い損失関数 “”” def __init__(self, smooth=1.0): super(DiceLoss, self).__init__() # smooth: ゼロ除算を防ぐための小さな値 self.smooth = smooth def forward(self, pred, target): “”” Args: pred: 予測 [batch, num_classes, H, W] target: 正解 [batch, H, W] “”” # Softmaxで確率に変換 pred = torch.softmax(pred, dim=1) # targetをone-hotに変換 num_classes = pred.shape[1] target_one_hot = torch.zeros_like(pred) target_one_hot.scatter_(1, target.unsqueeze(1), 1) # Dice係数を計算 # 各クラスについて計算 intersection = (pred * target_one_hot).sum(dim=(2, 3)) union = pred.sum(dim=(2, 3)) + target_one_hot.sum(dim=(2, 3)) dice = (2.0 * intersection + self.smooth) / (union + self.smooth) # 全クラスの平均を取り、1から引いて損失に dice_loss = 1.0 – dice.mean() return dice_loss # 使用例 criterion = DiceLoss() loss = criterion(output, target)

class CombinedLoss(nn.Module): “”” Cross Entropy Loss + Dice Loss の組み合わせ 両方の利点を活用 “”” def __init__(self, ce_weight=1.0, dice_weight=1.0): super(CombinedLoss, self).__init__() self.ce_loss = nn.CrossEntropyLoss() self.dice_loss = DiceLoss() self.ce_weight = ce_weight self.dice_weight = dice_weight def forward(self, pred, target): ce = self.ce_loss(pred, target) dice = self.dice_loss(pred, target) return self.ce_weight * ce + self.dice_weight * dice # 使用例 criterion = CombinedLoss(ce_weight=1.0, dice_weight=1.0) loss = criterion(output, target)

import torch from torch.utils.data import Dataset, DataLoader import torchvision.transforms as transforms from PIL import Image import numpy as np import os class SegmentationDataset(Dataset): “”” セグメンテーション用データセット ディレクトリ構造: dataset/ ├── images/ │ ├── train/ │ │ ├── img001.jpg │ │ └── … │ └── val/ ├── masks/ │ ├── train/ │ │ ├── img001.png (グレースケール、ピクセル値=クラスID) │ │ └── … │ └── val/ “”” def __init__(self, image_dir, mask_dir, transform=None, target_size=(256, 256)): “”” Args: image_dir: 画像ディレクトリのパス mask_dir: マスクディレクトリのパス transform: 画像の変換（Data Augmentation含む） target_size: リサイズ後のサイズ “”” self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.target_size = target_size # 画像ファイルのリストを取得 self.images = sorted(os.listdir(image_dir)) def __len__(self): return len(self.images) def __getitem__(self, idx): # 画像を読み込み img_name = self.images[idx] img_path = os.path.join(self.image_dir, img_name) image = Image.open(img_path).convert(‘RGB’) # マスクを読み込み（拡張子を変更） mask_name = img_name.replace(‘.jpg’, ‘.png’).replace(‘.jpeg’, ‘.png’) mask_path = os.path.join(self.mask_dir, mask_name) mask = Image.open(mask_path).convert(‘L’) # グレースケール # リサイズ image = image.resize(self.target_size, Image.BILINEAR) mask = mask.resize(self.target_size, Image.NEAREST) # NEARESTでクラスIDを保持 # NumPy配列に変換 image = np.array(image) mask = np.array(mask) # Data Augmentation（画像とマスクに同じ変換を適用） if self.transform: # Albumentationsを使用する場合 augmented = self.transform(image=image, mask=mask) image = augmented[‘image’] mask = augmented[‘mask’] # テンソルに変換 image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask = torch.from_numpy(mask).long() # 正規化 normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ) image = normalize(image) return image, mask

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader def train_model(model, train_loader, val_loader, num_epochs=50, device=’cuda’): “”” U-Netの訓練 Args: model: U-Netモデル train_loader: 訓練データローダー val_loader: 検証データローダー num_epochs: エポック数 device: ‘cuda’ or ‘cpu’ “”” model = model.to(device) # 損失関数（Cross EntropyとDiceの組み合わせ） criterion = CombinedLoss(ce_weight=1.0, dice_weight=1.0) # オプティマイザ optimizer = optim.Adam(model.parameters(), lr=1e-4) # 学習率スケジューラ scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode=’min’, factor=0.5, patience=5 ) best_val_iou = 0.0 for epoch in range(num_epochs): # =================================== # 訓練フェーズ # =================================== model.train() train_loss = 0.0 for images, masks in train_loader: images = images.to(device) masks = masks.to(device) # 順伝播 outputs = model(images) loss = criterion(outputs, masks) # 逆伝播 optimizer.zero_grad() loss.backward() optimizer.step() train_loss += loss.item() train_loss /= len(train_loader) # =================================== # 検証フェーズ # =================================== model.eval() val_loss = 0.0 val_iou = 0.0 with torch.no_grad(): for images, masks in val_loader: images = images.to(device) masks = masks.to(device) outputs = model(images) loss = criterion(outputs, masks) val_loss += loss.item() # IoU計算 preds = torch.argmax(outputs, dim=1) val_iou += compute_miou(preds, masks, num_classes=2) val_loss /= len(val_loader) val_iou /= len(val_loader) # 学習率の調整 scheduler.step(val_loss) # ベストモデルの保存 if val_iou > best_val_iou: best_val_iou = val_iou torch.save(model.state_dict(), ‘best_unet.pth’) print(f’Epoch [{epoch+1}/{num_epochs}]’) print(f’ Train Loss: {train_loss:.4f}’) print(f’ Val Loss: {val_loss:.4f}, Val mIoU: {val_iou:.4f}’) return model def compute_miou(pred, target, num_classes): “””mIoUを計算””” ious = [] for cls in range(num_classes): pred_cls = (pred == cls) target_cls = (target == cls) intersection = (pred_cls & target_cls).sum().item() union = (pred_cls | target_cls).sum().item() if union > 0: ious.append(intersection / union) return np.mean(ious) if ious else 0.0

import matplotlib.pyplot as plt def predict_and_visualize(model, image_path, device=’cuda’): “”” 画像を入力して予測結果を可視化 “”” model.eval() model = model.to(device) # 画像の読み込みと前処理 image = Image.open(image_path).convert(‘RGB’) original_size = image.size # リサイズ image_resized = image.resize((256, 256), Image.BILINEAR) # テンソルに変換 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) input_tensor = transform(image_resized).unsqueeze(0).to(device) # 推論 with torch.no_grad(): output = model(input_tensor) pred = torch.argmax(output, dim=1).squeeze().cpu().numpy() # 元のサイズにリサイズ pred_resized = Image.fromarray(pred.astype(np.uint8)) pred_resized = pred_resized.resize(original_size, Image.NEAREST) pred_resized = np.array(pred_resized) # 可視化 fig, axes = plt.subplots(1, 3, figsize=(15, 5)) # 元画像 axes[0].imshow(image) axes[0].set_title(‘Original Image’) axes[0].axis(‘off’) # 予測マスク axes[1].imshow(pred_resized, cmap=’jet’) axes[1].set_title(‘Prediction Mask’) axes[1].axis(‘off’) # オーバーレイ axes[2].imshow(image) axes[2].imshow(pred_resized, alpha=0.5, cmap=’jet’) axes[2].set_title(‘Overlay’) axes[2].axis(‘off’) plt.tight_layout() plt.savefig(‘segmentation_result.png’) plt.show() return pred_resized # 使用例 # model = UNet(in_channels=3, num_classes=2) # model.load_state_dict(torch.load(‘best_unet.pth’)) # result = predict_and_visualize(model, ‘test_image.jpg’)