add files

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+flagged
+*.pt
+DS_Store

README.md

@@ -9,4 +9,59 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Pytorch Image Segmentation
+
+## This repo contains the code for training a U-Net model for image segmentation on the Human Segmentation Dataset.
+
+<a href="https://colab.research.google.com/github/josebenitezg/Pytorch-Image-Segmentation/blob/main/Learn_PyTorch_ImageSegmentation.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"> 
+</a>
+
+## Usage :nut_and_bolt:
+
+1. Clone this repo
+
+```
+git clone https://github.com/josebenitezg/Pytorch-Image-Segmentation
+```
+
+2. Create a virtual enviroment
+
+```
+python -m venv env
+```
+
+3. Activate virtual enviroment
+
+- for linux
+
+```
+source env/bin/activate
+```
+
+- for windows
+
+```
+env\Scripts\Activate.bat
+```
+
+4. Install requirements
+
+```
+pip install -r requirements.txt
+```
+
+5. Train the model
+
+```
+python train.py
+```
+
+6. Run gradio inference app
+
+```
+python gradio_inference.py
+```
+
+This repo contains dataset files to train a small model.
+
+Dataset Credit : https://github.com/VikramShenoy97/Human-Segmentation-Datasets

model.py

@@ -0,0 +1,30 @@
+from torch import nn
+import segmentation_models_pytorch as smp
+from segmentation_models_pytorch.losses import DiceLoss
+
+ENCODER = 'timm-efficientnet-b0'
+WEIGHTS = 'imagenet'
+
+class SegmentationModel(nn.Module):
+
+  def __init__(self):
+    super(SegmentationModel, self).__init__()
+
+    self.arc = smp.Unet(
+        encoder_name = ENCODER,
+        encoder_weights = WEIGHTS,
+        in_channels = 3,
+        classes = 1,
+        activation = None
+    )
+
+  def forward(self, images, masks = None):
+
+    logits = self.arc(images)
+
+    if masks != None:
+      loss1 = DiceLoss(mode='binary')(logits, masks)
+      loss2 = nn.BCEWithLogitsLoss()(logits, masks)
+      return logits, loss1 + loss2
+    
+    return logits

requirements.txt

@@ -0,0 +1,3 @@
+albumentations==1.3.0
+segmentation-models-pytorch==0.3.2
+opencv-contrib-python

train.py

@@ -0,0 +1,55 @@
+import torch 
+import cv2
+
+import numpy as np 
+import pandas as pd
+from tqdm import tqdm
+import matplotlib.pyplot as plt 
+
+from utils import load_config, get_train_augs, get_valid_augs, train_fn, eval_fn, SegmentationDataset
+from model import SegmentationModel
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+
+# set device for training
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# load config file
+config = load_config()
+
+# load train files in dataframe
+df = pd.read_csv(config['files']['CSV_FILE'])
+
+train_df, valid_df = train_test_split(df, test_size = 0.2, random_state = 42)
+
+trainset = SegmentationDataset(train_df, get_train_augs(config['model']['IMAGE_SIZE']))
+
+validset = SegmentationDataset(valid_df, get_valid_augs(config['model']['IMAGE_SIZE']))
+
+print(f"Size of Trainset : {len(trainset)}")
+print(f"Size of Validset : {len(validset)}")
+
+trainloader = DataLoader(trainset, batch_size=config['model']['BATCH_SIZE'], shuffle = True)
+validloader = DataLoader(validset, batch_size=config['model']['BATCH_SIZE'])
+
+print(f"Total n of batches in trainloader: {len(trainloader)}")
+print(f"Total n of batches in validloader: {len(validloader)}")
+
+
+model = SegmentationModel()
+model.to(DEVICE)
+
+optimizer = torch.optim.Adam(model.parameters(), lr = config['model']['LR'])
+
+best_valid_loss = np.Inf
+
+for i in tqdm(range(config['model']['EPOCHS'])):
+
+  train_loss = train_fn(trainloader, model, optimizer, DEVICE)
+  valid_loss = eval_fn(validloader, model, DEVICE)
+
+  if valid_loss < best_valid_loss:
+    torch.save(model.state_dict(), 'best_model.pt')
+    print('SAVED-MODEL')
+    best_valid_loss = valid_loss
+  print(f"Epoch: {i+1} Train Loss: {train_loss} Valid Loss: {valid_loss}")

utils.py

@@ -0,0 +1,104 @@
+
+import cv2
+import torch
+import yaml
+import numpy as np
+import albumentations as A
+from torch.utils.data import Dataset
+
+
+def get_train_augs(IMAGE_SIZE):
+    
+  return A.Compose([
+      A.Resize(IMAGE_SIZE, IMAGE_SIZE),
+      A.HorizontalFlip(p = 0.5),
+      A.VerticalFlip(p = 0.5)
+  ])
+  
+def get_valid_augs(IMAGE_SIZE):
+    
+  return A.Compose([
+      A.Resize(IMAGE_SIZE, IMAGE_SIZE),
+  ])
+
+def train_fn(data_loader, model, optimizer, DEVICE):
+
+  model.train()
+  total_loss = 0.0
+
+  for images, masks in data_loader:
+
+    images = images.to(DEVICE)
+    masks = masks.to(DEVICE)
+
+    optimizer.zero_grad()
+    logits, loss = model(images, masks)
+    loss.backward()
+    optimizer.step()
+    total_loss += loss.item()
+
+    return total_loss / len(data_loader)
+
+def eval_fn(data_loader, model, DEVICE):
+
+  model.eval()
+  total_loss = 0.0
+  with torch.no_grad():
+    for images, masks in data_loader:
+
+      images = images.to(DEVICE)
+      masks = masks.to(DEVICE)
+
+      logits, loss = model(images, masks)
+
+      total_loss += loss.item()
+
+  return total_loss / len(data_loader)
+
+def load_config():
+    config_file = f'config/config.yaml'
+
+    with open(config_file, 'r') as file:
+        config = yaml.safe_load(file)
+
+    return config
+
+  
+class SegmentationDataset(Dataset):
+  
+  def __init__(self, df, augmentations):
+
+    self.df = df
+    self.augmentations = augmentations
+
+  def __len__(self):
+    return len(self.df)
+
+  def __getitem__(self, idx):
+
+    row = self.df.iloc[idx]
+
+    image_path = row.images
+    mask_path = row.masks
+
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) #(h, w, c)
+    # Resize the mask to the same dimensions as the image
+    mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST) # (h, w)
+    mask = np.expand_dims(mask, axis = -1)
+
+    if self.augmentations:
+      data = self.augmentations(image = image, mask = mask)
+      image = data['image']
+      mask = data['mask']
+
+    # (h, w, c) -> (c, h, w)
+    image = np.transpose(image, (2,0,1)).astype(np.float32)
+    mask = np.transpose(mask, (2,0,1)).astype(np.float32)
+
+    image = torch.Tensor(image) / 255.0
+    mask = torch.round(torch.Tensor(mask) / 255.0)
+
+    return image, mask