Upload folder using huggingface_hub

u2net/data_loader.py

@@ -3,12 +3,13 @@ import random
 from PIL import Image
 
 import torch
+import numpy as np
 from torchvision import transforms
 from sklearn.model_selection import train_test_split
-
+    
 
 class SaliencyDataset(torch.utils.data.Dataset):
-    def __init__(self, split, img_size=512, val_split_ratio=0.05):        
+    def __init__(self, split, img_size=512, val_split_ratio=0.05, subset_ratio=None):        
         self.img_size = img_size
         self.split = split
         self.image_dir, self.mask_dir = self.set_directories(split)
@@ -19,8 +20,14 @@ class SaliencyDataset(torch.utils.data.Dataset):
             self.images = train_imgs if split == 'train' else val_imgs
         else:
             self.images = all_images
+            
+        if subset_ratio: # subsampling
+            total_samples = len(self.images)
+            indices = np.random.choice(total_samples, int(total_samples * subset_ratio), replace=False)
+            self.images = [self.images[i] for i in indices]
         
-        self.resize = transforms.Resize((img_size, img_size))
+        self.img_resize = transforms.Resize((img_size, img_size), interpolation=transforms.InterpolationMode.BILINEAR)
+        self.mask_resize = transforms.Resize((img_size, img_size), interpolation=transforms.InterpolationMode.NEAREST)
         self.normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
 
     def __len__(self):
@@ -33,30 +40,38 @@ class SaliencyDataset(torch.utils.data.Dataset):
         mask_path = os.path.join(self.mask_dir, mask_filename)
         
         img = Image.open(img_path).convert('RGB')
-        mask = Image.open(mask_path).convert('L')
-        
-        img, mask = self.resize(img), self.resize(mask)
+        mask = Image.open(mask_path)
+        if mask.mode != 'L': mask = mask.convert('L')
+        mask = mask.point(lambda p: 255 if p > 128 else 0)
         
+        img, mask = self.img_resize(img), self.mask_resize(mask)
         if self.split == 'train':
             img, mask = self.apply_augmentations(img, mask)
         
         img = transforms.ToTensor()(img)
         img = self.normalize(img)
         mask = transforms.ToTensor()(mask).squeeze(0)
+        
         return img, mask
     
     def apply_augmentations(self, img, mask):
-        if random.random() > 0.5:  # horizontal flip
+        if random.random() > 0.5: # horizontal flip
             img = transforms.functional.hflip(img)
             mask = transforms.functional.hflip(mask)
         
-        if random.random() > 0.5:  # random resized crop
+        if random.random() > 0.5: # random resized crop
             resized_crop = transforms.RandomResizedCrop(self.img_size, scale=(0.8, 1.0))
             i, j, h, w = resized_crop.get_params(img, scale=(0.8, 1.0), ratio=(3/4, 4/3))
-            img = transforms.functional.resized_crop(img, i, j, h, w, (self.img_size, self.img_size))
-            mask = transforms.functional.resized_crop(mask, i, j, h, w, (self.img_size, self.img_size))
+            img = transforms.functional.resized_crop(
+                img, i, j, h, w, (self.img_size, self.img_size), 
+                interpolation=transforms.InterpolationMode.BILINEAR
+            )
+            mask = transforms.functional.resized_crop(
+                mask, i, j, h, w, (self.img_size, self.img_size), 
+                interpolation=transforms.InterpolationMode.NEAREST
+            )
         
-        if random.random() > 0.5:  # color jitter
+        if random.random() > 0.5: # color jitter
             color_jitter = transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.05)
             img = color_jitter(img)
         

u2net/evaluate.py

@@ -3,6 +3,7 @@ from tqdm import tqdm
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
+from safetensors.torch import load_file
 
 from data_loader import PASCALSDataset
 from model import U2Net
@@ -11,7 +12,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 print('Device:', device)
 
 def load_model(model, model_path):
-    state_dict = torch.load(model_path, map_location=device, weights_only=False)
+    state_dict = load_file(model_path, device=device.type)
     model.load_state_dict(state_dict)
     model.eval()
 
@@ -28,12 +29,14 @@ def eval(model, loader, criterion):
 
 
 if __name__ == '__main__':
-    batch_size = 40
+    batch_size = 1
 
+    model_type = input('Model type [b,f]: ')
+    model_name = 'best-u2net-duts-msra.safetensors' if model_type == 'b' else 'u2net-duts-msra.safetensors'
     loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
     model = U2Net().to(device)
     model = nn.DataParallel(model)
-    load_model(model, 'results/inter-u2net-duts.pt')
+    load_model(model, f'results/{model_name}')
     
     loader = DataLoader(PASCALSDataset(split='eval'), batch_size=batch_size, shuffle=False)
 

u2net/model.py

@@ -9,17 +9,19 @@ def init_weight(layer):
 
     
 class ConvBlock(nn.Module):
-    def __init__(self, in_channel, out_channel, dilation=1):
+    def __init__(self, in_channel, out_channel, dilation=1, dropout_rate=0.3):
         super(ConvBlock, self).__init__()
         self.conv = nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=dilation, dilation=dilation)
         self.bn = nn.BatchNorm2d(out_channel)
         self.relu = nn.ReLU(inplace=True)
+        self.dropout = nn.Dropout2d(p=dropout_rate) # custom - add dropout layer
         init_weight(self.conv)
         
     def forward(self, x):
         x = self.conv(x)
         x = self.bn(x)
         x = self.relu(x)
