Upload folder using huggingface_hub

u2net/data_loader.py

@@ -0,0 +1,85 @@
+import os
+import random
+from PIL import Image
+
+import torch
+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):        
+        self.img_size = img_size
+        self.split = split
+        self.image_dir, self.mask_dir = self.set_directories(split)
+        
+        all_images = [f for f in os.listdir(self.image_dir) if f.endswith('.jpg')]
+        if split in ['train', 'valid']:
+            train_imgs, val_imgs = train_test_split(all_images, test_size=val_split_ratio, random_state=42)
+            self.images = train_imgs if split == 'train' else val_imgs
+        else:
+            self.images = all_images
+        
+        self.resize = transforms.Resize((img_size, img_size))
+        self.normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+
+    def __len__(self):
+        return len(self.images)
+    
+    def __getitem__(self, idx):
+        img_filename = self.images[idx]
+        img_path = os.path.join(self.image_dir, img_filename)
+        mask_filename = img_filename.replace('.jpg', '.png')
+        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)
+        
+        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
+            img = transforms.functional.hflip(img)
+            mask = transforms.functional.hflip(mask)
+        
+        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))
+        
+        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)
+        
+        return img, mask
+
+
+class DUTSDataset(SaliencyDataset):
+    def set_directories(self, split):
+        train_or_test = 'train' if split in ['train', 'valid'] else 'test'
+        image_dir = f'/data/duts_{train_or_test}_data/images'
+        mask_dir = f'/data/duts_{train_or_test}_data/masks'
+        return image_dir, mask_dir
+    
+    
+class MSRADataset(SaliencyDataset):
+    def set_directories(self, split):
+        image_dir = '/data/msra_data/images'
+        mask_dir = '/data/msra_data/masks'
+        return image_dir, mask_dir
+    
+
+class PASCALSDataset(SaliencyDataset):
+    def set_directories(self, split):
+        image_dir = '/data/pascals_data/images'
+        mask_dir = '/data/pascals_data/masks'
+        return image_dir, mask_dir

u2net/download_data.sh

