Upload cloth_extraction.py

lib/common/cloth_extraction.py

@@ -0,0 +1,170 @@
+import numpy as np
+import json
+import os
+import itertools
+import trimesh
+from matplotlib.path import Path
+from collections import Counter
+from sklearn.neighbors import KNeighborsClassifier
+
+
+def load_segmentation(path, shape):
+    """
+    Get a segmentation mask for a given image
+    Arguments:
+        path: path to the segmentation json file
+        shape: shape of the output mask
+    Returns:
+        Returns a segmentation mask
+    """
+    with open(path) as json_file:
+        dict = json.load(json_file)
+        segmentations = []
+        for key, val in dict.items():
+            if not key.startswith('item'):
+                continue
+
+            # Each item can have multiple polygons. Combine them to one
+            # segmentation_coord = list(itertools.chain.from_iterable(val['segmentation']))
+            # segmentation_coord = np.round(np.array(segmentation_coord)).astype(int)
+
+            coordinates = []
+            for segmentation_coord in val['segmentation']:
+                # The format before is [x1,y1, x2, y2, ....]
+                x = segmentation_coord[::2]
+                y = segmentation_coord[1::2]
+                xy = np.vstack((x, y)).T
+                coordinates.append(xy)
+
+            segmentations.append(
+                {'type': val['category_name'], 'type_id': val['category_id'], 'coordinates': coordinates})
+
+        return segmentations
+
+
+def smpl_to_recon_labels(recon, smpl, k=1):
+    """
+    Get the bodypart labels for the recon object by using the labels from the corresponding smpl object
+    Arguments:
+        recon: trimesh object (fully clothed model)
+        shape: trimesh object (smpl model)
+        k: number of nearest neighbours to use
+    Returns:
+        Returns a dictionary containing the bodypart and the corresponding indices
+    """
+    smpl_vert_segmentation = json.load(
+        open(os.path.join(os.path.dirname(__file__), 'smpl_vert_segmentation.json')))
+    n = smpl.vertices.shape[0]
+    y = np.array([None] * n)
+    for key, val in smpl_vert_segmentation.items():
+        y[val] = key
+
+    classifier = KNeighborsClassifier(n_neighbors=1)
+    classifier.fit(smpl.vertices, y)
+
+    y_pred = classifier.predict(recon.vertices)
+
+    recon_labels = {}
+    for key in smpl_vert_segmentation.keys():
+        recon_labels[key] = list(np.argwhere(
+            y_pred == key).flatten().astype(int))
+
+    return recon_labels
+
+
+def extract_cloth(recon, segmentation, K, R, t, smpl=None):
+    """
+    Extract a portion of a mesh using 2d segmentation coordinates
+    Arguments:
+        recon: fully clothed mesh
+        seg_coord: segmentation coordinates in 2D (NDC)
+        K: intrinsic matrix of the projection
+        R: rotation matrix of the projection
+        t: translation vector of the projection
+    Returns:
+        Returns a submesh using the segmentation coordinates
+    """
+    seg_coord = segmentation['coord_normalized']
+    mesh = trimesh.Trimesh(recon.vertices, recon.faces)
+    extrinsic = np.zeros((3, 4))
+    extrinsic[:3, :3] = R
+    extrinsic[:, 3] = t
+    P = K[:3, :3] @ extrinsic
+
+    P_inv = np.linalg.pinv(P)
+
+    # Each segmentation can contain multiple polygons
+    # We need to check them separately
+    points_so_far = []
+    faces = recon.faces
+    for polygon in seg_coord:
+        n = len(polygon)
+        coords_h = np.hstack((polygon, np.ones((n, 1))))
+        # Apply the inverse projection on homogeneus 2D coordinates to get the corresponding 3d Coordinates
+        XYZ = P_inv @ coords_h[:, :, None]
+        XYZ = XYZ.reshape((XYZ.shape[0], XYZ.shape[1]))
+        XYZ = XYZ[:, :3] / XYZ[:, 3, None]
+
+        p = Path(XYZ[:, :2])
+
+        grid = p.contains_points(recon.vertices[:, :2])
+        indeces = np.argwhere(grid == True)
+        points_so_far += list(indeces.flatten())
+
+    if smpl is not None:
+        num_verts = recon.vertices.shape[0]
+        recon_labels = smpl_to_recon_labels(recon, smpl)
+        body_parts_to_remove = ['rightHand', 'leftToeBase', 'leftFoot', 'rightFoot', 'head',
+                                'leftHandIndex1', 'rightHandIndex1', 'rightToeBase', 'leftHand', 'rightHand']
+        type = segmentation['type_id']
+
+        # Remove additional bodyparts that are most likely not part of the segmentation but might intersect (e.g. hand in front of torso)
+        # https://github.com/switchablenorms/DeepFashion2
+        # Short sleeve clothes
+        if type == 1 or type == 3 or type == 10:
+            body_parts_to_remove += ['leftForeArm', 'rightForeArm']
+        # No sleeves at all or lower body clothes
+        elif type == 5 or type == 6 or type == 12 or type == 13 or type == 8 or type == 9:
+            body_parts_to_remove += ['leftForeArm',
+                                     'rightForeArm', 'leftArm', 'rightArm']
+        # Shorts
+        elif type == 7:
+            body_parts_to_remove += ['leftLeg', 'rightLeg',
+                                     'leftForeArm', 'rightForeArm', 'leftArm', 'rightArm']
+
+        verts_to_remove = list(itertools.chain.from_iterable(
+            [recon_labels[part] for part in body_parts_to_remove]))
+
+        label_mask = np.zeros(num_verts, dtype=bool)
+        label_mask[verts_to_remove] = True
+
+        seg_mask = np.zeros(num_verts, dtype=bool)
+        seg_mask[points_so_far] = True
+
+        # Remove points that belong to other bodyparts
+        # If a vertice in pointsSoFar is included in the bodyparts to remove, then these points should be removed
+        extra_verts_to_remove = np.array(list(seg_mask) and list(label_mask))
+
+        combine_mask = np.zeros(num_verts, dtype=bool)
+        combine_mask[points_so_far] = True
+        combine_mask[extra_verts_to_remove] = False
+
+        all_indices = np.argwhere(combine_mask == True).flatten()
+
+    i_x = np.where(np.in1d(faces[:, 0], all_indices))[0]
+    i_y = np.where(np.in1d(faces[:, 1], all_indices))[0]
+    i_z = np.where(np.in1d(faces[:, 2], all_indices))[0]
+
+    faces_to_keep = np.array(list(set(i_x).union(i_y).union(i_z)))
+    mask = np.zeros(len(recon.faces), dtype=bool)
+    if len(faces_to_keep) > 0:
+        mask[faces_to_keep] = True
+
+        mesh.update_faces(mask)
+        mesh.remove_unreferenced_vertices()
+
+        # mesh.rezero()
+
+        return mesh
+
+    return None