
%%capture
!pip install segmentation-models-pytorch
!pip install -U git+https://github.com/albumentations-team/albumentations
!pip install --upgrade opencv-contrib-python


!git clone https://github.com/parth1620/Human-Segmentation-Dataset-master.git

fatal: destination path 'Human-Segmentation-Dataset-master' already exists and is not an empty directory.


import sys
sys.path.append('/content/Human-Segmentation-Dataset-master')


import torch
import cv2

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from tqdm import tqdm
from skimage import io
import helper

from torch.utils.data import Dataset , DataLoader
from torchvision.transforms import v2
from torch import nn
import random
import os


df = pd.read_csv('/kaggle/working/Human-Segmentation-Dataset-master/train.csv')

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DATA_DIR = "/kaggle/working/"
IMG_SIZE = 224
class CFG : 
    epcohs = 100
    seed = 42
    batch_size = 16
    encoder = 'efficientnet-b2'
    encoder_weights = "imagenet"


def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(seed=CFG.seed)


df.head()


df.shape

(290, 2)


train_df , val_df = train_test_split(df , test_size= 0.2 , random_state= CFG.seed)


def show_data(idx) : 
    image , mask = DATA_DIR + df.iloc[idx] 

    image = io.imread(image)
    mask = io.imread(mask)

    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10,5))

    ax1.set_title('IMAGE')
    ax1.imshow(image)

    ax2.set_title('GROUND TRUTH')
    ax2.imshow(mask,cmap = 'gray')


show_data(99)


import albumentations as A
def get_train_augs() :
    return A.Compose([A.Resize(height=IMG_SIZE , width=IMG_SIZE) ,
                       A.Rotate(30) , 
                        A.HorizontalFlip(0.5) , 
                        A.VerticalFlip(0.5)] , is_check_shapes=False)


def get_val_augs() :
    return A.Compose([A.Resize(height=IMG_SIZE , width=IMG_SIZE) ,
                      ], is_check_shapes=False)


class Human_Segmentation_set(Dataset):
    def __init__(self , annotations_file , img_dir , transform = None , target_transform = None) :
        self.img_labels = annotations_file
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform
        
    def __len__(self) : 
        return len(self.img_labels)
    
    def __getitem__(self , idx) :
        
        mask_path , image_path  = self.img_dir + self.img_labels.iloc[idx] 
#         print(image.shape)
        # io.imread is not giving us the color_channels so we will preceed to reading images with cv2.imread()
        image = cv2.imread(image_path)
        mask = cv2.imread(mask_path)
        
        
        if self.transform : 
            data = self.transform(image = image , mask = mask)
            image = data['image']
            mask = data['mask']
#         print(image.shape)
        image_tensor = torch.from_numpy(image).permute(2,0,1)/255.
        mask_tensor = torch.from_numpy(mask).permute(2,0,1)/255.
        
        image_tensor = image_tensor.to(DEVICE)
        mask_tensor = mask_tensor.to(DEVICE)
        
        return image_tensor , mask_tensor
    
    
    def show_data(self , idx) : 
        mask_path , image_path = self.img_dir + self.img_labels.iloc[idx] 

        image = cv2.imread(image_path)
        mask = cv2.imread(mask_path)
        
        
        print(image.shape , mask.shape)
        if self.transform : 
            data = self.transform(image = image , mask = mask)
            image = data['image']
            mask = data['mask']
        f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10,5))

        ax1.set_title('IMAGE')
        ax1.imshow(image)

        ax2.set_title('GROUND TRUTH')
        ax2.imshow(mask,cmap = 'gray')


train_set = Human_Segmentation_set(train_df , DATA_DIR , get_train_augs())
val_set =  Human_Segmentation_set(val_df , DATA_DIR , get_val_augs())


train_set.show_data(9)

(640, 480, 3) (577, 433, 3)


print(f"Size of Trainset : {len(train_set)}")
print(f"Size of Validset : {len(val_set)}")

Size of Trainset : 232
Size of Validset : 58


train_loader = DataLoader(train_set , batch_size=CFG.batch_size , shuffle = True )
val_loader = DataLoader(val_set , batch_size=CFG.batch_size , shuffle = False )


# for image , mask in train_loader : 
#     break
    
# print(image.shape , mask.shape)


import segmentation_models_pytorch as smp
from segmentation_models_pytorch.losses import DiceLoss


class Human_seg_model(nn.Module):
    def __init__(self):
        super(Human_seg_model, self).__init__()
        self.model = smp.Unet(
            encoder_name=CFG.encoder,
            encoder_weights=CFG.encoder_weights,
            in_channels=3,
            classes=3,
            activation=None
        ).to(DEVICE) 
        
    def forward(self, images, masks=None):
        logits = self.model(images)
#         print(images.shape)
#         print(logits.shape)
#         print(masks.shape)
        if masks is not None:
            loss1 = DiceLoss(mode='binary')(logits, masks) 
            loss2 = nn.BCEWithLogitsLoss()(logits, masks)   
            return logits, loss1 + loss2
        
        return logits


