Hi Nishanth and ptrblck,
Thanks for the discussion. I have a similar problem of detecting thin wires (power lines on utility poles). In my case, I face the same problem of background pixels being in abundance and the masked pixels of wire being negligible compared to the background class. Could you please explain how to pass weights to the loss function? I’m currently using nn.NLLLoss().
Help needed urgently. Any help would be of great help.
Here’s the code I’m using.
from __future__ import print_function, division
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from PIL import Image # w, h when an image is opened
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torchvision import datasets, models
import time
import copy
from tqdm import tqdm

import warnings
warnings.filterwarnings("ignore")

class semanticSegmentationDataset(Dataset):
    """Semantic Segmentation Dataset."""    

    def __init__(self, root_dir, transformList=None):
        """
        Args:
            root_dir (string): train or test Directory with the images and masks folders.
            transformList (callable, optional): Optional transforms to be applied on a sample.
        """
        self.root_dir = root_dir
        self.transformList = transformList
        self.imgs = list(sorted(os.listdir(os.path.join(root_dir, "images"))))
        self.msks = list(sorted(os.listdir(os.path.join(root_dir, "masks"))))

    def __getitem__(self, idx):

        if torch.is_tensor(idx):
            idx = idx.tolist()

        img_path = os.path.join(self.root_dir, "images", self.imgs[idx])
        msk_path = os.path.join(self.root_dir, "masks", self.msks[idx])
        
        image = Image.open(img_path)
        mask = Image.open(msk_path)
                
        sample = {'image': image, 'mask': mask}

        if self.transformList:
            sample = self.transformList(sample)

        return sample

    def __len__(self):
        return len(self.imgs)


class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, sample):
        image, mask = sample['image'], sample['mask']
        
        w, h = image.size

        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size

        new_h, new_w = int(new_h), int(new_w)

        imageT = transforms.Compose([transforms.Resize((new_h,new_w))])
        maskT = transforms.Compose([transforms.Resize((new_h,new_w), interpolation = 0)])

        resizedImg = imageT(image)
        resizedMsk = maskT(mask)

        return {'image': resizedImg, 'mask': resizedMsk}


class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image, mask = sample['image'], sample['mask']

        imageT = transforms.Compose([transforms.ToTensor()])
        tensorImage = imageT(image)

        return {'image': tensorImage, 'mask': torch.from_numpy(np.array(mask)).long()}


class Normalize(object):
    """Convert pixel values (Tensors) to the range 0-1"""

    def __call__(self, sample):
        image, mask = sample['image'], sample['mask']
        
        imageT = transforms.Compose([transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                                          std =[0.229, 0.224, 0.225],
                                                          inplace=False)])
        normImage = imageT(image)
        
        return {'image': normImage, 'mask': mask}

composedTransformList = transforms.Compose([Rescale((576,378)),
                                            ToTensor(),
                                            Normalize()
                                           ])


from torchvision.models.segmentation.deeplabv3 import DeepLabHead

def createDeepLabv3(num_classes):
    model = models.segmentation.deeplabv3_resnet101(pretrained=True, progress=True)
    # Added a Sigmoid activation after the last convolution layer
    model.classifier = DeepLabHead(2048, num_classes)
    return model


def train_model(model, criterion, optimizer, dataloaders, device, scheduler, num_epochs=25, print_freq=1):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_loss = 1e15

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        loss_history = {'train': [], 'val': []}
        
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            # Iterate over data.
            for sample in tqdm(iter(dataloaders[phase])):
                imgs = sample['image'].to(device)
                msks = sample['mask'].to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(imgs)
                    loss = criterion(outputs['out'], msks)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

            if phase == 'train':
                scheduler.step()

            epoch_loss = np.float(loss.data)
            if (epoch + 1) % print_freq == 0:
                print('Epoch: [%d/%d], Loss: %.4f' %(epoch+1, num_epochs, epoch_loss))
                loss_history[phase].append(epoch_loss)

            # deep copy the model
            if phase == 'val' and epoch_loss < best_loss:
                best_loss = epoch_loss
                best_model_wts = copy.deepcopy(model.state_dict())


    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_loss))

    # load best model weights
    model.load_state_dict(best_model_wts)
    
    return model, loss_history

trainDataset = semanticSegmentationDataset(root_dir='data/train/', transformList=composedTransformList)
valDataset = semanticSegmentationDataset(root_dir='data/val/', transformList=composedTransformList)
testDataset = semanticSegmentationDataset(root_dir='data/test/', transformList=composedTransformList)

trainDataloader = DataLoader(trainDataset, batch_size=4, shuffle=True, num_workers=4)
valDataloader = DataLoader(valDataset, batch_size=4, shuffle=True, num_workers=4)

dataloaders = {'train': trainDataloader, 'val': valDataloader}

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

num_classes = 2 # background and wires

# get the model
model = createDeepLabv3(num_classes)

# move model to the right device
model.to(device)

#zerocount = 0
#onecount = 0
#for i in range(len(trainDataset)):
#    dat = trainDataset[i]
#    counts = dat['mask'].unique(return_counts=True)[1]
#    zerocount += counts[0]
#    onecount += counts[1]

weightForZeroClass = np.float((1/np.float(zerocount)))*0.5 # Also tried these values
weightForOneClass = np.float((1/np.float(onecount)))*0.5 # Also tried these values

criterion = nn.NLLLoss(reduction='mean', weight=torch.tensor([0.01, 1]).to(device)) # Tried these values randomly

# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]

#optimizer = torch.optim.SGD(params, lr=0.005,momentum=0.9, weight_decay=0.0005)
optimizer = optim.Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0005, amsgrad=False)

# and a learning rate scheduler which decreases the learning rate by
# 10x every 3 epochs
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=3,gamma=0.1)

total_epoch = 20

finalTrainedModelSpan, lossDictSpan = train_model(model, criterion, optimizer, dataloaders, device, scheduler, total_epoch)


## Code to predict and display the predicted mask

#i = 3
#model.eval()
#with torch.no_grad():
#    preds = model(testDataset[i]['image'].reshape(1,3,576,378).to(device))
#
#displayPreds = torch.argmax(preds['out'], 1)
#displayPreds.unique(return_counts=True)
#predMaskToDisplay = displayPreds.cpu().numpy().reshape(576,378)
#actualMaskToDisplay = testDataset[i]['mask'].numpy()
#
#plt.imshow(actualMaskToDisplay)
#plt.show()
#plt.imshow(predMaskToDisplay)