mkdir segmentation_tutorial
cd segmentation_tutorial
virtualenv venv
source venv/bin/activate

pip install albumentations # Библиотека для работы с аугментациями
pip install -U catalyst # Библиотека для постороения автоматизированных пайплайнов
pip install segmentation-models-pytorch # Библиотека в которой содержатся архитектуры и веса моделей нейронных сетей для решения задач сегментации
pip install ttach библиотека с оберткой для test-time аргументаций

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

download-gdrive 1iYaNijLmzsrMlAdMoUEhhJuo-5bkeAuj segmentation_data.zip
mkdir data
unzip segmentation_data.zip &>/dev/null 
mv segmentation_data/* data
rm segmentation_data.zip

bash download_dataset.sh

prepare_dataset.py
import pathlib
import pandas as pd
from sklearn import model_selection
def main(
    datapath: str = "./data", valid_size: float = 0.2, random_state: int = 42
):
    datapath = pathlib.Path(datapath)
    dataframe = pd.DataFrame({
        "image": sorted((datapath / "train").glob("*.jpg")),
        "mask": sorted((datapath / "train_masks").glob("*.gif"))
    })
    df_train, df_valid = model_selection.train_test_split(
        dataframe,
        test_size=valid_size,
        random_state=random_state,
        shuffle=False
    )
    for source, mode in zip((df_train, df_valid), ("train", "valid")):
        source.to_csv(datapath / f"dataset_{mode}.csv", index=False)
if __name__ == "__main__":

python prepare_dataset.py --datapath=./data

from .experiment import Experiment
from catalyst.dl import SupervisedRunner as Runner

from typing import Callable, Dict, List
import numpy as np
import imageio
from torch.utils.data import Dataset
class SegmentationDataset(Dataset):
    def __init__(self, list_data: List[Dict], dict_transform: Callable = None):
        self.data = list_data
        self.dict_transform = dict_transform
    def __getitem__(self, index: int) -> Dict:
        dict_ = self.data[index]
        dict_ = {
            "image": np.asarray(imageio.imread(dict_["image"], pilmode="RGB")),
            "mask": np.asarray(imageio.imread(dict_["mask"]))}
        if self.dict_transform is not None:
            dict_ = self.dict_transform(**dict_)
        return dict_
    def __len__(self) -> int:
        return len(self.data)

import collections
import pandas as pd
from catalyst.dl import ConfigExperiment
from .dataset import SegmentationDataset
from .transforms import (
    pre_transforms, post_transforms, hard_transform, Compose
)
class Experiment(ConfigExperiment):
    @staticmethod
    def get_transforms(
        stage: str = None, mode: str = None, image_size: int = 256,
    ):
        pre_transform_fn = pre_transforms(image_size=image_size)
        if mode == "train":
            post_transform_fn = Compose([
                hard_transform(image_size=image_size),
                post_transforms()
            ])
        elif mode in ["valid", "infer"]:
            post_transform_fn = post_transforms()
        else:
            raise NotImplementedError()
        
        transform_fn = Compose([pre_transform_fn, post_transform_fn])
        return transform_fn
    def get_datasets(
        self,
        stage: str,
        in_csv_train: str = None,
        in_csv_valid: str = None,
        image_size: int = 256,
    ):
        datasets = collections.OrderedDict()
        for source, mode in zip(
            (in_csv_train, in_csv_valid), ("train", "valid")
        ):
            dataframe = pd.read_csv(source)
            datasets[mode] = SegmentationDataset(
                dataframe.to_dict('records'),
                dict_transform=self.get_transforms(
                    stage=stage, mode=mode, image_size=image_size
                ),
            )
        return datasets

import cv2
from albumentations import (
    Compose, LongestMaxSize, PadIfNeeded, Normalize, ShiftScaleRotate,
    RandomGamma
)
from albumentations.pytorch import ToTensor
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)
def pre_transforms(image_size=256):
    return Compose([
        LongestMaxSize(max_size=image_size),
        PadIfNeeded(image_size, image_size, border_mode=cv2.BORDER_CONSTANT),
    ])
def post_transforms():
    return Compose([Normalize(), ToTensor()])
def hard_transform(image_size=256):
    return Compose([
        ShiftScaleRotate(
            shift_limit=0.1,
            scale_limit=0.1,
            rotate_limit=15,
            border_mode=cv2.BORDER_REFLECT,
            p=0.5
        ),
    ])

import torch.nn as nn
import torch.nn.functional as F
from catalyst.dl import reqistry
@registry.Model # Этот декоратор региструет модель 
class SimpleNet(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv1 = nn.Conv2d(3, 6, 5)
    self.pool = nn.MaxPool2D(2, 2)
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(16 * 5 * 5, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)
  def forward(self, x):
    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    x = x.view(-1, 16 * 5 * 5)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
  return x

model_params:                      
  model: ResnetFPNUnet             
  num_classes: 1                   
  arch: resnet18                   
  pretrained: True                 
args:
  expdir: "src"                    
  logdir: "logs" 
runner_params:  # OPTIONAL KEYWORD, params for Runner's init
  input_key: "image"
  output_key: "logits"                
stages:                                    
  state_params:                             
    main_metric: &reduce_metric dice        
    minimize_metric: False
  data_params:                                  
    num_workers: 0                              
    batch_size: 64                              
    in_csv_train: "./data/dataset_train.csv"   
    in_csv_valid: "./data/dataset_valid.csv"    
    image_size : 256                            
  criterion_params:
    criterion: DiceLoss                                            
  stage1:                                       
    state_params:                               
      num_epochs: 10                           
    optimizer_params:                           
      optimizer: Adam                          
      lr: 0.001                                
      weight_decay: 0.0003                      
    scheduler_params:                           
      scheduler: MultiStepLR                    
      milestones: [10]                          
      gamma: 0.3                                
    callbacks_params:                           
      loss_dice:                               
        callback: CriterionCallback             
        input_key: mask                         
        output_key: logits
        prefix: &loss loss_dice                  
      
      accuracy:                                 
        callback: DiceCallback                  
        input_key: mask
        output_key: logits
      optimizer:                           
        callback: OptimizerCallback
        loss_key: *loss
      scheduler:                              
        callback: SchedulerCallback
        reduce_metric: *reduce_metric
      saver:                                 
        callback: CheckpointCallback

catalyst-dl run --config=./config.yml --verbose

tensorboard --logdir ./logs

catalyst-dl trace ./logs

import numpy as np
import cv2
import torch.nn as nn
import torch 
import os
from catalyst.dl import utils
from transforms import (pre_transforms, post_transforms, Compose)
class Predictor:
    def __init__(self, model_path,image_size):
        self.augmentation =Compose([pre_transforms(image_size=image_size),post_transforms()])
        self.m = nn.Sigmoid()
        self.model = utils.load_traced_model(model_path).cuda()
        device_name = "cuda:0" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device_name)
        self.image_size = image_size
    def predict(self, image,threshold):
        augmented = self.augmentation(image=image)
        inputs = augmented['image']
        inputs = inputs.unsqueeze(0).to(self.device) 
        output = self.m(self.model(inputs)[0][0]).cpu().numpy()
        probability = (output > threshold).astype(np.uint8)
        return probability

from inferens import Predictor
import argparse
import pandas as pd
import os
import cv2
import numpy as np
def Car_Segmenatation(path_in, path_out, path_model,threshold,image_size):
    mask_predictor = RemoveBackground(path_model,image_size)
    data = pd.DataFrame({"image":os.listdir(path_in)})
    for image_file in data.iloc[:,0]:
        image = cv2.resize(cv2.imread(os.path.join(path_in,image_file)), (image_size,image_size))
        probability = mask_predictor.predict(image, threshold=0.5)
        image[probability==0]=[255,255,255]
        cv2.imwrite(os.path.join(path_out,image_file),image)
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = 'folders to files')
    parser.add_argument('path_in', type=str, help='dir with image')
    parser.add_argument('path_out', type=str, help='output dir')
    parser.add_argument('path_model', type = str, help = 'path_model')
    parser.add_argument('threshold', type = float, help = 'threshold')
    parser.add_argument('image_size', type = int, help = 'image_size')
    args = parser.parse_args()
    Car_Segmenatation(args.path_in,args.path_out,args.path_model,args.threshold,args.image_size)

python CarSegmentation.py path_in,path_out,path_model,treshold, image_size