def train_fn(model , train_loader , optimizer) : 
    model.train()
    total_loss = 0
    for images , masks in tqdm(train_loader) : 
        optimizer.zero_grad()
        
        images , masks = images.to(DEVICE) , masks.to(DEVICE)
        logits , loss = model(images , masks)
        
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
        
    return total_loss/len(train_loader)


def val_fn(model , val_loader) : 
    model.eval()
    total_loss = 0
    with torch.inference_mode() : 
        for image , mask in tqdm(val_loader) : 
            logits = model(image)
            loss = DiceLoss(mode = 'binary')(logits , mask)
            total_loss+=loss
        
    return total_loss/len(val_loader)


SegModel = Human_seg_model().to(DEVICE)
optimizer = torch.optim.Adam(SegModel.parameters())


for i in range(CFG.epcohs) : 
    train_loss = train_fn(SegModel , train_loader ,optimizer)
    val_loss = val_fn(SegModel , val_loader )

    print(f"EPOCH {i+1} : Train Loss : {train_loss} | Val LOSS : {val_loss}")

100%|██████████| 15/15 [00:03<00:00,  3.98it/s]
100%|██████████| 4/4 [00:00<00:00,  7.95it/s]

EPOCH 1 : Train Loss : 0.9997095545132955 | Val LOSS : 0.25976723432540894

100%|██████████| 15/15 [00:03<00:00,  4.17it/s]
100%|██████████| 4/4 [00:00<00:00,  9.41it/s]

EPOCH 2 : Train Loss : 0.5496830582618714 | Val LOSS : 0.20350058376789093

100%|██████████| 15/15 [00:03<00:00,  4.17it/s]
100%|██████████| 4/4 [00:00<00:00,  9.07it/s]

EPOCH 3 : Train Loss : 0.3670327146848043 | Val LOSS : 0.20375806093215942

100%|██████████| 15/15 [00:03<00:00,  4.19it/s]
100%|██████████| 4/4 [00:00<00:00,  9.55it/s]

EPOCH 4 : Train Loss : 0.2536358783642451 | Val LOSS : 0.2708975076675415

100%|██████████| 15/15 [00:03<00:00,  4.15it/s]
100%|██████████| 4/4 [00:00<00:00,  9.56it/s]

EPOCH 5 : Train Loss : 0.25740688145160673 | Val LOSS : 0.1315029114484787


val_set

<__main__.Human_Segmentation_set at 0x7e8326754670>


prediction.shape

(224, 224, 3)


def show_image(image,mask,pred_image = None):
    
    if pred_image.all() == None:
        
        f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10,5))
        
        ax1.set_title('IMAGE')
        ax1.imshow(image,cmap = 'gray')
        
        ax2.set_title('GROUND TRUTH')
        ax2.imshow(mas,cmap = 'gray')
        
    elif pred_image.all() != None :
        
        f, (ax1, ax2,ax3) = plt.subplots(1, 3, figsize=(10,5))
        
        ax1.set_title('IMAGE')
        ax1.imshow(image,cmap = 'gray')
        
        ax2.set_title('GROUND TRUTH')
        ax2.imshow(mask,cmap = 'gray')
        
        ax3.set_title('MODEL OUTPUT')
        ax3.imshow(pred_image,cmap = 'gray')


with torch.inference_mode() : 
    for image , mask in tqdm(val_set) : 

        image.to(DEVICE)
        
        logits = SegModel(image.unsqueeze(0))
        prediction = logits.squeeze().permute(1,2,0).detach().cpu().numpy()
        mask = mask.permute(1,2,0).detach().cpu().numpy()
        image = image.permute(1,2,0).detach().cpu().numpy()
        
        show_image(image ,mask , prediction)

 34%|███▍      | 20/58 [00:01<00:03, 12.53it/s]/tmp/ipykernel_34/3939027367.py:15: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`). Consider using `matplotlib.pyplot.close()`.
  f, (ax1, ax2,ax3) = plt.subplots(1, 3, figsize=(10,5))
100%|██████████| 58/58 [00:05<00:00, 10.31it/s]

	masks	images
0	Human-Segmentation-Dataset-master/Ground_Truth...	Human-Segmentation-Dataset-master/Training_Ima...
1	Human-Segmentation-Dataset-master/Ground_Truth...	Human-Segmentation-Dataset-master/Training_Ima...
2	Human-Segmentation-Dataset-master/Ground_Truth...	Human-Segmentation-Dataset-master/Training_Ima...
3	Human-Segmentation-Dataset-master/Ground_Truth...	Human-Segmentation-Dataset-master/Training_Ima...
4	Human-Segmentation-Dataset-master/Ground_Truth...	Human-Segmentation-Dataset-master/Training_Ima...

Task 1 : Set up colab gpu runtime environment¶

Download Dataset¶

Some Common Imports¶

Task : 2 Setup Configurations¶

Task 3 : Augmentation Functions¶

Task 4 : Create Custom Dataset¶

Task 5 : Load dataset into batches¶

Task 6 : Create Segmentation Model¶

Task 7 : Create Train and Validation Function¶

Task 8 : Train Model¶

Task 9 : Inference¶