initial commit

.gitignore

@@ -0,0 +1,141 @@
+# 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/
+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/
+cover/
+
+# 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
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# ignore tf models
+models/*

custom/core.py

@@ -0,0 +1,168 @@
+import numpy as np
+import tensorflow as tf
+
+
+def tflite_inference(inputs, model_path, dtype=np.float32):
+
+    if not isinstance(inputs, (list, tuple)):
+        inputs = (inputs,)
+
+    # Load the TFLite model and allocate tensors.
+    interpreter = tf.lite.Interpreter(model_path=model_path)
+    interpreter.allocate_tensors()
+
+    # Get input and output tensors.
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+
+    # Test the model on random input data.
+    for inp, inp_det in zip(inputs, input_details):
+        interpreter.set_tensor(inp_det["index"], np.array(inp[None, ...], dtype=dtype))
+
+    interpreter.invoke()
+
+    # The function `get_tensor()` returns a copy of the tensor data.
+    # Use `tensor()` in order to get a pointer to the tensor.
+    outputs = [interpreter.get_tensor(out["index"]) for out in output_details]
+
+    return outputs
+
+
+def landmarks_to_detections(landmarks):
+    """
+    landmarks: (3, N) landmarks
+    """
+    x_min = np.amin(landmarks[0, :])
+    x_max = np.amax(landmarks[0, :])
+    y_min = np.amin(landmarks[1, :])
+    y_max = np.amax(landmarks[1, :])
+
+    bbox = dict()
+    bbox["x_min"] = x_min
+    bbox["y_min"] = y_min
+    bbox["width"] = x_max - x_min
+    bbox["height"] = y_max - y_min
+
+    detections = dict()
+    detections["bboxs"] = bbox
+    detections["keypoints"] = landmarks[:2, :]
+
+    return detections
+
+
+def detections_to_rect(
+    detections,
+    image_size,
+    rotation_vector_start_end=None,
+    rotation_vector_target_angle=0,
+):
+
+    keypoints = detections["keypoints"]
+    x_min = np.amin(keypoints[0, :])
+    x_max = np.amax(keypoints[0, :])
+    y_min = np.amin(keypoints[1, :])
+    y_max = np.amax(keypoints[1, :])
+
+    rect = dict()
+    rect["x_center"] = (x_min + x_max) / 2
+    rect["y_center"] = (y_min + y_max) / 2
+    rect["width"] = x_max - x_min
+    rect["height"] = y_max - y_min
+
+    if rotation_vector_start_end is not None:
+        rect["rotation"] = compute_rotation(
+            detections,
+            image_size,
+            rotation_vector_start_end,
+            rotation_vector_target_angle,
+        )
+    else:
+        rect["rotation"] = None
+
+    return rect
+
+
+def compute_rotation(detections, image_size, rotation_vector_start_end, target_angle):
+
+    keypoints = detections["keypoints"]
+
+    x0 = keypoints[0, rotation_vector_start_end[0]] * image_size[0]
+    y0 = keypoints[1, rotation_vector_start_end[0]] * image_size[1]
+    x1 = keypoints[0, rotation_vector_start_end[1]] * image_size[0]
+    y1 = keypoints[1, rotation_vector_start_end[1]] * image_size[1]
+
+    rotation = normalize_radians(target_angle - np.arctan2(-(y1 - y0), x1 - x0))
+
+    return rotation
+
+
+def normalize_radians(angle):
+    return angle - 2 * np.pi * np.floor((angle - (-np.pi)) / (2 * np.pi))
+
+
+def transform_rect(
+    rect,
+    image_size,
+    scale_x=1,
+    scale_y=1,
+    shift_x=0,
+    shift_y=0,
+    square_long=True,
+    square_short=False,
+    opt_rotation=None,
+):
+    width = rect["width"]
+    height = rect["height"]
+    rotation = rect["rotation"]
+    image_width = image_size[0]
+    image_height = image_size[1]
+
+    if rotation is not None and opt_rotation is not None:
+        rotation += opt_rotation
+        rotation = normalize_radians(rotation)
+
+    if rotation is None:
+        rect["x_center"] = rect["x_center"] + width * shift_x
+        rect["y_center"] = rect["y_center"] + height * shift_y
+    else:
+        x_shift = (
+            image_width * width * shift_x * np.cos(rotation)
+            - image_height * height * shift_y * np.sin(rotation)
+        ) / image_width
+        y_shift = (
+            image_width * width * shift_x * np.sin(rotation)
+            + image_height * height * shift_y * np.cos(rotation)
+        ) / image_height
+
+        rect["x_center"] = rect["x_center"] + x_shift
+        rect["y_center"] = rect["y_center"] + y_shift
+
+    if square_long:
+        long_side = np.max((width * image_width, height * image_height))
+        width = long_side / image_width
+        height = long_side / image_height
+    elif square_short:
+        short_side = np.min((width * image_width, height * image_height))
+        width = short_side / image_width
+        height = short_side / image_height
+
+    rect["width"] = width * scale_x
+    rect["height"] = height * scale_y
+
+    return rect
+
+
+def slice_from_roi(roi, image_size, horizontal_side=True):
+    if horizontal_side:
+        center = roi["x_center"]
+        norm_side = roi["width"]
+        image_side = image_size[0]
+    else:
+        center = roi["y_center"]
+        norm_side = roi["height"]
+        image_side = image_size[1]
+
+    first_id = int((center - norm_side / 2) * image_side)
+    second_id = int((center + norm_side / 2) * image_side)
+
+    return (first_id, second_id)