Ali Mohsin
feat: Add virtual try-on system components including DensePose, SMPL, and pix2pixHD models, rendering, and utilities.
5db43ff
| import numpy as np | |
| import sys | |
| import cv2 | |
| sys.path.append('./util/cpp_extensions/build') | |
| #import example # Replace 'example' with the name of your module | |
| #1, 2 = Torso, 3 = Right Hand, 4 = Left Hand, 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right, | |
| # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right, 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left, | |
| # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left, 20, 22 = Lower Arm Right, 23, 24 = Head | |
| def IUV2UpperBodyMask(IUV: np.ndarray): | |
| upper_index = [1, 2, 15, 17, 16, 18, 19, 21, 20, 22, 23, 24]#[1, 2, 3, 4, 15, 17, 16, 18, 19, 21, 20, 22, 23, 24] | |
| h, w = IUV.shape[:2] | |
| mask = np.zeros((h, w), bool) | |
| for i in upper_index: | |
| mask[IUV[:, :, 0] == i] = True | |
| return mask | |
| def IUV2UpperBodyRoiTrans(IUV: np.ndarray,roi_size=1024,roi_ratio=1.2): | |
| mask = IUV2UpperBodyMask(IUV) | |
| h,w = IUV.shape[:2] | |
| y_indices, x_indices = np.where(mask == True) | |
| # Calculate the centroid | |
| centroid_x = np.mean(x_indices) | |
| centroid_y = np.mean(y_indices) | |
| # Finding the indices of non-zero elements | |
| rows = np.any(mask, axis=1) | |
| cols = np.any(mask, axis=0) | |
| ymin, ymax = np.where(rows)[0][[0, -1]] | |
| xmin, xmax = np.where(cols)[0][[0, -1]] | |
| size=max([ymax-centroid_y, centroid_y-ymin,xmax-centroid_x, centroid_x-xmin]) | |
| src = np.zeros([3, 2], np.float32) | |
| center = np.array([centroid_x,centroid_y],np.float32)#Joints[9] * 0.8 + Joints[12] * 0.2 | |
| size = size * roi_ratio | |
| src[0, :] = center + np.array([-size, size], np.float32) | |
| src[1, :] = center + np.array([-size, -size], np.float32) | |
| src[2, :] = center + np.array([size, -size], np.float32) | |
| dst = np.zeros([3, 2], np.float32) | |
| dst[0, :] = np.array([0, roi_size - 1], np.float32) | |
| dst[1, :] = np.array([0, 0], np.float32) | |
| dst[2, :] = np.array([roi_size - 1, 0], np.float32) | |
| trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) | |
| inv_trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) | |
| return trans, inv_trans | |
| def IUV2Img(IUV: np.ndarray): | |
| IUV = IUV.astype(np.float32) | |
| IUV[:, :, 0] /= 24.0 | |
| IUV[:, :, 0] *= 255 | |
| IUV = IUV.astype(np.uint8) | |
| return IUV | |
| def IUV2UpperBodyImg(IUV: np.ndarray): | |
| remove_index=[5,6,7,9,8,10,11,13,12,14,23,24] | |
| torso_index = [1, 2] | |
| result=IUV.copy() | |
| for i in remove_index: | |
| result[IUV[:,:,0]==i]=0 | |
| result = result.astype(np.float32) | |
| result[:, :, 0] /= 24.0 | |
| result[:, :, 0] *= 255 | |
| result = result.astype(np.uint8) | |
| for i in torso_index: | |
| result[IUV[:, :, 0] == i] = 255 | |
| return result | |
| def IUV2TorsoLeg(IUV: np.ndarray): | |
| torsoleg_index = [1, 2, 5, 6, 7, 9, 8, 10, 11, 13, 12, 14] | |
| result = IUV.copy() | |
| result = result.astype(np.float32) | |
| result[:, :, 0] /= 24 | |
| result[:, :, 0] *= 255 | |
| result = result.astype(np.uint8) | |
| for i in torsoleg_index: | |
| result[IUV[:, :, 0] == i] = 255 | |
| return result | |
| def IUV2SDP(IUV: np.ndarray): | |
| torsoleg_index = [1, 2, 5, 6, 7, 9, 8, 10, 11, 13, 12, 14] | |
| # Copy and normalize the IUV array | |
| result = IUV.copy().astype(np.float32) | |
| result[:, :, 0] = (result[:, :, 0] / 24) * 255 | |
| result = result.astype(np.uint8) | |
| # Create a mask for the torso and leg indices | |
| torsoleg_mask = np.isin(IUV[:, :, 0], torsoleg_index) | |
| # Apply the mask to set the corresponding pixels to 255 | |
| result[torsoleg_mask] = 255 | |
| return result | |
| #1, 2 = Torso, 3 = Right Hand, 4 = Left Hand, 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right, | |
| # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right, 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left, | |
| # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left, 20, 22 = Lower Arm Right, 23, 24 = Head | |
| def IUV2SSDP_old(IUV: np.ndarray): | |
| torsoleghead_index = [1, 2, 5, 6, 7, 9, 8, 10, 11, 13, 12, 14, 23,24] | |
| left_arm_index= [4,15,17,19,21] | |
| right_arm_index = [3,16,18,20,22] | |
| result = IUV.copy() | |
| result = result.astype(np.float32) | |
| result[:, :, 0] /= 24 | |
| result[:, :, 0] *= 255 | |
| result = result.astype(np.uint8) | |
| for i in torsoleghead_index: | |
| result[IUV[:, :, 0] == i] = 255 | |
| for i in left_arm_index: | |
| result[IUV[:, :, 0] == i] = np.array([255,0,0], np.uint8) | |
| for i in right_arm_index: | |
| result[IUV[:, :, 0] == i] = np.array([0,0, 255], np.uint8) | |
| return result | |
| def IUV2SSDP_deprecated(IUV: np.ndarray): | |
| torsoleghead_index = [1, 2, 5, 6, 7, 9, 8, 10, 11, 13, 12, 14, 23, 24] | |
| left_arm_index = [4, 15, 17, 19, 21] | |
| right_arm_index = [3, 16, 18, 20, 22] | |
| result = IUV.copy().astype(np.float32) | |
| result[:, :, 0] = (result[:, :, 0] / 24) * 255 | |
| result = result.astype(np.uint8) | |
| # Create masks for each index group | |
| torsoleghead_mask = np.isin(IUV[:, :, 0], torsoleghead_index) | |
| left_arm_mask = np.isin(IUV[:, :, 0], left_arm_index) | |
| right_arm_mask = np.isin(IUV[:, :, 0], right_arm_index) | |
| # Apply masks to result | |
| result[torsoleghead_mask] = 255 | |
| result[left_arm_mask] = np.array([255, 0, 0], np.uint8) | |
| result[right_arm_mask] = np.array([0, 0, 255], np.uint8) | |
| return result | |
| def IUV2SSDP(IUV: np.ndarray): | |
| return IUV2SSDP_new(IUV) | |
| #1, 2 = Torso, 3 = Right Hand, 4 = Left Hand, 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right, | |
| # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right, 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left, | |
| # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left, 20, 22 = Lower Arm Right, 23, 24 = Head | |
| def IUV2SSDP_new(IUV: np.ndarray): | |
| torsoleghead_index = [1, 2, 5, 6, 7, 9, 8, 10, 11, 13, 12, 14, 23, 24] | |
| left_upper_arm_index = [15, 17] | |
| left_lower_arm_index = [19, 21] | |
| left_hand_index = [4,] | |
| right_upper_arm_index = [16, 18] | |
| right_lower_arm_index = [20, 22] | |
| right_hand_index = [3,] | |
| result = IUV.copy().astype(np.float32) | |
| result[:, :, 0] = (result[:, :, 0] / 24) * 255 | |
| result = result.astype(np.uint8) | |
| # Create masks for each index group | |
| torsoleghead_mask = np.isin(IUV[:, :, 0], torsoleghead_index) | |
| left_upper_arm_mask = np.isin(IUV[:, :, 0], left_upper_arm_index) | |
| left_lower_arm_mask = np.isin(IUV[:, :, 0], left_lower_arm_index) | |
| left_hand_mask = np.isin(IUV[:, :, 0], left_hand_index) | |
| right_upper_arm_mask = np.isin(IUV[:, :, 0], right_upper_arm_index) | |
| right_lower_arm_mask = np.isin(IUV[:, :, 0], right_lower_arm_index) | |
| right_hand_mask = np.isin(IUV[:, :, 0], right_hand_index) | |
| # Apply masks to result | |
| result[torsoleghead_mask] = 255 | |
| result[left_upper_arm_mask] = np.array([255, 0, 0], np.uint8) | |
| result[left_lower_arm_mask] = np.array([255, 255, 0], np.uint8) | |
| result[left_hand_mask] = np.array([255, 0, 255], np.uint8) | |
| result[right_upper_arm_mask] = np.array([0, 0, 255], np.uint8) | |
| result[right_lower_arm_mask] = np.array([0, 255, 255], np.uint8) | |
| result[right_hand_mask] = np.array([0, 255, 0], np.uint8) | |
| return result |