Initial full code submit!

.gitignore

@@ -1,129 +1,11 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-pip-wheel-metadata/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-.python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
+*__pycache__*
+checkpoints/*
+tensorboard/*
+data/train/*
+data/test/*
+log/*
+*.sh
+result/*/*
+.idea/
+.idea/*
+*.pyc

body_binary_masking.py

@@ -0,0 +1,87 @@
+"""
+Make updated body shape from updated segmentation
+"""
+
+import os
+import numpy as np
+import cv2
+from PIL import Image
+import sys
+
+
+(cv_major, _, _) = cv2.__version__.split(".")
+if cv_major != '4' and cv_major != '3':
+    print('doesnot support opencv version')
+    sys.exit()
+
+
+# @TODO this is too simple and pixel based algorithm
+def body_detection(image, seg_mask):
+    # binary thresholding by blue ?
+    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+    lower_blue = np.array([0, 0, 120])
+    upper_blue = np.array([180, 38, 255])
+    mask = cv2.inRange(hsv, lower_blue, upper_blue)
+    result = cv2.bitwise_and(image, image, mask=mask)
+
+    # binary threshold by green ?
+    b, g, r = cv2.split(result)
+    filter = g.copy()
+    ret, mask = cv2.threshold(filter, 10, 255, 1)
+
+    # at least original segmentation is FG
+    mask[seg_mask] = 1
+
+    return mask
+
+
+def make_body_mask(data_dir, seg_dir, image_name, mask_name, save_dir=None):
+    print(image_name)
+
+    # define paths
+    img_pth = os.path.join(data_dir, image_name)
+    seg_pth = os.path.join(seg_dir, mask_name)
+
+    mask_path = None
+    if save_dir is not None:
+        mask_path = os.path.join(save_dir, mask_name)
+
+    # Load images
+    img = cv2.imread(img_pth)
+    # segm = Image.open(seg_pth)
+    gray = cv2.imread(seg_pth, cv2.IMREAD_GRAYSCALE)  # the png file should be 1-ch but it is 3 ch ^^;
+    _, seg_mask = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
+
+    body_mask = body_detection(img, seg_mask)
+    body_mask = body_mask + seg_mask
+    body_mask[seg_mask] = 1
+    cv2.imwrite(mask_path, body_mask)
+
+
+def main():
+    # define paths
+
+    root_dir = "D:/Datasets/viton_resize"
+    mask_folder = "image-mask"
+    seg_folder = "image-parse-new"
+
+    data_mode = "train"
+    # data_mode = "test"
+    image_folder = "image"
+
+    image_dir = os.path.join(os.path.join(root_dir, data_mode), image_folder)
+    seg_dir = os.path.join(os.path.join(root_dir, data_mode), seg_folder)
+
+    image_list = os.listdir(image_dir)
+    seg_list = os.listdir(seg_dir)
+
+    mask_dir = os.path.join(os.path.join(root_dir, data_mode), mask_folder)
+    if not os.path.exists(mask_dir):
+        os.makedirs(mask_dir)
+
+    for each in zip(image_list, seg_list):
+        make_body_mask(image_dir, seg_dir, each[0], each[1], mask_dir)
+
+
+if __name__ == '__main__':
+    main()