Upload 44 files

.gitattributes

@@ -33,3 +33,28 @@ 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
+cloth_segmentation/assets/000.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/001.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/002.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/003.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/004.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/005.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/006.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/007.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/008.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/009.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/010.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/011.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/012.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/013.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/014.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/015.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/016.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/017.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/018.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/019.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/020.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/021.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/022.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/023.png filter=lfs diff=lfs merge=lfs -text
+cloth_segmentation/assets/024.png filter=lfs diff=lfs merge=lfs -text

cloth_segmentation/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Levin Dabhi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cloth_segmentation/README.md

@@ -0,0 +1,56 @@
+# Clothes Segmentation using U2NET #
+
+![Python 3.8](https://img.shields.io/badge/python-3.8-green.svg)
+[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1EhEy3uQh-5oOSagUotVOJAf8m7Vqn0D6?usp=sharing)
+
+This repo contains training code, inference code and pre-trained model for Cloths Parsing from human portrait.</br>
+Here clothes are parsed into 3 category: Upper body(red), Lower body(green) and Full body(yellow)
+
+![Sample 000](assets/000.png)
+![Sample 024](assets/024.png)
+![Sample 018](assets/018.png)
+
+This model works well with any background and almost all poses. For more samples visit [samples.md](samples.md)
+
+# Techinal details
+
+* **U2NET** : This project uses an amazing [U2NET](https://arxiv.org/abs/2005.09007) as a deep learning model. Instead of having 1 channel output from u2net for typical salient object detection task it outputs 4 channels each respresting upper body cloth, lower body cloth, fully body cloth and background. Only categorical cross-entropy loss is used for a given version of the checkpoint.
+
+* **Dataset** : U2net is trained on 45k images [iMaterialist (Fashion) 2019 at FGVC6](https://www.kaggle.com/c/imaterialist-fashion-2019-FGVC6/data) dataset. To reduce complexity, I have clubbed the original 42 categories from dataset labels into 3 categories (upper body, lower body and full body). All images are resized into square `¯\_(ツ)_/¯` 768 x 768 px for training. (This experiment was conducted with 768 px but around 384 px will work fine too if one is retraining on another dataset).
+
+# Training 
+
+- For training this project requires,
+<ul>
+    <ul>
+    <li>&nbsp; PyTorch > 1.3.0</li>
+    <li>&nbsp; tensorboardX</li>
+    <li>&nbsp; gdown</li>
+    </ul>
+</ul>
+
+- Download dataset from this [link](https://www.kaggle.com/c/imaterialist-fashion-2019-FGVC6/data), extract all items.
+- Set path of `train` folder which contains training images and `train.csv` which is label csv file in `options/base_options.py`
+- To port original u2net of all layer except last layer please run `python setup_model_weights.py` and it will generate weights after model surgey in `prev_checkpoints` folder.
+- You can explore various options in `options/base_options.py` like checkpoint saving folder, logs folder etc.
+- For single gpu set `distributed = False` in `options/base_options.py`, for multi gpu set it to `True`.
+- For single gpu run `python train.py`
+- For multi gpu run <br>
+&nbsp;`python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=4 --use_env train.py` <br>
+Here command is for single node, 4 gpu. Tested only for single node.
+- You can watch loss graphs and samples in tensorboard by running tensorboard command in log folder.
+
+
+# Testing/Inference
+- Download pretrained model from this [link](https://drive.google.com/file/d/1mhF3yqd7R-Uje092eypktNl-RoZNuiCJ/view?usp=sharing)(165 MB) in `trained_checkpoint` folder.
+- Put input images in `input_images` folder
+- Run `python infer.py` for inference.
+- Output will be saved in `output_images`
+### OR 
+- Inference in colab from here [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1EhEy3uQh-5oOSagUotVOJAf8m7Vqn0D6?usp=sharing)
+
+# Acknowledgements
+- U2net model is from original [u2net repo](https://github.com/xuebinqin/U-2-Net). Thanks to Xuebin Qin for amazing repo.
+- Complete repo follows structure of [Pix2pixHD repo](https://github.com/NVIDIA/pix2pixHD)
+

cloth_segmentation/assets/label_descriptions.json

@@ -0,0 +1,842 @@
+{
+  "info": {
+    "year": 2019,
+    "version": "1.0",
+    "description": "The 2019 FGVC^6 iMaterialist Competition - Fashion track dataset.",
+    "contributor": "iMaterialist Fashion Competition group",
+    "url": "https://github.com/visipedia/imat_comp",
+    "date_created": "2019-04-19 12:38:27.493919"
+  },
+  "categories": [
+    {
+      "id": 0,
+      "name": "shirt, blouse",
+      "supercategory": "upperbody",
+      "level": 2
+    },
+    {
+      "id": 1,
+      "name": "top, t-shirt, sweatshirt",
+      "supercategory": "upperbody",
+      "level": 2
+    },
+    {
+      "id": 2,
+      "name": "sweater",
+      "supercategory": "upperbody",
+      "level": 2
+    },
+    {
+      "id": 3,
+      "name": "cardigan",
+      "supercategory": "upperbody",
+      "level": 2
+    },
+    {
+      "id": 4,
+      "name": "jacket",
+      "supercategory": "upperbody",
+      "level": 2
+    },
+    {
+      "id": 5,
+      "name": "vest",
+      "supercategory": "upperbody",
+      "level": 2
+    },
+    {
+      "id": 6,
+      "name": "pants",
+      "supercategory": "lowerbody",
+      "level": 2
+    },
+    {
+      "id": 7,
+      "name": "shorts",
+      "supercategory": "lowerbody",
+      "level": 2
+    },
+    {
+      "id": 8,
+      "name": "skirt",
+      "supercategory": "lowerbody",
+      "level": 2
+    },
+    {
+      "id": 9,
+      "name": "coat",
+      "supercategory": "wholebody",
+      "level": 2
+    },
+    {
+      "id": 10,
+      "name": "dress",
+      "supercategory": "wholebody",
+      "level": 2
+    },
+    {
+      "id": 11,
+      "name": "jumpsuit",
+      "supercategory": "wholebody",
+      "level": 2
+    },
+    {
+      "id": 12,
+      "name": "cape",
+      "supercategory": "wholebody",
+      "level": 2
+    },
+    {
+      "id": 13,
+      "name": "glasses",
+      "supercategory": "head",
+      "level": 2
+    },
+    {
+      "id": 14,
+      "name": "hat",
+      "supercategory": "head",
+      "level": 2
+    },
+    {
+      "id": 15,
+      "name": "headband, head covering, hair accessory",
+      "supercategory": "head",
+      "level": 2
+    },
+    {
+      "id": 16,
+      "name": "tie",
+      "supercategory": "neck",
+      "level": 2
+    },
+    {
+      "id": 17,
+      "name": "glove",
+      "supercategory": "arms and hands",
+      "level": 2
+    },
+    {
+      "id": 18,
+      "name": "watch",
+      "supercategory": "arms and hands",
+      "level": 2
+    },
+    {
+      "id": 19,
+      "name": "belt",
+      "supercategory": "waist",
+      "level": 2
+    },
+    {
+      "id": 20,
+      "name": "leg warmer",
+      "supercategory": "legs and feet",
+      "level": 2
+    },
+    {
+      "id": 21,
+      "name": "tights, stockings",
+      "supercategory": "legs and feet",
+      "level": 2
+    },
+    {
+      "id": 22,
+      "name": "sock",
+      "supercategory": "legs and feet",
+      "level": 2
+    },
+    {
+      "id": 23,
+      "name": "shoe",
+      "supercategory": "legs and feet",
+      "level": 2
+    },
+    {
+      "id": 24,
+      "name": "bag, wallet",
+      "supercategory": "others",
+      "level": 2
+    },
+    {
+      "id": 25,
+      "name": "scarf",
+      "supercategory": "others",
+      "level": 2
+    },
+    {
+      "id": 26,
+      "name": "umbrella",
+      "supercategory": "others",
+      "level": 2
+    },
+    {
+      "id": 27,
+      "name": "hood",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 28,
+      "name": "collar",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 29,
+      "name": "lapel",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 30,
+      "name": "epaulette",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 31,
+      "name": "sleeve",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 32,
+      "name": "pocket",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 33,
+      "name": "neckline",
+      "supercategory": "garment parts",
+      "level": 2
+    },
+    {
+      "id": 34,
+      "name": "buckle",
+      "supercategory": "closures",
+      "level": 2
+    },
+    {
+      "id": 35,
+      "name": "zipper",
+      "supercategory": "closures",
+      "level": 2
+    },
+    {
+      "id": 36,
+      "name": "applique",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 37,
+      "name": "bead",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 38,
+      "name": "bow",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 39,
+      "name": "flower",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 40,
+      "name": "fringe",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 41,
+      "name": "ribbon",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 42,
+      "name": "rivet",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 43,
+      "name": "ruffle",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 44,
+      "name": "sequin",
+      "supercategory": "decorations",
+      "level": 2
+    },
+    {
+      "id": 45,
+      "name": "tassel",
+      "supercategory": "decorations",
+      "level": 2
+    }
+  ],
+  "attributes": [
+    {
+      "id": 0,
+      "name": "above-the-hip (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 1,
+      "name": "hip (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 2,
+      "name": "micro (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 3,
+      "name": "mini (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 4,
+      "name": "above-the-knee (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 5,
+      "name": "knee (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 6,
+      "name": "below the knee (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 7,
+      "name": "midi",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 8,
+      "name": "maxi (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 9,
+      "name": "floor (length)",
+      "supercategory": "length",
+      "level": 1
+    },
+    {
+      "id": 10,
+      "name": "single breasted",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 11,
+      "name": "double breasted",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 12,
+      "name": "lace up",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 13,
+      "name": "wrapping",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 14,
+      "name": "zip-up",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 15,
+      "name": "fly (opening)",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 16,
+      "name": "buckled (opening)",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 17,
+      "name": "toggled (opening)",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 18,
+      "name": "no opening",
+      "supercategory": "opening type",
+      "level": 1
+    },
+    {
+      "id": 19,
+      "name": "asymmetrical",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 20,
+      "name": "symmetrical",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 21,
+      "name": "peplum",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 22,
+      "name": "circle",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 23,
+      "name": "flare",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 24,
+      "name": "fit and flare",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 25,
+      "name": "trumpet",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 26,
+      "name": "mermaid",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 27,
+      "name": "balloon",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 28,
+      "name": "bell",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 29,
+      "name": "bell bottom",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 30,
+      "name": "bootcut",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 31,
+      "name": "peg",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 32,
+      "name": "pencil",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 33,
+      "name": "straight",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 34,
+      "name": "a-line",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 35,
+      "name": "tent",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 36,
+      "name": "baggy",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 37,
+      "name": "wide leg",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 38,
+      "name": "high low",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 39,
+      "name": "curved (fit)",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 40,
+      "name": "tight (fit)",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 41,
+      "name": "regular (fit)",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 42,
+      "name": "loose (fit)",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 43,
+      "name": "oversized",
+      "supercategory": "silhouette",
+      "level": 1
+    },
+    {
+      "id": 44,
+      "name": "burnout",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 45,
+      "name": "distressed",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 46,
+      "name": "washed",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 47,
+      "name": "embossed",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 48,
+      "name": "frayed",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 49,
+      "name": "printed",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 50,
+      "name": "ruched",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 51,
+      "name": "quilted",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 52,
+      "name": "pleat",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 53,
+      "name": "gathering",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 54,
+      "name": "smocking",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 55,
+      "name": "tiered",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 56,
+      "name": "cutout",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 57,
+      "name": "slit",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 58,
+      "name": "perforated",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 59,
+      "name": "lining",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 60,
+      "name": "no special manufacturing technique",
+      "supercategory": "textile finishing, manufacturing techniques",
+      "level": 1
+    },
+    {
+      "id": 61,
+      "name": "plain (pattern)",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 62,
+      "name": "abstract",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 63,
+      "name": "cartoon",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 64,
+      "name": "letters, numbers",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 65,
+      "name": "camouflage",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 66,
+      "name": "check",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 67,
+      "name": "dot",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 68,
+      "name": "fair isle",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 69,
+      "name": "floral",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 70,
+      "name": "geometric",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 71,
+      "name": "paisley",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 72,
+      "name": "stripe",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 73,
+      "name": "houndstooth (pattern)",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 74,
+      "name": "herringbone (pattern)",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 75,
+      "name": "chevron",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 76,
+      "name": "argyle",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 77,
+      "name": "leopard",
+      "supercategory": "animal",
+      "level": 2
+    },
+    {
+      "id": 78,
+      "name": "snakeskin (pattern)",
+      "supercategory": "animal",
+      "level": 2
+    },
+    {
+      "id": 79,
+      "name": "cheetah",
+      "supercategory": "animal",
+      "level": 2
+    },
+    {
+      "id": 80,
+      "name": "peacock",
+      "supercategory": "animal",
+      "level": 2
+    },
+    {
+      "id": 81,
+      "name": "zebra",
+      "supercategory": "animal",
+      "level": 2
+    },
+    {
+      "id": 82,
+      "name": "giraffe",
+      "supercategory": "animal",
+      "level": 2
+    },
+    {
+      "id": 83,
+      "name": "toile de jouy",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 84,
+      "name": "plant",
+      "supercategory": "textile pattern",
+      "level": 1
+    },
+    {
+      "id": 85,
+      "name": "empire waistline",
+      "supercategory": "waistline",
+      "level": 1
+    },
+    {
+      "id": 86,
+      "name": "dropped waistline",
+      "supercategory": "waistline",
+      "level": 1
+    },
+    {
+      "id": 87,
+      "name": "high waist",
+      "supercategory": "waistline",
+      "level": 1
+    },
+    {
+      "id": 88,
+      "name": "normal waist",
+      "supercategory": "waistline",
+      "level": 1
+    },
+    {
+      "id": 89,
+      "name": "low waist",
+      "supercategory": "waistline",
+      "level": 1
+    },
+    {
+      "id": 90,
+      "name": "basque (wasitline)",
+      "supercategory": "waistline",
+      "level": 1
+    },
+    {
+      "id": 91,
+      "name": "no waistline",
+      "supercategory": "waistline",
+      "level": 1
+    }
+  ]
+}

cloth_segmentation/data/aligned_dataset.py

@@ -0,0 +1,169 @@
+from data.base_dataset import BaseDataset, Rescale_fixed, Normalize_image
+from data.image_folder import make_dataset, make_dataset_test
+
+import os
+import cv2
+import json
+import itertools
+import collections
+from tqdm import tqdm
+
+import pandas as pd
+import numpy as np
+from PIL import Image
+
+import torch
+import torchvision.transforms as transforms
+
+
+class AlignedDataset(BaseDataset):
+    def initialize(self, opt):
+        self.opt = opt
+        self.image_dir = opt.image_folder
+        self.df_path = opt.df_path
+        self.width = opt.fine_width
+        self.height = opt.fine_height
+
+        # for rgb imgs
+
+        transforms_list = []
+        transforms_list += [transforms.ToTensor()]
+        transforms_list += [Normalize_image(opt.mean, opt.std)]
+        self.transform_rgb = transforms.Compose(transforms_list)
+
+        self.df = pd.read_csv(self.df_path)
+        self.image_info = collections.defaultdict(dict)
+        self.df["CategoryId"] = self.df.ClassId.apply(lambda x: str(x).split("_")[0])
+        temp_df = (
+            self.df.groupby("ImageId")["EncodedPixels", "CategoryId"]
+            .agg(lambda x: list(x))
+            .reset_index()
+        )
+        size_df = self.df.groupby("ImageId")["Height", "Width"].mean().reset_index()
+        temp_df = temp_df.merge(size_df, on="ImageId", how="left")
+        for index, row in tqdm(temp_df.iterrows(), total=len(temp_df)):
+            image_id = row["ImageId"]
+            image_path = os.path.join(self.image_dir, image_id)
+            self.image_info[index]["image_id"] = image_id
+            self.image_info[index]["image_path"] = image_path
+            self.image_info[index]["width"] = self.width
+            self.image_info[index]["height"] = self.height
+            self.image_info[index]["labels"] = row["CategoryId"]
+            self.image_info[index]["orig_height"] = row["Height"]
+            self.image_info[index]["orig_width"] = row["Width"]
+            self.image_info[index]["annotations"] = row["EncodedPixels"]
+
+        self.dataset_size = len(self.image_info)
+
+    def __getitem__(self, index):
+        # load images ad masks
+        idx = index
+        img_path = self.image_info[idx]["image_path"]
+        img = Image.open(img_path).convert("RGB")
+        img = img.resize((self.width, self.height), resample=Image.BICUBIC)
+        image_tensor = self.transform_rgb(img)
+
+        info = self.image_info[idx]
+        mask = np.zeros(
+            (len(info["annotations"]), self.width, self.height), dtype=np.uint8
+        )
+        labels = []
+        for m, (annotation, label) in enumerate(
+            zip(info["annotations"], info["labels"])
+        ):
+            sub_mask = self.rle_decode(
+                annotation, (info["orig_height"], info["orig_width"])
+            )
+            sub_mask = Image.fromarray(sub_mask)
+            sub_mask = sub_mask.resize(
+                (self.width, self.height), resample=Image.BICUBIC
+            )
+            mask[m, :, :] = sub_mask
+            labels.append(int(label) + 1)
+
+        num_objs = len(labels)
+        boxes = []
+        new_labels = []
+        new_masks = []
+
+        for i in range(num_objs):
+            try:
+                pos = np.where(mask[i, :, :])
+                xmin = np.min(pos[1])
+                xmax = np.max(pos[1])
+                ymin = np.min(pos[0])
+                ymax = np.max(pos[0])
+                if abs(xmax - xmin) >= 20 and abs(ymax - ymin) >= 20:
+                    boxes.append([xmin, ymin, xmax, ymax])
+                    new_labels.append(labels[i])
+                    new_masks.append(mask[i, :, :])
+            except ValueError:
+                continue
+
+        if len(new_labels) == 0:
+            boxes.append([0, 0, 20, 20])
+            new_labels.append(0)
+            new_masks.append(mask[0, :, :])
+
+        nmx = np.zeros((len(new_masks), self.width, self.height), dtype=np.uint8)
+        for i, n in enumerate(new_masks):
+            nmx[i, :, :] = n
+
+        boxes = torch.as_tensor(boxes, dtype=torch.float32)
+        labels = torch.as_tensor(new_labels, dtype=torch.int64)
+        masks = torch.as_tensor(nmx, dtype=torch.uint8)
+
+        final_label = np.zeros((self.width, self.height), dtype=np.uint8)
+        first_channel = np.zeros((self.width, self.height), dtype=np.uint8)
+        second_channel = np.zeros((self.width, self.height), dtype=np.uint8)
+        third_channel = np.zeros((self.width, self.height), dtype=np.uint8)
+
+        upperbody = [0, 1, 2, 3, 4, 5]
+        lowerbody = [6, 7, 8]
+        wholebody = [9, 10, 11, 12]
+
+        for i in range(len(labels)):
+            if labels[i] in upperbody:
+                first_channel += new_masks[i]
+            elif labels[i] in lowerbody:
+                second_channel += new_masks[i]
+            elif labels[i] in wholebody:
+                third_channel += new_masks[i]
+
+        first_channel = (first_channel > 0).astype("uint8")
+        second_channel = (second_channel > 0).astype("uint8")
+        third_channel = (third_channel > 0).astype("uint8")
+
+        final_label = first_channel + second_channel * 2 + third_channel * 3
+        conflict_mask = (final_label <= 3).astype("uint8")
+        final_label = (conflict_mask) * final_label + (1 - conflict_mask) * 1
+        target_tensor = torch.as_tensor(final_label, dtype=torch.int64)
+
+        return image_tensor, target_tensor
+
+    def __len__(self):
+        return len(self.image_info)
+
+    def name(self):
+        return "AlignedDataset"
+
+    def rle_decode(self, mask_rle, shape):
+        """
+        mask_rle: run-length as string formated: [start0] [length0] [start1] [length1]... in 1d array
+        shape: (height,width) of array to return
+        Returns numpy array according to the shape, 1 - mask, 0 - background
+        """
+        shape = (shape[1], shape[0])
+        s = mask_rle.split()
+        # gets starts & lengths 1d arrays
+        starts, lengths = [np.asarray(x, dtype=int) for x in (s[0::2], s[1::2])]
+        starts -= 1
+        # gets ends 1d array
+        ends = starts + lengths
+        # creates blank mask image 1d array
+        img = np.zeros(shape[0] * shape[1], dtype=np.uint8)
+        # sets mark pixles
+        for lo, hi in zip(starts, ends):
+            img[lo:hi] = 1
+        # reshape as a 2d mask image
+        return img.reshape(shape).T  # Needed to align to RLE direction

cloth_segmentation/data/base_data_loader.py

@@ -0,0 +1,10 @@
+class BaseDataLoader:
+    def __init__(self):
+        pass
+
+    def initialize(self, opt):
+        self.opt = opt
+        pass
+
+    def load_data():
+        return None

cloth_segmentation/data/base_dataset.py

@@ -0,0 +1,189 @@
+import os
+from PIL import Image
+import cv2
+import numpy as np
+import random
+
+import torch
+import torch.utils.data as data
+import torchvision.transforms as transforms
+
+
+class BaseDataset(data.Dataset):
+    def __init__(self):
+        super(BaseDataset, self).__init__()
+
+    def name(self):
+        return "BaseDataset"
+
+    def initialize(self, opt):
+        pass
+
+
+class Rescale_fixed(object):
+    """Rescale the input image into given size.
+
+    Args:
+        (w,h) (tuple): output size or x (int) then resized will be done in (x,x).
+    """
+
+    def __init__(self, output_size):
+        self.output_size = output_size
+
+    def __call__(self, image):
+        return image.resize(self.output_size, Image.BICUBIC)
+
+
+class Rescale_custom(object):
+    """Rescale the input image and target image into randomly selected size with lower bound of min_size arg.
+
+    Args:
+        min_size (int): Minimum desired output size.
+    """
+
+    def __init__(self, min_size, max_size):
+        assert isinstance(min_size, (int, float))
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, sample):
+
+        input_image, target_image = sample["input_image"], sample["target_image"]
+
+        assert input_image.size == target_image.size
+        w, h = input_image.size
+
+        # Randomly select size to resize
+        if min(self.max_size, h, w) > self.min_size:
+            self.output_size = np.random.randint(
+                self.min_size, min(self.max_size, h, w)
+            )
+        else:
+            self.output_size = self.min_size
+
+        # calculate new size by keeping aspect ratio same
+        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
+
+        new_w, new_h = int(new_w), int(new_h)
+        input_image = input_image.resize((new_w, new_h), Image.BICUBIC)
+        target_image = target_image.resize((new_w, new_h), Image.BICUBIC)
+        return {"input_image": input_image, "target_image": target_image}
+
+
+class ToTensor(object):
+    """Convert ndarrays in sample to Tensors."""
+
+    def __init__(self):
+        self.totensor = transforms.ToTensor()
+
+    def __call__(self, sample):
+        input_image, target_image = sample["input_image"], sample["target_image"]
+
+        return {
+            "input_image": self.totensor(input_image),
+            "target_image": self.totensor(target_image),
+        }
+
+
+class RandomCrop_custom(object):
+    """Crop randomly the image in a sample.
+
+    Args:
+        output_size (tuple or int): Desired output size. If int, square crop
+            is made.
+    """
+
+    def __init__(self, output_size):
+        assert isinstance(output_size, (int, tuple))
+        if isinstance(output_size, int):
+            self.output_size = (output_size, output_size)
+        else:
+            assert len(output_size) == 2
+            self.output_size = output_size
+
+        self.randomcrop = transforms.RandomCrop(self.output_size)
+
+    def __call__(self, sample):
+        input_image, target_image = sample["input_image"], sample["target_image"]
+        cropped_imgs = self.randomcrop(torch.cat((input_image, target_image)))
+
+        return {
+            "input_image": cropped_imgs[
+                :3,
+                :,
+            ],
+            "target_image": cropped_imgs[
+                3:,
+                :,
+            ],
+        }
+
+
+class Normalize_custom(object):
+    """Normalize given dict into given mean and standard dev
+
+    Args:
+        mean (tuple or int): Desired mean to substract from dict's tensors
+        std (tuple or int): Desired std to divide from dict's tensors
+    """
+
+    def __init__(self, mean, std):
+        assert isinstance(mean, (float, tuple))
+        if isinstance(mean, float):
+            self.mean = (mean, mean, mean)
+        else:
+            assert len(mean) == 3
+            self.mean = mean
+
+        if isinstance(std, float):
+            self.std = (std, std, std)
+        else:
+            assert len(std) == 3
+            self.std = std
+
+        self.normalize = transforms.Normalize(self.mean, self.std)
+
+    def __call__(self, sample):
+        input_image, target_image = sample["input_image"], sample["target_image"]
+
+        return {
+            "input_image": self.normalize(input_image),
+            "target_image": self.normalize(target_image),
+        }
+
+
+class Normalize_image(object):
+    """Normalize given tensor into given mean and standard dev
+
+    Args:
+        mean (float): Desired mean to substract from tensors
+        std (float): Desired std to divide from tensors
+    """
+
+    def __init__(self, mean, std):
+        assert isinstance(mean, (float))
+        if isinstance(mean, float):
+            self.mean = mean
+
+        if isinstance(std, float):
+            self.std = std
+
+        self.normalize_1 = transforms.Normalize(self.mean, self.std)
+        self.normalize_3 = transforms.Normalize([self.mean] * 3, [self.std] * 3)
+        self.normalize_18 = transforms.Normalize([self.mean] * 18, [self.std] * 18)
+
+    def __call__(self, image_tensor):
+        if image_tensor.shape[0] == 1:
+            return self.normalize_1(image_tensor)
+
+        elif image_tensor.shape[0] == 3:
+            return self.normalize_3(image_tensor)
+
+        elif image_tensor.shape[0] == 18:
+            return self.normalize_18(image_tensor)
+
+        else:
+            assert "Please set proper channels! Normlization implemented only for 1, 3 and 18"

cloth_segmentation/data/custom_dataset_data_loader.py

@@ -0,0 +1,71 @@
+import torch.utils.data
+from data.base_data_loader import BaseDataLoader
+
+
+def CreateDataset(opt):
+    dataset = None
+    from data.aligned_dataset import AlignedDataset
+    dataset = AlignedDataset()
+
+    print("dataset [%s] was created" % (dataset.name()))
+    dataset.initialize(opt)
+    return dataset
+
+
+class CustomDatasetDataLoader(BaseDataLoader):
+    def name(self):
+        return 'CustomDatasetDataLoader'
+
+    def initialize(self, opt):
+        BaseDataLoader.initialize(self, opt)
+        self.dataset = CreateDataset(opt)
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batchSize,
+            sampler=data_sampler(self.dataset,
+                                 not opt.serial_batches, opt.distributed),
+            num_workers=int(opt.nThreads),
+            pin_memory=True)
+
+    def get_loader(self):
+        return self.dataloader
+
+    def __len__(self):
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+
+def data_sampler(dataset, shuffle, distributed):
+    if distributed:
+        return torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)
+
+    if shuffle:
+        return torch.utils.data.RandomSampler(dataset)
+
+    else:
+        return torch.utils.data.SequentialSampler(dataset)
+
+
+def sample_data(loader):
+    while True:
+        for batch in loader:
+            yield batch
+
+
+class CustomTestDataLoader(BaseDataLoader):
+    def name(self):
+        return 'CustomDatasetDataLoader'
+
+    def initialize(self, opt):
+        BaseDataLoader.initialize(self, opt)
+        self.dataset = CreateDataset(opt)
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batchSize,
+            num_workers=int(opt.nThreads),
+            pin_memory=True)
+
+    def get_loader(self):
+        return self.dataloader
+
+    def __len__(self):
+        return min(len(self.dataset), self.opt.max_dataset_size)

cloth_segmentation/data/data_loader.py

@@ -0,0 +1,7 @@
+def CreateDataLoader(opt):
+    from data.custom_dataset_data_loader import CustomDatasetDataLoader
+
+    data_loader = CustomDatasetDataLoader()
+    print(data_loader.name())
+    data_loader.initialize(opt)
+    return data_loader

cloth_segmentation/data/image_folder.py

@@ -0,0 +1,81 @@
+###############################################################################
+# Code from
+# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
+# Modified the original code so that it also loads images from the current
+# directory as well as the subdirectories
+###############################################################################
+import torch.utils.data as data
+from PIL import Image
+import os
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    f = dir.split('/')[-1].split('_')[-1]
+    print(dir, f)
+    dirs = os.listdir(dir)
+    for img in dirs:
+        path = os.path.join(dir, img)
+        images.append(path)
+    return images
+
+
+def make_dataset_test(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    f = dir.split('/')[-1].split('_')[-1]
+    for i in range(len([name for name in os.listdir(dir) if os.path.isfile(os.path.join(dir, name))])):
+        if f == 'label' or f == 'labelref':
+            img = str(i) + '.png'
+        else:
+            img = str(i) + '.jpg'
+        path = os.path.join(dir, img)
+        # print(path)
+        images.append(path)
+    return images
+
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False,
+                 loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise(RuntimeError("Found 0 images in: " + root + "\n"
+                               "Supported image extensions are: " +
+                               ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)

cloth_segmentation/infer.py

@@ -0,0 +1,86 @@
+import os
+
+from tqdm import tqdm
+from PIL import Image
+import numpy as np
+
+import warnings
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+
+from data.base_dataset import Normalize_image
+from utils.saving_utils import load_checkpoint_mgpu
+
+from networks import U2NET
+
+device = "cuda"
+
+image_dir = "input_images"
+result_dir = "output_images"
+checkpoint_path = os.path.join("trained_checkpoint", "cloth_segm_u2net_latest.pth")
+do_palette = True
+
+
+def get_palette(num_cls):
+    """Returns the color map for visualizing the segmentation mask.
+    Args:
+        num_cls: Number of classes
+    Returns:
+        The color map
+    """
+    n = num_cls
+    palette = [0] * (n * 3)
+    for j in range(0, n):
+        lab = j
+        palette[j * 3 + 0] = 0
+        palette[j * 3 + 1] = 0
+        palette[j * 3 + 2] = 0
+        i = 0
+        while lab:
+            palette[j * 3 + 0] |= ((lab >> 0) & 1) << (7 - i)
+            palette[j * 3 + 1] |= ((lab >> 1) & 1) << (7 - i)
+            palette[j * 3 + 2] |= ((lab >> 2) & 1) << (7 - i)
+            i += 1
+            lab >>= 3
+    return palette
+
+
+transforms_list = []
+transforms_list += [transforms.ToTensor()]
+transforms_list += [Normalize_image(0.5, 0.5)]
+transform_rgb = transforms.Compose(transforms_list)
+
+net = U2NET(in_ch=3, out_ch=4)
+net = load_checkpoint_mgpu(net, checkpoint_path)
+net = net.to(device)
+net = net.eval()
+
+palette = get_palette(4)
+
+images_list = sorted(os.listdir(image_dir))
+pbar = tqdm(total=len(images_list))
+for image_name in images_list:
+    img = Image.open(os.path.join(image_dir, image_name)).convert("RGB")
+    image_tensor = transform_rgb(img)
+    image_tensor = torch.unsqueeze(image_tensor, 0)
+
+    output_tensor = net(image_tensor.to(device))
+    output_tensor = F.log_softmax(output_tensor[0], dim=1)
+    output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
+    output_tensor = torch.squeeze(output_tensor, dim=0)
+    output_tensor = torch.squeeze(output_tensor, dim=0)
+    output_arr = output_tensor.cpu().numpy()
+
+    output_img = Image.fromarray(output_arr.astype("uint8"), mode="L")
+    if do_palette:
+        output_img.putpalette(palette)
+    output_img.save(os.path.join(result_dir, image_name[:-3] + "png"))
+
+    pbar.update(1)
+
+pbar.close()

cloth_segmentation/model_surgery.py

@@ -0,0 +1,51 @@
+import os
+import gdown
+import torch
+
+from networks import U2NET
+from utils.saving_utils import save_checkpoint
+
+os.makedirs("prev_checkpoints", exist_ok=True)
+gdown.download(
+    "https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ",
+    "./prev_checkpoints/u2net.pth",
+    quiet=False,
+)
+
+u_net = U2NET(in_ch=3, out_ch=4)
+save_checkpoint(u_net, os.path.join("prev_checkpoints", "u2net_random.pth"))
+
+# u2net.pth contains trained weights
+trained_net_pth = os.path.join("prev_checkpoints", "u2net.pth")
+# u2net_random.pth contains random weights
+custom_net_pth = os.path.join("prev_checkpoints", "u2net_random.pth")
+
+net_state_dict = torch.load(trained_net_pth)
+count = 0
+for k, v in net_state_dict.items():
+    count += 1
+print("Total number of layers in trained model are: {}".format(count))
+
+custom_state_dict = torch.load(custom_net_pth)
+count = 0
+for k, v in custom_state_dict.items():
+    count += 1
+print("Total number of layers in trained model are: {}".format(count))
+
+total_count = 0
+update_count = 0
+for k, v in net_state_dict.items():
+    total_count += 1
+    if custom_state_dict[k].shape == v.shape:
+        update_count += 1
+        custom_state_dict[k] = v
+
+print(
+    "Out of {} layers in custom network, {} layers weights are recovered from trained model".format(
+        total_count, update_count
+    )
+)
+torch.save(
+    custom_state_dict, os.path.join("prev_checkpoints", "cloth_segm_unet_surgery.pth")
+)
+print("cloth_segm_unet_surgery.pth is generated in prev_checkpoints directory!")

cloth_segmentation/networks/__init__.py

@@ -0,0 +1 @@
+from .u2net import U2NET

cloth_segmentation/networks/u2net.py

@@ -0,0 +1,565 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(
+            in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate
+        )
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src, tar):
+
+    src = F.upsample(src, size=tar.shape[2:], mode="bilinear")
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):  # UNet07DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
+        hx6dup = _upsample_like(hx6d, hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6, hx7
+        del hx6d, hx5d, hx3d, hx2d
+        del hx2dup, hx3dup, hx4dup, hx5dup, hx6dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):  # UNet06DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6
+        del hx5d, hx4d, hx3d, hx2d
+        del hx2dup, hx3dup, hx4dup, hx5dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-5 ###
+class RSU5(nn.Module):  # UNet05DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5
+        del hx4d, hx3d, hx2d
+        del hx2dup, hx3dup, hx4dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-4 ###
+class RSU4(nn.Module):  # UNet04DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4
+        del hx3d, hx2d
+        del hx2dup, hx3dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-4F ###
+class RSU4F(nn.Module):  # UNet04FRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4
+        del hx3d, hx2d
+        """
+
+        return hx1d + hxin
+
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NET, self).__init__()
+
+        self.stage1 = RSU7(in_ch, 32, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 32, 128)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(128, 64, 256)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(256, 128, 512)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(512, 256, 512)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(512, 256, 512)
+
+        # decoder
+        self.stage5d = RSU4F(1024, 256, 512)
+        self.stage4d = RSU4(1024, 128, 256)
+        self.stage3d = RSU5(512, 64, 128)
+        self.stage2d = RSU6(256, 32, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
+
+        self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1)
+
+    def forward(self, x):
+
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # -------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6
+        del hx5d, hx4d, hx3d, hx2d, hx1d
+        del hx6up, hx5dup, hx4dup, hx3dup, hx2dup
+        """
+
+        return d0, d1, d2, d3, d4, d5, d6
+
+
+### U^2-Net small ###
+class U2NETP(nn.Module):
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NETP, self).__init__()
+
+        self.stage1 = RSU7(in_ch, 16, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 16, 64)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(64, 16, 64)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(64, 16, 64)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(64, 16, 64)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(64, 16, 64)
+
+        # decoder
+        self.stage5d = RSU4F(128, 16, 64)
+        self.stage4d = RSU4(128, 16, 64)
+        self.stage3d = RSU5(128, 16, 64)
+        self.stage2d = RSU6(128, 16, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(64, out_ch, 3, padding=1)
+
+        self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1)
+
+    def forward(self, x):
+
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # decoder
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6
+        del hx5d, hx4d, hx3d, hx2d, hx1d
+        del hx6up, hx5dup, hx4dup, hx3dup, hx2dup
+        """
+
+        return d0, d1, d2, d3, d4, d5, d6

cloth_segmentation/options/base_options.py

@@ -0,0 +1,38 @@
+import os.path as osp
+import os
+
+
+class parser(object):
+    def __init__(self):
+        self.name = "training_cloth_segm_u2net_exp1"  # Expriment name
+        self.image_folder = "../imaterialist/train/"  # image folder path
+        self.df_path = "../imaterialist/train.csv"  # label csv path
+        self.distributed = False  # True for multi gpu training
+        self.isTrain = True
+
+        self.fine_width = 192 * 4
+        self.fine_height = 192 * 4
+
+        # Mean std params
+        self.mean = 0.5
+        self.std = 0.5
+
+        self.batchSize = 2  # 12
+        self.nThreads = 2  # 3
+        self.max_dataset_size = float("inf")
+
+        self.serial_batches = False
+        self.continue_train = True
+        if self.continue_train:
+            self.unet_checkpoint = "prev_checkpoints/cloth_segm_unet_surgery.pth"
+
+        self.save_freq = 1000
+        self.print_freq = 10
+        self.image_log_freq = 100
+
+        self.iter = 100000
+        self.lr = 0.0002
+        self.clip_grad = 5
+
+        self.logs_dir = osp.join("logs", self.name)
+        self.save_dir = osp.join("results", self.name)

cloth_segmentation/samples.md

@@ -0,0 +1,33 @@
+![Sample 001](assets/001.png)
+![Sample 002](assets/002.png)
+![Sample 003](assets/003.png)
+![Sample 004](assets/004.png)
+![Sample 007](assets/007.png)
+![Sample 008](assets/008.png)
+![Sample 009](assets/009.png)
+![Sample 010](assets/010.png)
+![Sample 011](assets/011.png)
+![Sample 012](assets/012.png)
+![Sample 013](assets/013.png)
+![Sample 014](assets/014.png)
+![Sample 015](assets/015.png)
+![Sample 016](assets/016.png)
+![Sample 017](assets/017.png)
+![Sample 022](assets/022.png)
+![Sample 023](assets/023.png)
+
+## With different poses
+
+This model works well with different kind of poses too.
+![Sample 019](assets/019.png)
+![Sample 021](assets/021.png)
+
+## Limitations
+
+This model doesn't work in the following condition
+- Image containing multiple people
+- Dress or cloth style which are extremely different from trained dataset <br>
+
+![Sample 005](assets/005.png)
+![Sample 006](assets/006.png)
+![Sample 020](assets/020.png)

cloth_segmentation/train.py

@@ -0,0 +1,190 @@
+import os
+import sys
+import time
+import yaml
+import cv2
+import pprint
+import traceback
+import numpy as np
+
+import warnings
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.autograd import Variable
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.cuda.amp import autocast
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from torchvision import models
+
+from data.custom_dataset_data_loader import CustomDatasetDataLoader, sample_data
+
+
+from options.base_options import parser
+from utils.tensorboard_utils import board_add_images
+from utils.saving_utils import save_checkpoints
+from utils.saving_utils import load_checkpoint, load_checkpoint_mgpu
+from utils.distributed import get_world_size, set_seed, synchronize, cleanup
+
+from networks import U2NET
+
+
+def options_printing_saving(opt):
+    os.makedirs(opt.logs_dir, exist_ok=True)
+    os.makedirs(opt.save_dir, exist_ok=True)
+    os.makedirs(os.path.join(opt.save_dir, "images"), exist_ok=True)
+    os.makedirs(os.path.join(opt.save_dir, "checkpoints"), exist_ok=True)
+
+    # Saving options in yml file
+    option_dict = vars(opt)
+    with open(os.path.join(opt.save_dir, "training_options.yml"), "w") as outfile:
+        yaml.dump(option_dict, outfile)
+
+    for key, value in option_dict.items():
+        print(key, value)
+
+
+def training_loop(opt):
+
+    if opt.distributed:
+        local_rank = int(os.environ.get("LOCAL_RANK"))
+        # Unique only on individual node.
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cuda:0")
+        local_rank = 0
+
+    u_net = U2NET(in_ch=3, out_ch=4)
+    if opt.continue_train:
+        u_net = load_checkpoint(u_net, opt.unet_checkpoint)
+    u_net = u_net.to(device)
+    u_net.train()
+
+    if local_rank == 0:
+        with open(os.path.join(opt.save_dir, "networks.txt"), "w") as outfile:
+            print("<----U-2-Net---->", file=outfile)
+            print(u_net, file=outfile)
+
+    if opt.distributed:
+        u_net = nn.parallel.DistributedDataParallel(
+            u_net,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            broadcast_buffers=False,
+        )
+        print("Going super fast with DistributedDataParallel")
+
+    # initialize optimizer
+    optimizer = optim.Adam(
+        u_net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0
+    )
+
+    custom_dataloader = CustomDatasetDataLoader()
+    custom_dataloader.initialize(opt)
+    loader = custom_dataloader.get_loader()
+
+    if local_rank == 0:
+        dataset_size = len(custom_dataloader)
+        print("Total number of images avaliable for training: %d" % dataset_size)
+        writer = SummaryWriter(opt.logs_dir)
+        print("Entering training loop!")
+
+    # loss function
+    weights = np.array([1, 1.5, 1.5, 1.5], dtype=np.float32)
+    weights = torch.from_numpy(weights).to(device)
+    loss_CE = nn.CrossEntropyLoss(weight=weights).to(device)
+
+    pbar = range(opt.iter)
+    get_data = sample_data(loader)
+
+    start_time = time.time()
+    # Main training loop
+    for itr in pbar:
+        data_batch = next(get_data)
+        image, label = data_batch
+        image = Variable(image.to(device))
+        label = label.type(torch.long)
+        label = Variable(label.to(device))
+
+        d0, d1, d2, d3, d4, d5, d6 = u_net(image)
+
+        loss0 = loss_CE(d0, label)
+        loss1 = loss_CE(d1, label)
+        loss2 = loss_CE(d2, label)
+        loss3 = loss_CE(d3, label)
+        loss4 = loss_CE(d4, label)
+        loss5 = loss_CE(d5, label)
+        loss6 = loss_CE(d6, label)
+        del d1, d2, d3, d4, d5, d6
+
+        total_loss = loss0 * 1.5 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
+
+        for param in u_net.parameters():
+            param.grad = None
+
+        total_loss.backward()
+        if opt.clip_grad != 0:
+            nn.utils.clip_grad_norm_(u_net.parameters(), opt.clip_grad)
+        optimizer.step()
+
+        if local_rank == 0:
+            # printing and saving work
+            if itr % opt.print_freq == 0:
+                pprint.pprint(
+                    "[step-{:08d}] [time-{:.3f}] [total_loss-{:.6f}]  [loss0-{:.6f}]".format(
+                        itr, time.time() - start_time, total_loss, loss0
+                    )
+                )
+
+            if itr % opt.image_log_freq == 0:
+                d0 = F.log_softmax(d0, dim=1)
+                d0 = torch.max(d0, dim=1, keepdim=True)[1]
+                visuals = [[image, torch.unsqueeze(label, dim=1) * 85, d0 * 85]]
+                board_add_images(writer, "grid", visuals, itr)
+
+            writer.add_scalar("total_loss", total_loss, itr)
+            writer.add_scalar("loss0", loss0, itr)
+
+            if itr % opt.save_freq == 0:
+                save_checkpoints(opt, itr, u_net)
+
+    print("Training done!")
+    if local_rank == 0:
+        itr += 1
+        save_checkpoints(opt, itr, u_net)
+
+
+if __name__ == "__main__":
+
+    opt = parser()
+
+    if opt.distributed:
+        if int(os.environ.get("LOCAL_RANK")) == 0:
+            options_printing_saving(opt)
+    else:
+        options_printing_saving(opt)
+
+    try:
+        if opt.distributed:
+            print("Initialize Process Group...")
+            torch.distributed.init_process_group(backend="nccl", init_method="env://")
+            synchronize()
+
+        set_seed(1000)
+        training_loop(opt)
+        cleanup(opt.distributed)
+        print("Exiting..............")
+
+    except KeyboardInterrupt:
+        cleanup(opt.distributed)
+
+    except Exception:
+        traceback.print_exc(file=sys.stdout)
+        cleanup(opt.distributed)

cloth_segmentation/utils/distributed.py

@@ -0,0 +1,47 @@
+import os
+import math
+import numpy as np
+import random
+import pickle
+import torch
+from torch import distributed as dist
+from torch.utils.data.sampler import Sampler
+
+
+def set_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True
+    np.random.seed(seed)
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+
+
+def synchronize():
+    if not dist.is_available():
+        return
+
+    if not dist.is_initialized():
+        return
+
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return
+
+    dist.barrier()
+
+
+def cleanup(distributed):
+    if distributed:
+        dist.destroy_process_group()
+
+
+def get_world_size():
+    if not dist.is_available():
+        return 1
+
+    if not dist.is_initialized():
+        return 1
+
+    return dist.get_world_size()

cloth_segmentation/utils/saving_utils.py

@@ -0,0 +1,45 @@
+import os
+import copy
+import cv2
+import numpy as np
+from collections import OrderedDict
+
+import torch
+
+
+def load_checkpoint(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        print("----No checkpoints at given path----")
+        return
+    model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device("cpu")))
+    print("----checkpoints loaded from path: {}----".format(checkpoint_path))
+    return model
+
+
+def load_checkpoint_mgpu(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        print("----No checkpoints at given path----")
+        return
+    model_state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    new_state_dict = OrderedDict()
+    for k, v in model_state_dict.items():
+        name = k[7:]  # remove `module.`
+        new_state_dict[name] = v
+
+    model.load_state_dict(new_state_dict)
+    print("----checkpoints loaded from path: {}----".format(checkpoint_path))
+    return model
+
+
+def save_checkpoint(model, save_path):
+    print(save_path)
+    if not os.path.exists(os.path.dirname(save_path)):
+        os.makedirs(os.path.dirname(save_path))
+    torch.save(model.state_dict(), save_path)
+
+
+def save_checkpoints(opt, itr, net):
+    save_checkpoint(
+        net,
+        os.path.join(opt.save_dir, "checkpoints", "itr_{:08d}_u2net.pth".format(itr)),
+    )

cloth_segmentation/utils/tensorboard_utils.py

@@ -0,0 +1,54 @@
+import os
+import numpy as np
+from PIL import Image
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+
+# Adding image in tensorboardX
+
+
+def tensor_for_board(img_tensor):
+    # map into [0,1]
+    tensor = (img_tensor.clone()+1) * 0.5
+    tensor.cpu().clamp(0, 1)
+
+    if tensor.size(1) == 1:
+        tensor = tensor.repeat(1, 3, 1, 1)
+
+    return tensor
+
+
+def tensor_list_for_board(img_tensors_list):
+    grid_h = len(img_tensors_list)
+    grid_w = max(len(img_tensors) for img_tensors in img_tensors_list)
+
+    batch_size, channel, height, width = tensor_for_board(
+        img_tensors_list[0][0]).size()
+    canvas_h = grid_h * height
+    canvas_w = grid_w * width
+    canvas = torch.FloatTensor(
+        batch_size, channel, canvas_h, canvas_w).fill_(0.5)
+    for i, img_tensors in enumerate(img_tensors_list):
+        for j, img_tensor in enumerate(img_tensors):
+            offset_h = i * height
+            offset_w = j * width
+            tensor = tensor_for_board(img_tensor)
+            canvas[:, :, offset_h: offset_h + height,
+                   offset_w: offset_w + width].copy_(tensor)
+
+    return canvas
+
+
+def board_add_image(board, tag_name, img_tensor, step_count):
+    tensor = tensor_for_board(img_tensor)
+
+    for i, img in enumerate(tensor):
+        board.add_image('%s/%03d' % (tag_name, i), img, step_count)
+
+
+def board_add_images(board, tag_name, img_tensors_list, step_count):
+    tensor = tensor_list_for_board(img_tensors_list)
+
+    for i, img in enumerate(tensor):
+        board.add_image('%s/%03d' % (tag_name, i), img, step_count)