Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+test.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,19 @@
+# Clothes Segmentation
+This is my implementation of [this project](https://github.com/ternaus/cloths_segmentation)! 
+
+## Dependencies
+-   python >= 3.6 
+-   [pytorch](https://pytorch.org/) >= 1.2
+-   opencv
+-   matplotlib
+-   albumentations, iglovikov_helper_functions, pytorch_lightning, pytorch_toolbelt, segmentation-models-pytorch, tqdm, wandb
+
+## Installation
+1. Download & install cuda 10.2 toolkit [here](https://developer.nvidia.com/cuda-10.2-download-archive?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=1804&target_type=debnetwork)
+2. Download & install anaconda python 3.7 version 
+3. Install Dependencies
+4. Run `main.py`
+
+## A example
+<a href="url"><img src="test.jpg" align="left" style="width=50%" ></a>
+<a href="url"><img src="test.png" align="left" style="width=50%" ></a>

cloths_segmentation/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.2"

cloths_segmentation/configs/2020-10-29.yaml

@@ -0,0 +1,122 @@
+---
+seed: 1984
+
+num_workers: 4
+experiment_name: "2020-10-29"
+
+val_split: 0.2
+
+model:
+  type: segmentation_models_pytorch.Unet
+  encoder_name: timm-efficientnet-b3
+  classes: 1
+  encoder_weights: noisy-student
+
+trainer:
+  type: pytorch_lightning.Trainer
+  gpus: 4
+  max_epochs: 30
+  distributed_backend: ddp
+  progress_bar_refresh_rate: 1
+  benchmark: True
+  precision: 16
+  gradient_clip_val: 5.0
+  num_sanity_val_steps: 2
+  sync_batchnorm: True
+
+
+scheduler:
+  type: torch.optim.lr_scheduler.CosineAnnealingWarmRestarts
+  T_0: 10
+  T_mult: 2
+
+train_parameters:
+  batch_size: 8
+
+checkpoint_callback:
+  type: pytorch_lightning.callbacks.ModelCheckpoint
+  filepath: "2020-10-29"
+  monitor: val_iou
+  verbose: True
+  mode: max
+  save_top_k: -1
+
+val_parameters:
+  batch_size: 2
+
+optimizer:
+  type: adamp.AdamP
+  lr: 0.0001
+
+
+train_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.PadIfNeeded
+        always_apply: False
+        min_height: 800
+        min_width: 800
+        border_mode: 0 # cv2.BORDER_CONSTANT
+        value: 0
+        mask_value: 0
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.RandomCrop
+        always_apply: False
+        height: 512
+        width: 512
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.HorizontalFlip
+        always_apply: False
+        p: 0.5
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225
+
+val_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.PadIfNeeded
+        always_apply: False
+        min_height: 800
+        min_width: 800
+        border_mode: 0 # cv2.BORDER_CONSTANT
+        value: 0
+        mask_value: 0
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225

cloths_segmentation/configs/2020-10-29a.yaml

@@ -0,0 +1,124 @@
+---
+seed: 1984
+
+num_workers: 4
+experiment_name: "2020-10-29a"
+
+val_split: 0.1
+
+resume_from_checkpoint: 2020-10-29/epoch=4.ckpt
+
+model:
+  type: segmentation_models_pytorch.Unet
+  encoder_name: timm-efficientnet-b3
+  classes: 1
+  encoder_weights: noisy-student
+
+trainer:
+  type: pytorch_lightning.Trainer
+  gpus: 4
+  max_epochs: 30
+  distributed_backend: ddp
+  progress_bar_refresh_rate: 1
+  benchmark: True
+  precision: 16
+  gradient_clip_val: 5.0
+  num_sanity_val_steps: 2
+  sync_batchnorm: True
+
+
+scheduler:
+  type: torch.optim.lr_scheduler.CosineAnnealingWarmRestarts
+  T_0: 10
+  T_mult: 2
+
+train_parameters:
+  batch_size: 8
+
+checkpoint_callback:
+  type: pytorch_lightning.callbacks.ModelCheckpoint
+  filepath: "2020-10-29a"
+  monitor: val_iou
+  verbose: True
+  mode: max
+  save_top_k: -1
+
+val_parameters:
+  batch_size: 2
+
+optimizer:
+  type: adamp.AdamP
+  lr: 0.0001
+
+
+train_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.PadIfNeeded
+        always_apply: False
+        min_height: 800
+        min_width: 800
+        border_mode: 0 # cv2.BORDER_CONSTANT
+        value: 0
+        mask_value: 0
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.RandomCrop
+        always_apply: False
+        height: 512
+        width: 512
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.HorizontalFlip
+        always_apply: False
+        p: 0.5
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225
+
+val_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.PadIfNeeded
+        always_apply: False
+        min_height: 800
+        min_width: 800
+        border_mode: 0 # cv2.BORDER_CONSTANT
+        value: 0
+        mask_value: 0
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225

cloths_segmentation/configs/2020-10-30.yaml

@@ -0,0 +1,147 @@
+---
+seed: 1984
+
+num_workers: 4
+experiment_name: "2020-10-30"
+
+val_split: 0.1
+
+model:
+  type: segmentation_models_pytorch.Unet
+  encoder_name: timm-efficientnet-b3
+  classes: 1
+  encoder_weights: noisy-student
+
+trainer:
+  type: pytorch_lightning.Trainer
+  gpus: 4
+  max_epochs: 70
+  distributed_backend: ddp
+  progress_bar_refresh_rate: 1
+  benchmark: True
+  precision: 16
+  gradient_clip_val: 5.0
+  num_sanity_val_steps: 2
+  sync_batchnorm: True
+#  resume_from_checkpoint: 2020-10-30/epoch=67.ckpt
+
+
+scheduler:
+  type: torch.optim.lr_scheduler.CosineAnnealingWarmRestarts
+  T_0: 10
+  T_mult: 2
+
+train_parameters:
+  batch_size: 8
+
+checkpoint_callback:
+  type: pytorch_lightning.callbacks.ModelCheckpoint
+  filepath: "2020-10-30"
+  monitor: val_iou
+  verbose: True
+  mode: max
+  save_top_k: -1
+
+val_parameters:
+  batch_size: 2
+
+optimizer:
+  type: adamp.AdamP
+  lr: 0.0001
+
+
+train_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.PadIfNeeded
+        always_apply: False
+        min_height: 800
+        min_width: 800
+        border_mode: 0 # cv2.BORDER_CONSTANT
+        value: 0
+        mask_value: 0
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.RandomCrop
+        always_apply: False
+        height: 512
+        width: 512
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.HorizontalFlip
+        always_apply: False
+        p: 0.5
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225
+
+val_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.PadIfNeeded
+        always_apply: False
+        min_height: 800
+        min_width: 800
+        border_mode: 0 # cv2.BORDER_CONSTANT
+        value: 0
+        mask_value: 0
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225
+
+test_aug:
+  transform:
+    __class_fullname__: albumentations.core.composition.Compose
+    bbox_params: null
+    keypoint_params: null
+    p: 1
+    transforms:
+      - __class_fullname__: albumentations.augmentations.transforms.LongestMaxSize
+        always_apply: False
+        max_size: 800
+        p: 1
+      - __class_fullname__: albumentations.augmentations.transforms.Normalize
+        always_apply: false
+        max_pixel_value: 255.0
+        mean:
+          - 0.485
+          - 0.456
+          - 0.406
+        p: 1
+        std:
+          - 0.229
+          - 0.224
+          - 0.225

cloths_segmentation/dataloaders.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+from typing import List, Dict, Any, Tuple
+
+import albumentations as albu
+import numpy as np
+import torch
+from iglovikov_helper_functions.utils.image_utils import load_rgb, load_grayscale
+from pytorch_toolbelt.utils.torch_utils import tensor_from_rgb_image
+from torch.utils.data import Dataset
+
+
+class SegmentationDataset(Dataset):
+    def __init__(
+        self,
+        samples: List[Tuple[Path, Path]],
+        transform: albu.Compose,
+        length: int = None,
+    ) -> None:
+        self.samples = samples
+        self.transform = transform
+
+        if length is None:
+            self.length = len(self.samples)
+        else:
+            self.length = length
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        idx = idx % len(self.samples)
+
+        image_path, mask_path = self.samples[idx]
+
+        image = load_rgb(image_path, lib="cv2")
+        mask = load_grayscale(mask_path)
+
+        # apply augmentations
+        sample = self.transform(image=image, mask=mask)
+        image, mask = sample["image"], sample["mask"]
+
+        mask = (mask > 0).astype(np.uint8)
+
+        mask = torch.from_numpy(mask)
+
+        return {
+            "image_id": image_path.stem,
+            "features": tensor_from_rgb_image(image),
+            "masks": torch.unsqueeze(mask, 0).float(),
+        }

cloths_segmentation/inference.py

@@ -0,0 +1,161 @@
+import argparse
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+
+import albumentations as albu
+import cv2
+import numpy as np
+import torch
+import torch.nn.parallel
+import torch.utils.data
+import torch.utils.data.distributed
+import yaml
+from albumentations.core.serialization import from_dict
+from iglovikov_helper_functions.config_parsing.utils import object_from_dict
+from iglovikov_helper_functions.dl.pytorch.utils import state_dict_from_disk, tensor_from_rgb_image
+from iglovikov_helper_functions.utils.image_utils import load_rgb, pad_to_size, unpad_from_size
+from torch.utils.data import Dataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    arg = parser.add_argument
+    arg("-i", "--input_path", type=Path, help="Path with images.", required=True)
+    arg("-c", "--config_path", type=Path, help="Path to config.", required=True)
+    arg("-o", "--output_path", type=Path, help="Path to save masks.", required=True)
+    arg("-b", "--batch_size", type=int, help="batch_size", default=1)
+    arg("-j", "--num_workers", type=int, help="num_workers", default=12)
+    arg("-w", "--weight_path", type=str, help="Path to weights.", required=True)
+    arg("--world_size", default=-1, type=int, help="number of nodes for distributed training")
+    arg("--local_rank", default=-1, type=int, help="node rank for distributed training")
+    arg("--fp16", action="store_true", help="Use fp6")
+    return parser.parse_args()
+
+
+class InferenceDataset(Dataset):
+    def __init__(self, file_paths: List[Path], transform: albu.Compose) -> None:
+        self.file_paths = file_paths
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return len(self.file_paths)
+
+    def __getitem__(self, idx: int) -> Optional[Dict[str, Any]]:
+        image_path = self.file_paths[idx]
+
+        image = load_rgb(image_path)
+        height, width = image.shape[:2]
+
+        image = self.transform(image=image)["image"]
+        pad_dict = pad_to_size((max(image.shape[:2]), max(image.shape[:2])), image)
+
+        return {
+            "torched_image": tensor_from_rgb_image(pad_dict["image"]),
+            "image_path": str(image_path),
+            "pads": pad_dict["pads"],
+            "original_width": width,
+            "original_height": height,
+        }
+
+
+def main():
+    args = get_args()
+    torch.distributed.init_process_group(backend="nccl")
+
+    with open(args.config_path) as f:
+        hparams = yaml.load(f, Loader=yaml.SafeLoader)
+
+    hparams.update(
+        {
+            "local_rank": args.local_rank,
+            "fp16": args.fp16,
+        }
+    )
+
+    output_mask_path = args.output_path
+    output_mask_path.mkdir(parents=True, exist_ok=True)
+    hparams["output_mask_path"] = output_mask_path
+
+    device = torch.device("cuda", args.local_rank)
+
+    model = object_from_dict(hparams["model"])
+    model = model.to(device)
+
+    if args.fp16:
+        model = model.half()
+
+    corrections: Dict[str, str] = {"model.": ""}
+    state_dict = state_dict_from_disk(file_path=args.weight_path, rename_in_layers=corrections)
+    model.load_state_dict(state_dict)
+
+    model = torch.nn.parallel.DistributedDataParallel(
+        model, device_ids=[args.local_rank], output_device=args.local_rank
+    )
+
+    file_paths = []
+
+    for regexp in ["*.jpg", "*.png", "*.jpeg", "*.JPG"]:
+        file_paths += sorted([x for x in tqdm(args.input_path.rglob(regexp))])
+
+    # Filter file paths for which we already have predictions
+    file_paths = [x for x in file_paths if not (args.output_path / x.parent.name / f"{x.stem}.png").exists()]
+
+    dataset = InferenceDataset(file_paths, transform=from_dict(hparams["test_aug"]))
+
+    sampler = DistributedSampler(dataset, shuffle=False)
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        shuffle=False,
+        drop_last=False,
+        sampler=sampler,
+    )
+
+    predict(dataloader, model, hparams, device)
+
+
+def predict(dataloader, model, hparams, device):
+    model.eval()
+
+    if hparams["local_rank"] == 0:
+        loader = tqdm(dataloader)
+    else:
+        loader = dataloader
+
+    with torch.no_grad():
+        for batch in loader:
+            torched_images = batch["torched_image"]  # images that are rescaled and padded
+
+            if hparams["fp16"]:
+                torched_images = torched_images.half()
+
+            image_paths = batch["image_path"]
+            pads = batch["pads"]
+            heights = batch["original_height"]
+            widths = batch["original_width"]
+
+            batch_size = torched_images.shape[0]
+
+            predictions = model(torched_images.to(device))
+
+            for batch_id in range(batch_size):
+                file_id = Path(image_paths[batch_id]).stem
+                folder_name = Path(image_paths[batch_id]).parent.name
+
+                mask = (predictions[batch_id][0].cpu().numpy() > 0).astype(np.uint8) * 255
+                mask = unpad_from_size(pads, image=mask)["image"]
+                mask = cv2.resize(
+                    mask, (widths[batch_id].item(), heights[batch_id].item()), interpolation=cv2.INTER_NEAREST
+                )
+
+                (hparams["output_mask_path"] / folder_name).mkdir(exist_ok=True, parents=True)
+                cv2.imwrite(str(hparams["output_mask_path"] / folder_name / f"{file_id}.png"), mask)
+
+
+if __name__ == "__main__":
+    main()

cloths_segmentation/metrics.py

@@ -0,0 +1,18 @@
+import torch
+
+EPSILON = 1e-15
+
+
+def binary_mean_iou(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+    output = (logits > 0).int()
+
+    if output.shape != targets.shape:
+        targets = torch.squeeze(targets, 1)
+
+    intersection = (targets * output).sum()
+
+    union = targets.sum() + output.sum() - intersection
+
+    result = (intersection + EPSILON) / (union + EPSILON)
+
+    return result

cloths_segmentation/pre_trained_models.py

@@ -0,0 +1,23 @@
+from collections import namedtuple
+from torch import nn
+from torch.utils import model_zoo
+from iglovikov_helper_functions.dl.pytorch.utils import rename_layers
+
+from segmentation_models_pytorch import Unet
+
+model = namedtuple("model", ["url", "model"])
+
+models = {
+    "Unet_2020-10-30": model(
+        url="https://github.com/ternaus/cloths_segmentation/releases/download/0.0.1/weights.zip",
+        model=Unet(encoder_name="timm-efficientnet-b3", classes=1, encoder_weights=None),
+    )
+}
+
+
+def create_model(model_name: str) -> nn.Module:
+    model = models[model_name].model
+    state_dict = model_zoo.load_url(models[model_name].url, progress=True, map_location="cpu")["state_dict"]
+    state_dict = rename_layers(state_dict, {"model.": ""})
+    model.load_state_dict(state_dict)
+    return model

cloths_segmentation/train.py

@@ -0,0 +1,180 @@
+import argparse
+import os
+from pathlib import Path
+from typing import Dict
+
+import pytorch_lightning as pl
+import torch
+import yaml
+from albumentations.core.serialization import from_dict
+from iglovikov_helper_functions.config_parsing.utils import object_from_dict
+from iglovikov_helper_functions.dl.pytorch.lightning import find_average
+from iglovikov_helper_functions.dl.pytorch.utils import state_dict_from_disk
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_toolbelt.losses import JaccardLoss, BinaryFocalLoss
+from torch.utils.data import DataLoader
+
+from cloths_segmentation.dataloaders import SegmentationDataset
+from cloths_segmentation.metrics import binary_mean_iou
+from cloths_segmentation.utils import get_samples
+
+image_path = Path(os.environ["IMAGE_PATH"])
+mask_path = Path(os.environ["MASK_PATH"])
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    arg = parser.add_argument
+    arg("-c", "--config_path", type=Path, help="Path to the config.", required=True)
+    return parser.parse_args()
+
+
+class SegmentPeople(pl.LightningModule):
+    def __init__(self, hparams):
+        super().__init__()
+        self.hparams = hparams
+
+        self.model = object_from_dict(self.hparams["model"])
+        if "resume_from_checkpoint" in self.hparams:
+            corrections: Dict[str, str] = {"model.": ""}
+
+            state_dict = state_dict_from_disk(
+                file_path=self.hparams["resume_from_checkpoint"],
+                rename_in_layers=corrections,
+            )
+            self.model.load_state_dict(state_dict)
+
+        self.losses = [
+            ("jaccard", 0.1, JaccardLoss(mode="binary", from_logits=True)),
+            ("focal", 0.9, BinaryFocalLoss()),
+        ]
+
+    def forward(self, batch: torch.Tensor) -> torch.Tensor:  # type: ignore
+        return self.model(batch)
+
+    def setup(self, stage=0):
+        samples = get_samples(image_path, mask_path)
+
+        num_train = int((1 - self.hparams["val_split"]) * len(samples))
+
+        self.train_samples = samples[:num_train]
+        self.val_samples = samples[num_train:]
+
+        print("Len train samples = ", len(self.train_samples))
+        print("Len val samples = ", len(self.val_samples))
+
+    def train_dataloader(self):
+        train_aug = from_dict(self.hparams["train_aug"])
+
+        if "epoch_length" not in self.hparams["train_parameters"]:
+            epoch_length = None
+        else:
+            epoch_length = self.hparams["train_parameters"]["epoch_length"]
+
+        result = DataLoader(
+            SegmentationDataset(self.train_samples, train_aug, epoch_length),
+            batch_size=self.hparams["train_parameters"]["batch_size"],
+            num_workers=self.hparams["num_workers"],
+            shuffle=True,
+            pin_memory=True,
+            drop_last=True,
+        )
+
+        print("Train dataloader = ", len(result))
+        return result
+
+    def val_dataloader(self):
+        val_aug = from_dict(self.hparams["val_aug"])
+
+        result = DataLoader(
+            SegmentationDataset(self.val_samples, val_aug, length=None),
+            batch_size=self.hparams["val_parameters"]["batch_size"],
+            num_workers=self.hparams["num_workers"],
+            shuffle=False,
+            pin_memory=True,
+            drop_last=False,
+        )
+
+        print("Val dataloader = ", len(result))
+
+        return result
+
+    def configure_optimizers(self):
+        optimizer = object_from_dict(
+            self.hparams["optimizer"],
+            params=[x for x in self.model.parameters() if x.requires_grad],
+        )
+
+        scheduler = object_from_dict(self.hparams["scheduler"], optimizer=optimizer)
+        self.optimizers = [optimizer]
+
+        return self.optimizers, [scheduler]
+
+    def training_step(self, batch, batch_idx):
+        features = batch["features"]
+        masks = batch["masks"]
+
+        logits = self.forward(features)
+
+        total_loss = 0
+        logs = {}
+        for loss_name, weight, loss in self.losses:
+            ls_mask = loss(logits, masks)
+            total_loss += weight * ls_mask
+            logs[f"train_mask_{loss_name}"] = ls_mask
+
+        logs["train_loss"] = total_loss
+
+        logs["lr"] = self._get_current_lr()
+
+        return {"loss": total_loss, "log": logs}
+
+    def _get_current_lr(self) -> torch.Tensor:
+        lr = [x["lr"] for x in self.optimizers[0].param_groups][0]  # type: ignore
+        return torch.Tensor([lr])[0].cuda()
+
+    def validation_step(self, batch, batch_id):
+        features = batch["features"]
+        masks = batch["masks"]
+
+        logits = self.forward(features)
+
+        result = {}
+        for loss_name, _, loss in self.losses:
+            result[f"val_mask_{loss_name}"] = loss(logits, masks)
+
+        result["val_iou"] = binary_mean_iou(logits, masks)
+
+        return result
+
+    def validation_epoch_end(self, outputs):
+        logs = {"epoch": self.trainer.current_epoch}
+
+        avg_val_iou = find_average(outputs, "val_iou")
+
+        logs["val_iou"] = avg_val_iou
+
+        return {"val_iou": avg_val_iou, "log": logs}
+
+
+def main():
+    args = get_args()
+
+    with open(args.config_path) as f:
+        hparams = yaml.load(f, Loader=yaml.SafeLoader)
+
+    pipeline = SegmentPeople(hparams)
+
+    Path(hparams["checkpoint_callback"]["filepath"]).mkdir(exist_ok=True, parents=True)
+
+    trainer = object_from_dict(
+        hparams["trainer"],
+        logger=WandbLogger(hparams["experiment_name"]),
+        checkpoint_callback=object_from_dict(hparams["checkpoint_callback"]),
+    )
+
+    trainer.fit(pipeline)
+
+
+if __name__ == "__main__":
+    main()

cloths_segmentation/utils.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+from typing import Union, Dict, List, Tuple
+
+
+def get_id2_file_paths(path: Union[str, Path]) -> Dict[str, Path]:
+    return {x.stem: x for x in Path(path).glob("*.*")}
+
+
+def get_samples(image_path: Path, mask_path: Path) -> List[Tuple[Path, Path]]:
+    """Couple masks and images.
+
+    Args:
+        image_path:
+        mask_path:
+
+    Returns:
+    """
+
+    image2path = get_id2_file_paths(image_path)
+    mask2path = get_id2_file_paths(mask_path)
+
+    return [(image_file_path, mask2path[file_id]) for file_id, image_file_path in image2path.items()]

rb.py

@@ -0,0 +1,34 @@
+import numpy as np
+import cv2
+import torch
+import albumentations as albu
+from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
+from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
+
+from cloths_segmentation.pre_trained_models import create_model
+model = create_model("Unet_2020-10-30")
+model.eval()
+
+image = cv2.imread(str(r"test.jpg"))
+image_2_extract = image
+image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+transform = albu.Compose([albu.Normalize(p=1)], p=1)
+padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
+x = transform(image=padded_image)["image"]
+x = torch.unsqueeze(tensor_from_rgb_image(x), 0)
+
+with torch.no_grad():
+  prediction = model(x)[0][0]
+  mask = (prediction > 0).cpu().numpy().astype(np.uint8)
+  mask = unpad(mask, pads)
+  rmask = (cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) * 255).astype(np.uint8)
+  mask2 = np.where((rmask < 255), 0, 1).astype('uint8')
+  image_2_extract = image_2_extract * mask2[:, :, 1, np.newaxis]
+
+  tmp = cv2.cvtColor(image_2_extract, cv2.COLOR_BGR2GRAY)
+  _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY)
+  b, g, r = cv2.split(image_2_extract)
+  rgba = [b, g, r, alpha]
+  dst = cv2.merge(rgba, 4)
+  cv2.imwrite("test.png", dst)
+  # cv2.waitKey(0)

semgent_from_folder.py

@@ -0,0 +1,105 @@
+# import os
+# import cv2
+# import numpy as np
+# import torch
+# import albumentations as albu
+# from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
+# from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
+# from cloths_segmentation.pre_trained_models import create_model
+# from tqdm import tqdm
+
+
+# model = create_model("Unet_2020-10-30")
+# model.to("cuda")
+# model.eval()
+
+# input_dir = "../../image"
+# output_dir = "../../seg_masks"
+
+# os.makedirs(output_dir, exist_ok=True)
+
+# for image_filename in tqdm(os.listdir(input_dir), colour="green"):
+#     image_path = os.path.join(input_dir, image_filename)
+
+#     image = cv2.imread(image_path)
+#     image_2_extract = image
+#     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+#     transform = albu.Compose([albu.Normalize(p=1)], p=1)
+#     padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
+#     x = transform(image=padded_image)["image"]
+#     x = torch.unsqueeze(tensor_from_rgb_image(x), 0).to("cuda")
+
+#     with torch.no_grad():
+#         prediction = model(x)[0][0]
+#         mask = (prediction > 0).cpu().numpy().astype(np.uint8)
+#         mask = unpad(mask, pads)
+#         rmask = (cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) * 255).astype(np.uint8)
+#         mask2 = np.where((rmask < 255), 0, 1).astype('uint8')
+#         image_2_extract = image_2_extract * mask2[:, :, 1, np.newaxis]
+
+#         tmp = cv2.cvtColor(image_2_extract, cv2.COLOR_BGR2GRAY)
+#         _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY)
+#         b, g, r = cv2.split(image_2_extract)
+#         rgba = [b, g, r, alpha]
+#         dst = cv2.merge(rgba, 4)
+
+#         output_image_path = os.path.join(output_dir, image_filename.replace(".jpg", ".png"))
+#         cv2.imwrite(output_image_path, dst)
+
+#     # break
+
+
+
+import os
+import cv2
+import numpy as np
+import torch
+import albumentations as albu
+from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
+from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
+from cloths_segmentation.pre_trained_models import create_model
+from tqdm import tqdm
+
+# Create the model and wrap it with DataParallel
+model = create_model("Unet_2020-10-30")
+model = torch.nn.DataParallel(model)
+
+# Move the model to CUDA devices
+model.to("cuda")
+model.eval()
+
+input_dir = "../../image"
+output_dir = "../../seg_masks"
+
+os.makedirs(output_dir, exist_ok=True)
+
+for image_filename in tqdm(os.listdir(input_dir), colour="green"):
+    image_path = os.path.join(input_dir, image_filename)
+
+    image = cv2.imread(image_path)
+    image_2_extract = image
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    transform = albu.Compose([albu.Normalize(p=1)], p=1)
+    padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
+    x = transform(image=padded_image)["image"]
+    x = torch.unsqueeze(tensor_from_rgb_image(x), 0).to("cuda")
+
+    with torch.no_grad():
+        # Use DataParallel to perform inference on all 4 GPUs
+        prediction = model(x)[0][0]
+        mask = (prediction > 0).cpu().numpy().astype(np.uint8)
+        mask = unpad(mask, pads)
+        rmask = (cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) * 255).astype(np.uint8)
+        mask2 = np.where((rmask < 255), 0, 1).astype('uint8')
+        image_2_extract = image_2_extract * mask2[:, :, 1, np.newaxis]
+
+        tmp = cv2.cvtColor(image_2_extract, cv2.COLOR_BGR2GRAY)
+        _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY)
+        b, g, r = cv2.split(image_2_extract)
+        rgba = [b, g, r, alpha]
+        dst = cv2.merge(rgba, 4)
+
+        output_image_path = os.path.join(output_dir, image_filename.replace(".jpg", ".png"))
+        cv2.imwrite(output_image_path, dst)
+
+    # break