@@ -0,0 +1,91 @@
+#!/bin/bash
+
+DUTS_TRAIN_URL="http://saliencydetection.net/duts/download/DUTS-TR.zip"
+DUTS_TEST_URL="http://saliencydetection.net/duts/download/DUTS-TE.zip"
+MSRA10K_URL="http://mftp.mmcheng.net/Data/MSRA10K_Imgs_GT.zip"
+PASCALS_URL="https://cbs.ic.gatech.edu/salobj/download/salObj.zip"
+
+dataset_dir="data_temp"
+mkdir -p $dataset_dir
+
+download_and_extract() {
+  url=$1
+  filename=$(basename $url)
+
+  echo "Downloading $filename..."
+  curl -L -o "$dataset_dir/$filename" $url
+
+  echo "Extracting $filename..."
+  unzip -q "$dataset_dir/$filename" -d $dataset_dir
+  rm "$dataset_dir/$filename"
+}
+
+download_duts() {
+  download_and_extract $DUTS_TRAIN_URL
+  mv "$dataset_dir/DUTS-TR/DUTS-TR-Image" "$dataset_dir/DUTS-TR/images"
+  mv "$dataset_dir/DUTS-TR/DUTS-TR-Mask" "$dataset_dir/DUTS-TR/masks"
+  mv "$dataset_dir/DUTS-TR" "$dataset_dir/duts_train_data"
+  
+  download_and_extract $DUTS_TEST_URL
+  mv "$dataset_dir/DUTS-TE/DUTS-TE-Image" "$dataset_dir/DUTS-TE/images"
+  mv "$dataset_dir/DUTS-TE/DUTS-TE-Mask" "$dataset_dir/DUTS-TE/masks"
+  mv "$dataset_dir/DUTS-TE" "$dataset_dir/duts_test_data"
+}
+
+download_msra() {
+  download_and_extract $MSRA10K_URL
+  rm -f "$dataset_dir/Readme.txt"
+  mkdir -p "$dataset_dir/MSRA10K_Imgs_GT/masks"
+  mv "$dataset_dir/MSRA10K_Imgs_GT/Imgs/"*.png "$dataset_dir/MSRA10K_Imgs_GT/masks"
+  mv "$dataset_dir/MSRA10K_Imgs_GT/Imgs" "$dataset_dir/MSRA10K_Imgs_GT/images"
+  mv "$dataset_dir/MSRA10K_Imgs_GT" "$dataset_dir/msra_data"
+}
+
+download_pascals() {
+  download_and_extract $PASCALS_URL
+  rm -rf "$dataset_dir/algmaps" "$dataset_dir/benchmark" "$dataset_dir/code" "$dataset_dir/results" \
+        "$dataset_dir/readme.pdf" "$dataset_dir/tips_for_matlab.txt" "$dataset_dir/datasets/fixations" \
+        "$dataset_dir/datasets/segments" "$dataset_dir/datasets/imgs/bruce" "$dataset_dir/datasets/imgs/cerf" \
+        "$dataset_dir/datasets/imgs/ft" "$dataset_dir/datasets/imgs/judd" "$dataset_dir/datasets/imgs/pascal" \
+        "$dataset_dir/datasets/masks/bruce" "$dataset_dir/datasets/masks/ft" "$dataset_dir/datasets/masks/pascal"
+  mv "$dataset_dir/datasets/imgs/imgsal"/* "$dataset_dir/datasets/imgs"
+  mv "$dataset_dir/datasets/masks/imgsal"/* "$dataset_dir/datasets/masks"
+  rm -rf "$dataset_dir/datasets/imgs/imgsal" "$dataset_dir/datasets/imgs/Thumbs.db" "$dataset_dir/datasets/masks/imgsal"
+  mv "$dataset_dir/datasets/imgs" "$dataset_dir/datasets/images"
+  mv "$dataset_dir/datasets" "$dataset_dir/pascals_data"
+}
+
+usage() {
+  echo "Usage: $0 [-d] [-m] [-p]"
+  echo "  -d    Download DUTS dataset (train and test)"
+  echo "  -m    Download MSRA10K dataset"
+  echo "  -p    Download Pascal-S dataset"
+  echo "If no options are provided, all datasets will be downloaded."
+  exit 1
+}
+
+all=false
+while getopts "dmp" opt; do
+  case $opt in
+    d)
+      download_duts
+      ;;
+    m)
+      download_msra
+      ;;
+    p)
+      download_pascals
+      ;;
+    *)
+      usage
+      ;;
+  esac
+done
+
+# Check if no options were provided
+if [ $OPTIND -eq 1 ]; then
+  echo "No options provided; downloading all datasets."
+  download_duts
+  download_msra
+  download_pascals
+fi

u2net/evaluate.py

@@ -0,0 +1,41 @@
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from data_loader import PASCALSDataset
+from model import U2Net
+
+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=True)
+    model.load_state_dict(state_dict)
+    model.eval()
+
+def eval(model, loader, criterion):
+    model.eval()
+    running_loss = 0.
+    with torch.no_grad():
+        for images, masks in tqdm(loader, desc='Testing'):
+            images, masks = images.to(device), masks.to(device)
+            outputs = model(images)
+            loss = sum([criterion(output, masks) for output in outputs])
+            running_loss += loss.item()
+    return running_loss / len(loader)
+
+
+if __name__ == '__main__':
+    batch_size = 40
+
+    loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
+    model = U2Net().to(device)
+    model = nn.DataParallel(model)
+    load_model(model, 'results/inter-u2net-duts.pt')
+    
+    loader = DataLoader(PASCALSDataset(split='eval'), batch_size=batch_size, shuffle=False)
+
+    loss = eval(model, loader, loss_fn)
+    print(f'Loss: {loss:.4f}')

u2net/inference.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+
+from model import U2Net
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def preprocess_image(image_path):
+    img = Image.open(image_path).convert('RGB')
+    preprocess = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    img = preprocess(img).unsqueeze(0).to(device)
+    return img
+
+def run_inference(model, image_path, threshold=0.5):
+    input_img = preprocess_image(image_path)
+    with torch.no_grad():
+        d1, *_ = model(input_img)
+        pred = torch.sigmoid(d1)
+        pred = pred[0, :, :].cpu().numpy()
+    
+    pred = (pred - pred.min()) / (pred.max() - pred.min())
+    if threshold is not None:
+        pred = (pred > threshold).astype(np.uint8) * 255
+    else:
+        pred = (pred * 255).astype(np.uint8)
+    return pred
+
+def overlay_segmentation(original_image, binary_mask, alpha=0.5):
+    original_image = Image.open(original_image).convert('RGB').resize((512, 512), Image.BILINEAR)
+    original_image_np = np.array(original_image)
+    overlay = np.zeros_like(original_image_np)
+    overlay[:, :, 0] = binary_mask
+    overlay_image = (1 - alpha) * original_image_np + alpha * overlay
+    overlay_image = overlay_image.astype(np.uint8)
+    return overlay_image
+
+
+if __name__ == '__main__':
+    # ---
+    model_path = 'results/inter-u2net-duts.pt'
+    image_path = 'images/ladies.jpg'
+    # ---
+    model = U2Net().to(device)
+    model = nn.DataParallel(model)
+    state_dict = torch.load(model_path, map_location=device, weights_only=True)
+    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
+    model.eval()
+
+    mask = run_inference(model, image_path, threshold=None)
+    mask_with_threshold = run_inference(model, image_path, threshold=0.7)
+    
+    fig = plt.figure(figsize=(10, 10))
+    gs = GridSpec(2, 2, figure=fig, wspace=0, hspace=0)
+    
+    images = [
+        Image.open(image_path).resize((512, 512)),
+        mask,
+        overlay_segmentation(image_path, mask_with_threshold),
+        mask_with_threshold
+    ]
+    
+    for i, img in enumerate(images):
+        ax = fig.add_subplot(gs[i // 2, i % 2])
+        ax.imshow(img, cmap='gray' if i % 2 != 0 else None)
+        ax.axis('off')
+
+    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+    plt.savefig('inference-output.jpg', format='jpg', bbox_inches='tight', pad_inches=0)

u2net/model.py

@@ -29,7 +29,7 @@ class RSU(nn.Module):
         self.conv = ConvBlock(C_in, C_out)
         
         self.enc = nn.ModuleList([ConvBlock(C_out, M)])
-        for i in range(L-2):
+        for _ in range(L-2):
             self.enc.append(ConvBlock(M, M))
             
         self.mid = ConvBlock(M, M, dilation=2)
@@ -148,4 +148,5 @@ class U2Net(nn.Module):
     
     side_out.append(self.lastconv(torch.cat(side_out, dim=1)))       
 
-    return [torch.sigmoid(s.squeeze(1)) for s in side_out]
+    # logits
+    return [s.squeeze(1) for s in side_out]

u2net/train.py

@@ -1 +1,85 @@
-# for training u2net
+import pickle
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, ConcatDataset
+from torch.amp import autocast, GradScaler
+
+from data_loader import DUTSDataset, MSRADataset
+from model import U2Net
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+scaler = GradScaler()
+
+def train_one_epoch(model, loader, criterion, optimizer):
+    model.train()
+    running_loss = 0.
+    for images, masks in tqdm(loader, desc='Training', leave=False):
+        images, masks = images.to(device, non_blocking=True), masks.to(device, non_blocking=True)
+        
+        optimizer.zero_grad()
+        with autocast(device_type='cuda'):
+            outputs = model(images)
+            loss = sum([criterion(output, masks) for output in outputs])
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+
+        running_loss += loss.item()
+    return running_loss / len(loader)
+
+def validate(model, loader, criterion):
+    model.eval()
+    running_loss = 0.
+    with torch.no_grad():
+        for images, masks in tqdm(loader, desc='Validating', leave=False):
+            images, masks = images.to(device, non_blocking=True), masks.to(device, non_blocking=True)
+            outputs = model(images)
+            loss = sum([criterion(output, masks) for output in outputs])
+            running_loss += loss.item()
+    avg_loss = running_loss / len(loader)
+    return avg_loss
+
+
+if __name__ == '__main__':
+    batch_size = 40
+    valid_batch_size = 80
+    epochs = 100
+    
+    lr = 1e-4
+    loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
+    
+    model_name = 'u2net-duts'
+    model = U2Net()
+    model = torch.nn.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
+    )
+    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
+    )
+    
+    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 (epoch + 1) % 10 == 0:
+            torch.save(model.state_dict(), f'results/inter-{model_name}.pt')
+            
+        print(f'Epoch [{epoch+1}/{epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')
+
+    torch.save(model.state_dict(), f'results/{model_name}.pt')
+    with open('results/loss.txt', 'wb') as f:
+        pickle.dump(losses, f)