Ali Mohsin
feat: Add virtual try-on system components including DensePose, SMPL, and pix2pixHD models, rendering, and utilities.
5db43ff
| import cv2 | |
| import numpy as np | |
| from Graphonomy.human_parser import HumanParser | |
| from Graphonomy.dataset_settings import dataset_settings | |
| cihp_label = dataset_settings['cihp']['label'] | |
| import random | |
| def slide_window_garment_mask_refine(human_parser: HumanParser, num_window, raw_image): | |
| trans2roi=np.zeros((2,3),np.float32) | |
| trans2roi[0,0]=1 | |
| trans2roi[1, 1] = 1 | |
| height = raw_image.shape[0] | |
| width = raw_image.shape[1] | |
| trans_list = [] | |
| bias = width * 0.1 | |
| txs = np.linspace(int(-bias), int(bias), num_window).tolist() | |
| tys = np.linspace(int(-bias), int(bias), num_window).tolist() | |
| for i in range(num_window): | |
| for j in range(num_window): | |
| tx = int(txs[i]) | |
| ty = int(tys[j]) | |
| new_trans = trans2roi.copy() | |
| new_trans[0, 2] += tx | |
| new_trans[1, 2] += ty | |
| trans_list.append(new_trans) | |
| raw_mask_list = [] | |
| raw_valid_mask_list = [] | |
| for i in range(len(trans_list)): | |
| roi_img = cv2.warpAffine(raw_image, trans_list[i], (width, height), flags=cv2.INTER_LINEAR, | |
| borderMode=cv2.BORDER_CONSTANT, | |
| borderValue=(0, 0, 0)) | |
| garment_mask = get_garment_mask(human_parser, roi_img).astype(np.uint8) * 255 | |
| raw_mask = cv2.warpAffine(garment_mask, get_inverse_trans(trans_list[i]), (width, height), | |
| flags=cv2.INTER_LINEAR, | |
| borderMode=cv2.BORDER_CONSTANT, | |
| borderValue=(0)) | |
| raw_valid_mask = get_trans_mask(trans_list[i], height, width) | |
| raw_mask_list.append((raw_mask > 128).astype(np.uint8)) | |
| raw_valid_mask_list.append(raw_valid_mask) | |
| sum_mask = np.zeros((height, width), dtype=np.uint8) | |
| sum_valid_mask = np.zeros((height, width), dtype=np.uint8) | |
| for i in range(num_window*num_window): | |
| sum_mask += raw_mask_list[i] | |
| sum_valid_mask += raw_valid_mask_list[i] | |
| result_mask = sum_mask.astype(np.float32) > (sum_valid_mask.astype(np.float32) * 0.5) | |
| #result_mask = raw_valid_mask_list[-1].astype(np.float32) * 0.5 > 0 | |
| return result_mask | |
| def get_trans_mask(trans, height, width): | |
| valid_mask = np.ones((height, width), dtype=np.uint8) | |
| raw_valid_mask = cv2.warpAffine(valid_mask, get_inverse_trans(trans), (width, height), flags=cv2.INTER_LINEAR, | |
| borderMode=cv2.BORDER_CONSTANT, | |
| borderValue=(0)) | |
| return raw_valid_mask | |
| def get_garment_mask(human_parser: HumanParser, roi_img): | |
| garment_mask = human_parser.GetAtrGarmentMask(roi_img,isRGB=False) | |
| return garment_mask | |
| def get_inverse_trans(trans): | |
| full_matrix = np.vstack([trans, [0, 0, 1]]) | |
| # Compute the inverse of the 3x3 matrix | |
| inverse_matrix = np.linalg.inv(full_matrix) | |
| # Extract the top 2 rows for use with cv2.warpAffine | |
| inverse_transform_matrix = inverse_matrix[:2, :] | |
| return inverse_transform_matrix | |