+        self.dropout(x)
         return x
 
 
@@ -93,7 +95,7 @@ class RSU4F(nn.Module):
     
     
 class U2Net(nn.Module):
-  def __init__(self):
+  def __init__(self, dropout_rate=0.3):
     super(U2Net, self).__init__()
     self.enc = nn.ModuleList([
       RSU(L=7, C_in=3, C_out=64, M=32),
@@ -123,6 +125,7 @@ class U2Net(nn.Module):
 
     self.lastconv = nn.Conv2d(6, 1, 1)
     self.downsample = nn.MaxPool2d(2, stride=2)
+    self.dropout = nn.Dropout(p=dropout_rate) # custom - add dropout layer
 
     init_weight(self.lastconv)
     for conv in self.convs:
@@ -143,10 +146,10 @@ class U2Net(nn.Module):
     
     side_out = []
     for i, conv in enumerate(self.convs):
-      if i == 0: side_out.append(conv(dec_out[5]))
-      else: side_out.append(self.upsample(conv(dec_out[5-i]), side_out[0]))
+      if i == 0: side_out.append(self.dropout(conv(dec_out[5])))
+      else: side_out.append(self.upsample(self.dropout(conv(dec_out[5-i])), side_out[0]))
     
     side_out.append(self.lastconv(torch.cat(side_out, dim=1)))       
 
-    # logits
+    # logits (no sigmoid)
     return [s.squeeze(1) for s in side_out]

u2net/train.py

@@ -14,6 +14,20 @@ from model import U2Net
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 scaler = GradScaler()
 
+
+class DiceLoss(nn.Module):
+    def __init__(self):
+        super(DiceLoss, self).__init__()
+    
+    def forward(self, inputs, targets, smooth=1):
+        inputs = torch.sigmoid(inputs)
+        inputs = inputs.view(-1)
+        targets = targets.view(-1)
+        intersection = (inputs * targets).sum()
+        dice = (2. * intersection + smooth) / (inputs.sum() + targets.sum() + smooth)
+        return 1 - dice
+
+
 def train_one_epoch(model, loader, criterion, optimizer):
     model.train()
     running_loss = 0.
@@ -43,47 +57,54 @@ def validate(model, loader, criterion):
     avg_loss = running_loss / len(loader)
     return avg_loss
 
+def save(model, model_name, losses):
+    save_file(model.state_dict(), f'results/{model_name}.safetensors')
+    with open('results/loss.txt', 'wb') as f:
+        pickle.dump(losses, f)
+
 
 if __name__ == '__main__':
     batch_size = 40
     valid_batch_size = 80
-    epochs = 100
+    epochs = 200
     
     lr = 1e-3
-    loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
+    loss_fn_bce = nn.BCEWithLogitsLoss(reduction='mean')
+    loss_fn_dice = DiceLoss()
+    alpha = 0.6
+    loss_fn = lambda o, m: alpha * loss_fn_bce(o, m) + (1 - alpha) * loss_fn_dice(o, m)
     
-    model_name = 'u2net-duts'
+    model_name = 'u2net-duts-msra'
     model = U2Net()
-    model = torch.nn.DataParallel(model.to(device))
-    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
+    model = torch.nn.parallel.DataParallel(model.to(device))
+    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
 
     train_loader = DataLoader(
         ConcatDataset([DUTSDataset(split='train'), MSRADataset(split='train')]), 
         batch_size=batch_size, shuffle=True, pin_memory=True, 
-        num_workers=16, persistent_workers=True
+        num_workers=8, persistent_workers=True
     )
     valid_loader = DataLoader(
         ConcatDataset([DUTSDataset(split='valid'), MSRADataset(split='valid')]), 
         batch_size=valid_batch_size, shuffle=False, pin_memory=True, 
-        num_workers=16, persistent_workers=True
+        num_workers=8, persistent_workers=True
     )
     
     best_val_loss = float('inf')
     losses = {'train': [], 'val': []}
-    for epoch in tqdm(range(epochs), desc='Epochs'):
-        torch.cuda.empty_cache()
-        train_loss = train_one_epoch(model, train_loader, loss_fn, optimizer)
-        val_loss = validate(model, valid_loader, loss_fn)
-        losses['train'].append(train_loss)
-        losses['val'].append(val_loss)
-
-        if val_loss < best_val_loss:
-            best_val_loss = val_loss
-            save_file(model.state_dict(), f'results/best-{model_name}.safetensors')
-            print('Best model saved.')
-            
-        print(f'Epoch [{epoch+1}/{epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f} (Best: {best_val_loss:.4f})')
+    
+    # training loop
+    try:
+        for epoch in range(epochs):
+            train_loss = train_one_epoch(model, train_loader, loss_fn, optimizer)
+            val_loss = validate(model, valid_loader, loss_fn)
+            losses['train'].append(train_loss)
+            losses['val'].append(val_loss)
 
-    save_file(model.state_dict(), f'results/{model_name}.safetensors')
-    with open('results/loss.txt', 'wb') as f:
-        pickle.dump(losses, f)
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                save_file(model.state_dict(), f'results/best-{model_name}.safetensors')
+                
+            print(f'Epoch [{epoch+1}/{epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f} (Best: {best_val_loss:.4f})')
+    finally:
+        save(model, model_name, losses)