初始化换发型项目：3个微服务代码 + 部署脚本

包含： - hair_service_sd: 换发型/换发色算法服务 (端口 8801) - photo_service: LoRA 训练调度服务 (端口 32678) - stable-diffusion-webui: SD WebUI 推理服务 (端口 57860) - kohya_ss_home: 训练环境代码 - meidaojia: 监控测试脚本 - setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查) - start_all_services.sh: 启动3个服务 - configure.ini.template: 路径模板化 (BASE_DIR自动推导) - conda_envs/py310.yml: py310 环境定义大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除，由网盘单独上传。
2026-07-11 18:11:49 +08:00
commit 0eb61f3e60
628 changed files with 120882 additions and 0 deletions
@@ -0,0 +1,462 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+import torch
+import numpy as np
+import os
+from utils import landmark_processor
+from algorithm_conf import ConfFactory
+from utils.umeyama import umeyama
+
+import cv2
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, is_1k=False, zero_init_residual=False):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        if is_1k:
+            self.fc = nn.Sequential(*[nn.Linear(512 * block.expansion, num_classes), nn.Tanh()])
+        else:
+            self.fc_key = nn.Sequential(*[nn.Linear(256 * block.expansion, 45 * 2), nn.Tanh()])
+            self.fc_ctrl = nn.Sequential(*[nn.Linear(256 * block.expansion, 48 * 2), nn.Tanh()])
+        self.is_1k = is_1k
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        if self.is_1k:
+            key = self.avgpool(x)
+            key = key.view(key.size(0), -1)
+            key = self.fc(key)
+            return key
+        else:
+            key, ctrl = torch.chunk(x, 2, dim=1)
+
+            key = self.avgpool(key)
+            key = key.view(key.size(0), -1)
+            key = self.fc_key(key)
+
+            ctrl = self.avgpool(ctrl)
+            ctrl = ctrl.view(ctrl.size(0), -1)
+            ctrl = self.fc_ctrl(ctrl)
+
+            return key, ctrl
+
+def resnet18(pretrained=False, **kwargs):
+    """Constructs a ResNet-18 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
+    return model
+
+
+def resnet34(pretrained=False, **kwargs):
+    """Constructs a ResNet-34 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+    return model
+
+
+def resnet50(pretrained=False, **kwargs):
+    """Constructs a ResNet-50 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def resnet101(pretrained=False, **kwargs):
+    """Constructs a ResNet-101 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model
+
+
+def resnet152(pretrained=False, **kwargs):
+    """Constructs a ResNet-152 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+    return model
+
+class Model1k(nn.Module):
+    def __init__(self, gpu_id=None):
+        super(Model1k, self).__init__()
+
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.face_alignment_net = resnet18(pretrained=False, num_classes=1000 * 2, is_1k=True)
+
+        self.model_dir = ConfFactory.getModelValue("model_dir")
+        weights = torch.load(os.path.join(self.model_dir, 'face_alignment_1k.pth'), map_location=lambda storage, loc: storage)
+
+        self.load_state_dict(weights)
+        self.to(self.device)
+        self.eval()
+
+    def forward(self, imgs):
+        pred_key_pts = self.face_alignment_net(imgs)
+        pred_key_pts = pred_key_pts + 0.5
+        return pred_key_pts
+
+class MomocvFaceAlignment1K(object):
+    def __init__(self, gpu_id=None):
+        self.gpu_id = gpu_id
+
+        self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
+
+        self.face_alignment_net = Model1k(gpu_id)
+
+        self.trackingFaceRects = []
+
+        print('MomocvFaceAlignment1K success')
+
+    def forward(self, img_tensor):
+        fullyconnected1 = self.face_alignment_net(img_tensor).detach().cpu().numpy()
+        return fullyconnected1
+
+    def detect(self, img, landmarks):
+        dst_size = 256
+        landmarks_res = []
+        with torch.no_grad():
+            input_numpy = np.zeros((len(landmarks), 3, dst_size, dst_size), dtype=np.float32)
+            all_mat = []
+            for ix, landmark in enumerate(landmarks):
+                M = landmark_processor.get_transform_mat_full_face(landmark, dst_size)
+                all_mat.append(M)
+                tmp = cv2.warpAffine(img, M, (dst_size, dst_size))
+
+                # cv2.imshow('inp', tmp)
+                # cv2.waitKey()
+
+                input_numpy[ix, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            for ix, pts in enumerate(fullyconnected1):
+                orig_pts = (np.reshape(pts, (2, 1000)).transpose((1, 0)) * dst_size)
+                orig_pts = landmark_processor.transform_points(orig_pts, all_mat[ix], invert=True)
+                landmarks_res.append(orig_pts)
+            return landmarks_res
+
+    def detect_single_face(self, img, crop_M):
+        dst_size = 256
+        with torch.no_grad():
+
+            #  add for change crop for 1024
+            h, w, _ = img.shape
+            if h == 768:
+                crop_img = img[104:img.shape[0] - 104, 104:img.shape[1] - 104, :]
+                tmp = cv2.resize(crop_img, (dst_size, dst_size))
+
+            else:
+                tmp = cv2.warpAffine(img, crop_M, (dst_size, dst_size), flags=cv2.INTER_CUBIC)
+
+                # cv2.imshow("img_paf_test_crop: ", tmp)
+                # cv2.waitKey()
+                # crop_img = img[220:img.shape[0] - 220, 266:img.shape[1] - 266, :]   # 220 266
+                # tmp = cv2.resize(crop_img, (dst_size, dst_size))
+
+            # cv2.imshow("detect face: h:{:d}".format(h), tmp)
+            # cv2.waitKey()
+
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+
+            orig_pts = np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0))
+
+            if h == 768:
+                orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1]
+                orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0]
+                orig_pts[:, 0] += 104
+                orig_pts[:, 1] += 104
+            else:
+                orig_pts[:, 0] = orig_pts[:, 0] * dst_size
+                orig_pts[:, 1] = orig_pts[:, 1] * dst_size
+                orig_pts = landmark_processor.transform_points(orig_pts, crop_M, invert=True)
+                # orig_pts[:, 0] = orig_pts[:, 0] * crop_img.shape[1]
+                # orig_pts[:, 1] = orig_pts[:, 1] * crop_img.shape[0]
+                # orig_pts[:, 0] += 266
+                # orig_pts[:, 1] += 220
+
+            return orig_pts
+
+    def detect_single_face_old(self, img):
+        dst_size = 256
+        with torch.no_grad():
+            tmp = cv2.resize(img, (dst_size, dst_size))
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * img.shape[0])
+            return orig_pts
+
+    def detect_according_5pts(self, img, pts5):
+        dst_size = 256
+        with torch.no_grad():
+            input_numpy = np.zeros((1, 3, dst_size, dst_size), dtype=np.float32)
+            eye_dis = 0.34
+            mouth_dis = 0.34
+            g_Average_5point_180 = np.array([
+                eye_dis, 0.3,
+                1 - eye_dis, 0.3,
+                0.5, 0.6,
+                mouth_dis, 0.63,
+                1 - mouth_dis, 0.63
+            ])
+            # print(g_Average_5point_180)
+            left_eye = np.array([pts5[0], pts5[5]])
+            right_eye = np.array([pts5[1], pts5[6]])
+            nose = np.array([pts5[2], pts5[7]])
+            left_mouth = np.array([pts5[3], pts5[8]])
+            right_mouth = np.array([pts5[4], pts5[9]])
+
+            pts5_src = np.vstack((left_eye, right_eye,
+                                  nose,
+                                  left_mouth, right_mouth))
+            pts5_src = np.array(pts5_src).astype(np.int32)
+            pts5_dst = g_Average_5point_180.reshape((5, -1)) * dst_size
+
+            mat = umeyama(pts5_src, pts5_dst, True)[0:2]
+
+            tmp = cv2.warpAffine(img, mat, (dst_size, dst_size))
+            # cv2.imshow("tmp", tmp)
+            # cv2.waitKey()
+            input_numpy[0, :, :, :] = tmp.transpose((2, 0, 1)).astype(np.float32) / 255
+            in_tensor = torch.from_numpy(input_numpy)
+            in_tensor = in_tensor.to(self.device)
+            fullyconnected1 = self.face_alignment_net(in_tensor).detach().cpu().numpy()
+            orig_pts = (np.reshape(fullyconnected1[0], (2, 1000)).transpose((1, 0)) * dst_size)
+            orig_pts = landmark_processor.transform_points(orig_pts, mat, invert=True)
+        return orig_pts
+
+
+    def stable_forward(self, image, detected_faces, reset=False):
+        if reset is True:
+            self.trackingFaceRects = []
+
+        if len(self.trackingFaceRects) == 0:
+            for face_rect in detected_faces:
+                    new_tracking_rect = [face_rect, True, [0, 0], 0, None]
+                    self.trackingFaceRects.append(new_tracking_rect)
+
+        with torch.no_grad():
+            landmarks = []
+            for ix, tracking_face_rect in enumerate(self.trackingFaceRects):
+                if tracking_face_rect[1] == True:
+                    d = tracking_face_rect[0]
+                    src_center = np.array([d[2] - (d[2] - d[0]) / 2.0, d[3] - (d[3] - d[1]) / 2.0])
+                    rotate_degree = tracking_face_rect[3]
+                    scale = 256 * 0.6 / min(d[2] - d[0], d[3] - d[1])
+                    dst_center = np.array([0.5, 0.5]) * 256
+                    offset = dst_center - src_center
+                    M = cv2.getRotationMatrix2D((src_center[0], src_center[1]), rotate_degree, scale)
+                    M[:, 2] += offset
+                else:
+                    rotate_degree = 0
+                    M = landmark_processor.get_transform_mat_mmcv_bigger(tracking_face_rect[4], 256)
+                inp = cv2.warpAffine(image, M, (256, 256))
+
+                # cv2.imshow('inp_{}'.format(ix), inp)
+                # cv2.waitKey()
+                orig_inp = inp
+
+                inp = inp.transpose((2, 0, 1)).astype(np.float32)
+                inp = inp[np.newaxis, :, :, :] / 255
+
+                in_tensor = torch.from_numpy(inp)
+                in_tensor = in_tensor.cuda(0)
+                fullyconnected1 = self.forward(in_tensor)
+                fullyconnected1 = fullyconnected1[0]
+                orig_pts = (np.reshape(fullyconnected1, (2, 1000)).transpose((1, 0))) * 256
+
+                t2 = cv2.getTickCount()
+                orig_pts = landmark_processor.transform_points(orig_pts, M, invert=True)
+
+                # orig_pts = orig_pts.transpose((1, 0))
+                fullyconnected1 = orig_pts
+
+                # update tracking infos
+                tracking_face_rect[1] = False
+                tracking_face_rect[2] = None
+                tracking_face_rect[3] = rotate_degree
+                tracking_face_rect[4] = fullyconnected1
+
+                # fullyconnected1 = landmark_processor.pts_1k_to_137(fullyconnected1)
+
+                # eye_landmark = self.detect_eye(image, fullyconnected1)
+                # fullyconnected1[87:104] = eye_landmark[0]
+                # fullyconnected1[104:121] = eye_landmark[1]
+
+                landmarks.append(fullyconnected1)
+        return landmarks
@@ -0,0 +1,330 @@
+import torch
+import numpy as np
+
+
+def point_form(boxes):
+    """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
+    representation for comparison to point form ground truth data.
+    Args:
+        boxes: (tensor) center-size default boxes from priorbox layers.
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return torch.cat((boxes[:, :2] - boxes[:, 2:]/2,     # xmin, ymin
+                     boxes[:, :2] + boxes[:, 2:]/2), 1)  # xmax, ymax
+
+
+def center_size(boxes):
+    """ Convert prior_boxes to (cx, cy, w, h)
+    representation for comparison to center-size form ground truth data.
+    Args:
+        boxes: (tensor) point_form boxes
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return torch.cat((boxes[:, 2:] + boxes[:, :2])/2,  # cx, cy
+                     boxes[:, 2:] - boxes[:, :2], 1)  # w, h
+
+
+def intersect(box_a, box_b):
+    """ We resize both tensors to [A,B,2] without new malloc:
+    [A,2] -> [A,1,2] -> [A,B,2]
+    [B,2] -> [1,B,2] -> [A,B,2]
+    Then we compute the area of intersect between box_a and box_b.
+    Args:
+      box_a: (tensor) bounding boxes, Shape: [A,4].
+      box_b: (tensor) bounding boxes, Shape: [B,4].
+    Return:
+      (tensor) intersection area, Shape: [A,B].
+    """
+    A = box_a.size(0)
+    B = box_b.size(0)
+    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
+                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+    inter = torch.clamp((max_xy - min_xy), min=0)
+    return inter[:, :, 0] * inter[:, :, 1]
+
+
+def jaccard(box_a, box_b):
+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+    is simply the intersection over union of two boxes.  Here we operate on
+    ground truth boxes and default boxes.
+    E.g.:
+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+    Args:
+        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
+        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
+    Return:
+        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
+    """
+    inter = intersect(box_a, box_b)
+    area_a = ((box_a[:, 2]-box_a[:, 0]) *
+              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
+    area_b = ((box_b[:, 2]-box_b[:, 0]) *
+              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
+    union = area_a + area_b - inter
+    return inter / union  # [A,B]
+
+
+def matrix_iou(a, b):
+    """
+    return iou of a and b, numpy version for data augenmentation
+    """
+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
+
+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
+    area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
+    return area_i / (area_a[:, np.newaxis] + area_b - area_i)
+
+
+def matrix_iof(a, b):
+    """
+    return iof of a and b, numpy version for data augenmentation
+    """
+    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
+    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
+
+    area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
+    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
+    return area_i / np.maximum(area_a[:, np.newaxis], 1)
+
+
+def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx):
+    """Match each prior box with the ground truth box of the highest jaccard
+    overlap, encode the bounding boxes, then return the matched indices
+    corresponding to both confidence and location preds.
+    Args:
+        threshold: (float) The overlap threshold used when mathing boxes.
+        truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
+        priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
+        variances: (tensor) Variances corresponding to each prior coord,
+            Shape: [num_priors, 4].
+        labels: (tensor) All the class labels for the image, Shape: [num_obj].
+        landms: (tensor) Ground truth landms, Shape [num_obj, 10].
+        loc_t: (tensor) Tensor to be filled w/ endcoded location targets.
+        conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
+        landm_t: (tensor) Tensor to be filled w/ endcoded landm targets.
+        idx: (int) current batch index
+    Return:
+        The matched indices corresponding to 1)location 2)confidence 3)landm preds.
+    """
+    # jaccard index
+    overlaps = jaccard(
+        truths,
+        point_form(priors)
+    )
+    # (Bipartite Matching)
+    # [1,num_objects] best prior for each ground truth
+    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
+
+    # ignore hard gt
+    valid_gt_idx = best_prior_overlap[:, 0] >= 0.2
+    best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]
+    if best_prior_idx_filter.shape[0] <= 0:
+        loc_t[idx] = 0
+        conf_t[idx] = 0
+        return
+
+    # [1,num_priors] best ground truth for each prior
+    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
+    best_truth_idx.squeeze_(0)
+    best_truth_overlap.squeeze_(0)
+    best_prior_idx.squeeze_(1)
+    best_prior_idx_filter.squeeze_(1)
+    best_prior_overlap.squeeze_(1)
+    best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2)  # ensure best prior
+    # TODO refactor: index  best_prior_idx with long tensor
+    # ensure every gt matches with its prior of max overlap
+    for j in range(best_prior_idx.size(0)):     # 判别此anchor是预测哪一个boxes
+        best_truth_idx[best_prior_idx[j]] = j
+    matches = truths[best_truth_idx]            # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来
+    conf = labels[best_truth_idx]               # Shape: [num_priors]      此处为每一个anchor对应的label取出来
+    conf[best_truth_overlap < threshold] = 0    # label as background   overlap<0.35的全部作为负样本
+    loc = encode(matches, priors, variances)
+
+    matches_landm = landms[best_truth_idx]
+    landm = encode_landm(matches_landm, priors, variances)
+    loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn
+    conf_t[idx] = conf  # [num_priors] top class label for each prior
+    landm_t[idx] = landm
+
+
+def encode(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 4].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded boxes (tensor), Shape: [num_priors, 4]
+    """
+
+    # dist b/t match center and prior's center
+    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, 2:])
+    # match wh / prior wh
+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
+    g_wh = torch.log(g_wh) / variances[1]
+    # return target for smooth_l1_loss
+    return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]
+
+def encode_landm(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth boxes
+    we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 10].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded landm (tensor), Shape: [num_priors, 10]
+    """
+
+    # dist b/t match center and prior's center
+    matched = torch.reshape(matched, (matched.size(0), 5, 2))
+    priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
+    priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2)
+    g_cxcy = matched[:, :, :2] - priors[:, :, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, :, 2:])
+    # g_cxcy /= priors[:, :, 2:]
+    g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1)
+    # return target for smooth_l1_loss
+    return g_cxcy
+
+
+# Adapted from https://github.com/Hakuyume/chainer-ssd
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = torch.cat((
+        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes
+
+def decode_landm(pre, priors, variances):
+    """Decode landm from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        pre (tensor): landm predictions for loc layers,
+            Shape: [num_priors,10]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded landm predictions
+    """
+    landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+                        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+                        ), dim=1)
+    return landms
+
+
+def log_sum_exp(x):
+    """Utility function for computing log_sum_exp while determining
+    This will be used to determine unaveraged confidence loss across
+    all examples in a batch.
+    Args:
+        x (Variable(tensor)): conf_preds from conf layers
+    """
+    x_max = x.data.max()
+    return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max
+
+
+# Original author: Francisco Massa:
+# https://github.com/fmassa/object-detection.torch
+# Ported to PyTorch by Max deGroot (02/01/2017)
+def nms(boxes, scores, overlap=0.5, top_k=200):
+    """Apply non-maximum suppression at test time to avoid detecting too many
+    overlapping bounding boxes for a given object.
+    Args:
+        boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
+        scores: (tensor) The class predscores for the img, Shape:[num_priors].
+        overlap: (float) The overlap thresh for suppressing unnecessary boxes.
+        top_k: (int) The Maximum number of box preds to consider.
+    Return:
+        The indices of the kept boxes with respect to num_priors.
+    """
+
+    keep = torch.Tensor(scores.size(0)).fill_(0).long()
+    if boxes.numel() == 0:
+        return keep
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    area = torch.mul(x2 - x1, y2 - y1)
+    v, idx = scores.sort(0)  # sort in ascending order
+    # I = I[v >= 0.01]
+    idx = idx[-top_k:]  # indices of the top-k largest vals
+    xx1 = boxes.new()
+    yy1 = boxes.new()
+    xx2 = boxes.new()
+    yy2 = boxes.new()
+    w = boxes.new()
+    h = boxes.new()
+
+    # keep = torch.Tensor()
+    count = 0
+    while idx.numel() > 0:
+        i = idx[-1]  # index of current largest val
+        # keep.append(i)
+        keep[count] = i
+        count += 1
+        if idx.size(0) == 1:
+            break
+        idx = idx[:-1]  # remove kept element from view
+        # load bboxes of next highest vals
+        torch.index_select(x1, 0, idx, out=xx1)
+        torch.index_select(y1, 0, idx, out=yy1)
+        torch.index_select(x2, 0, idx, out=xx2)
+        torch.index_select(y2, 0, idx, out=yy2)
+        # store element-wise max with next highest score
+        xx1 = torch.clamp(xx1, min=x1[i])
+        yy1 = torch.clamp(yy1, min=y1[i])
+        xx2 = torch.clamp(xx2, max=x2[i])
+        yy2 = torch.clamp(yy2, max=y2[i])
+        w.resize_as_(xx2)
+        h.resize_as_(yy2)
+        w = xx2 - xx1
+        h = yy2 - yy1
+        # check sizes of xx1 and xx2.. after each iteration
+        w = torch.clamp(w, min=0.0)
+        h = torch.clamp(h, min=0.0)
+        inter = w*h
+        # IoU = i / (area(a) + area(b) - i)
+        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
+        union = (rem_areas - inter) + area[i]
+        IoU = inter/union  # store result in iou
+        # keep only elements with an IoU <= overlap
+        idx = idx[IoU.le(overlap)]
+    return keep, count
+
+
@@ -0,0 +1,125 @@
+import requests
+import json
+from common.logger import config
+
+version = config.get('default', 'version')
+if version == "local":
+    current_photo_service_url = 'http://192.168.1.57:32678/'
+else:
+    current_photo_service_url = 'http://0.0.0.0:32678/'
+
+def call_hair_infer(task_id, hair_id, hair_material_dir, infer_req, is_hr, inference_port):
+    url = f"{current_photo_service_url}api/hair/inference"
+    payload = json.dumps({
+        "task_id": task_id,
+        "hair_id": hair_id,
+        "hd_version_flag":is_hr,
+        "hair_material_dir": hair_material_dir,
+        "request_json": infer_req,
+        "inference_port": inference_port
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    # print("---call infer payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    # print(response.json())
+
+    return response.json()
+
+def call_hair_infer_diy(task_id, infer_req, inference_port):
+    url = f"{current_photo_service_url}api/hair/inference_diy"
+    payload = json.dumps({
+        "task_id": task_id,
+        "request_json": infer_req,
+        "inference_port": inference_port
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    # print("---call infer payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    # print(response.json())
+
+    return response.json()
+
+
+
+def call_hair_enhance(img_path, mask_path, task_id, in_gender):
+    url = "http://127.0.0.1:7393/hairEnhance/v1"
+    payload = json.dumps({
+        "img_path": img_path,
+        "mask_path": mask_path,
+        "req_id": task_id,
+        "gender": in_gender
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    # print("---call infer payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    # print(response.json())
+
+    return response.json()
+
+def face_info(img_path):
+    url = "http://127.0.0.1:7393/faceInfo/v1"
+    payload = json.dumps({
+        "img": img_path,
+        "userId": 'fff',
+        "isLocal": True,
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    # print("---call infer payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    print(response.json())
+
+    return response.json()
+
+def is_same(img_url1, img_url2):
+    url = "http://127.0.0.1:7393/faceInfo/same"
+    payload = json.dumps({
+        "img_url1": img_url1,
+        "img_url2": img_url2,
+        # "isLocal": True,
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    # print("---call infer payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    print(response.json())
+
+    return response.json()
+
+def hair_select(img_url):
+    url = "http://127.0.0.1:7393/hairStyle/hair_select"
+    payload = json.dumps({
+        "img_url": img_url,
+        # "isLocal": True,
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    # print("---call infer payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    print(response.json())
+
+    return response.json()
+# test_url = 'https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/hairstyle/fb7d7a58-3228-40f2-84a6-337088ee31e2/2023031623425988.jpg'
+# test_url = 'https://ydapp-1317132355.cos.ap-beijing.myqcloud.com/hair_mz/images/vaffflue12.png'
+# hair_select(test_url)
@@ -0,0 +1,30 @@
+import requests
+import json
+from common.logger import config
+
+version = config.get('default', 'version')
+if version == "local":
+    current_photo_service_url = 'http://192.168.1.57:32678/'
+else:
+    current_photo_service_url = 'http://0.0.0.0:32678/'
+
+def call_hair_train(task_id, hair_id, hair_material_dir, tag, is_tj="0", webui_addr=None, device_id=None):
+    url = f"{webui_addr}api/hair/train"
+    payload = json.dumps({
+        "task_id": task_id,
+        "hair_id": hair_id,
+        "hair_material_dir": hair_material_dir,
+        "tag": tag,
+        "is_tj": is_tj,
+        "webui_addr": webui_addr,
+        "device_id":device_id
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    print("---call train payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    print(response.json())
+
@@ -0,0 +1,20 @@
+import requests
+import json
+
+def recall(req_id, state, message, clothId):
+    url = "http://192.168.11.220:9281/api/cloth/callBack"
+    payload = json.dumps({
+      "taskId": req_id,
+      "status": state,
+      "clothId": clothId,
+      "msg": message
+    })
+    headers = {
+      'Content-Type': 'application/json'
+    }
+    print("---payload:", payload)
+
+    response = requests.request("POST", url, headers=headers, data=payload)
+
+    print(response.json())
+
@@ -0,0 +1,27 @@
+import os
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+]
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+def mkdirs(paths):
+    """create empty directories if they don't exist
+    Parameters:
+        paths (str list) -- a list of directory paths
+    """
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            make_dir(path)
+    else:
+        make_dir(paths)
+
+def make_dir(target_dir):
+    """
+    Create dir if not exists
+    """
+    if not os.path.exists(target_dir):
+        os.makedirs(target_dir)
@@ -0,0 +1,80 @@
+import cv2
+import numpy as np
+
+import base64
+import requests
+from common.logger import config
+
+version = config.get('default', 'version')
+if version == "local":
+    webui_url = 'http://192.168.1.57:57860/'
+else:
+    webui_url = 'http://0.0.0.0:57860/'
+
+def encode_numpy_to_base64(img):
+    retval, bytes = cv2.imencode('.png', img)
+    encoded_image = base64.b64encode(bytes).decode('utf-8')
+    return encoded_image
+
+
+def webui_img2img(img, mask, prompt=''):
+    url = f"{webui_url}sdapi/v1/img2img"
+    request_dict = {
+            "prompt": prompt,
+            "negative_prompt": '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, bad_pictures,easynegative',
+            "sampler_name": "DPM++ 2M Karras",
+            "batch_size": 1,
+            "steps": 30,
+            "width": img.shape[1],
+            "height": img.shape[0],
+            "cfg_scale": 7.0,
+            "seed": 123456789,
+            "mask_blur": 5,
+            "init_images": [
+                encode_numpy_to_base64(img)
+            ],
+            "inpaint_full_res": False,
+            "inpainting_fill": 1,
+            "inpainting_mask_invert": 0,
+            "mask": encode_numpy_to_base64(mask),
+            # "refiner_checkpoint":"majicmixRealistic_v7.safetensors",
+            # "refiner_switch_at": 0.4,
+            "denoising_strength": 0.35,
+            "alwayson_scripts": {
+            }
+        }
+    response = requests.post(url=url, json=request_dict)
+    ret_json = response.json()
+    result = ret_json['images'][0]
+    img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR)
+    return img
+
+def webui_super_res_img(img, ratio):
+    url = f"{webui_url}sdapi/v1/extra-single-image"
+    request_dict = {
+                      "resize_mode": 0,
+                      "show_extras_results": False,
+                      "gfpgan_visibility": 0,
+                      "codeformer_visibility": 1,
+                      "codeformer_weight": 1,
+                      "upscaling_resize": ratio,
+                      "upscaler_1": "8x_NMKD-Superscale_150000_G",
+                      "upscale_first": False,
+                      "image": encode_numpy_to_base64(img)
+                    }
+    response = requests.post(url=url, json=request_dict)
+    ret_json = response.json()
+    result = ret_json['image']
+    img = cv2.imdecode(np.frombuffer(base64.b64decode(result), np.uint8), cv2.IMREAD_COLOR)
+    return img
+
+
+def webui_tag_by_clip(img):
+    url = f"{webui_url}sdapi/v1/interrogate"
+    request_dict = {
+                      "image": encode_numpy_to_base64(img),
+                      "model": "clip"
+                    }
+    response = requests.post(url=url, json=request_dict)
+    ret_json = response.json()
+    return ret_json['caption']
@@ -0,0 +1,25 @@
+import os
+import torch
+
+def load_model_by_path(model_save_path, model, gpu_id = None):
+    if not os.path.exists(model_save_path): return
+    loc = 'cpu' if gpu_id is None else 'cuda:{}'.format(gpu_id)
+    pretrained_dict = torch.load(model_save_path, map_location=loc)
+    if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict']
+    model_dict = model.state_dict()
+    # pretrained_dict.pop('netG.model.1.weight', '404')
+    pretrained_dict_new = {k: v for k, v in pretrained_dict.items()
+                       if (k in model_dict and model_dict[k].data.shape == v.data.shape)}
+    # if 'netG.model.1.weight' not in pretrained_dict_new and 'netG.model.1.weight' in model_dict:
+    #     pretrained_dict_new['netG.model.1.weight'] = model_dict['netG.model.1.weight']
+    #     pretrained_dict_new['netG.model.1.weight'][:,:12] = pretrained_dict['netG.model.1.weight']
+
+    model_dict.update(pretrained_dict_new)
+    model.load_state_dict(model_dict)
+    print("load model: ", model_save_path)
+
+def save_model_by_path(model_save_path, model):
+    save_dir, _ = os.path.split(model_save_path)
+    if not os.path.isdir(save_dir): os.makedirs(save_dir)
+    model_dic={k.replace('.module', ''):v for k,v in model.state_dict().items()}
+    torch.save(model_dic, model_save_path)
@@ -0,0 +1,326 @@
+import pickle
+import sys
+
+import cv2
+from PIL import Image
+from OpenGL.GL import *
+from OpenGL.GLU import *
+from OpenGL.GLUT import *
+import numpy as np
+
+# define shader code
+vertex_code='''
+uniform float scale;
+attribute vec2 position;
+attribute vec4 color;
+varying vec4 v_color;
+attribute vec2 TexCoordIn; 
+varying vec2 TexCoordOut;
+
+void main()
+{
+    gl_Position = vec4(position*scale, 0.0, 1.0);
+    v_color = color;
+    TexCoordOut = TexCoordIn;
+
+}'''
+
+fragment_code='''
+varying vec4 v_color;
+varying vec2 TexCoordOut; 
+uniform sampler2D Texture; 
+uniform vec2 originPosition;
+uniform vec2 targetPosition;
+vec2 curveWarp(vec2 textureCoord, vec2 originPosition, vec2 targetPosition, float radius)
+{
+    vec2 offset = vec2(0.0);
+    vec2 result = vec2(0.0);
+
+    vec2 direction = targetPosition - originPosition;
+
+    float infect = distance(textureCoord, originPosition)/radius;
+
+    infect = 1.0 - infect;
+    infect = clamp(infect, 0.0, 1.0);
+    offset = direction * infect;
+
+    result = textureCoord - offset;
+    return result;
+}
+
+void main()
+{
+    vec2 coordinate = vec2(0.0);
+    float radius = 0.5;
+    coordinate = curveWarp(TexCoordOut,originPosition,targetPosition,radius);
+    gl_FragColor = v_color*0.000000000001 + texture2D(Texture, coordinate);
+
+}'''
+
+
+
+#define useful function
+def display():
+    glClear(GL_COLOR_BUFFER_BIT)
+    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
+
+    glutSwapBuffers()
+
+def reshape(width,height):
+    glViewport(0, 0, width, height)
+
+#step1 init the context
+def  init(img_path):
+    image = Image.open(img_path)
+    glutInit()
+    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB)
+    glutCreateWindow('Hello world!')
+    glutReshapeWindow(image.width,image.height)
+    glutReshapeFunc(reshape)
+    glutDisplayFunc(display)
+    return image
+
+#step 2
+def initShaderProgram():
+    program  = glCreateProgram()
+    vertex   = glCreateShader(GL_VERTEX_SHADER)
+    fragment = glCreateShader(GL_FRAGMENT_SHADER)
+
+    # Set shaders source
+    glShaderSource(vertex, vertex_code)
+    glShaderSource(fragment, fragment_code)
+
+    # Compile shaders
+    glCompileShader(vertex)
+    glCompileShader(fragment)
+
+    fragSuccess = glGetShaderiv(fragment, GL_COMPILE_STATUS)
+    vertSuccess = glGetShaderiv(vertex, GL_COMPILE_STATUS)
+    print("vertext shader compile  success  [%s]" % (vertSuccess,))
+    print("fragment shader compile  success  [%s]" % (fragSuccess,))
+
+    if vertSuccess == 0:
+        print(glGetShaderInfoLog(vertex))
+        sys.exit(0)
+
+    if fragSuccess == 0:
+        print(glGetShaderInfoLog(fragment))
+        sys.exit(0)
+
+
+    glAttachShader(program, vertex)
+    glAttachShader(program, fragment)
+    glLinkProgram(program)
+    linksucc=glGetProgramiv(program, GL_LINK_STATUS)
+    print("link  program success [%s]" % (linksucc,))
+    glUseProgram(program)
+
+    return program
+
+#step 3.1 optional setup texture
+def getTextureFromFile(image_file):
+    #convert file to bytes
+    image = image_file.transpose(Image.FLIP_TOP_BOTTOM)
+    image = image.convert("RGBA")
+    byteImage =np.array(list(image.getdata()), np.uint8)
+
+    #setup texture
+    texIndex=glGenTextures(1)
+    glEnable( GL_TEXTURE_2D )
+    glBindTexture(GL_TEXTURE_2D,texIndex)
+
+    glPixelStorei(GL_UNPACK_ALIGNMENT,1)
+    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP)
+    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP)
+    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
+    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
+     #make the texture the default
+    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texIndex, 0)
+    glTexImage2D(GL_TEXTURE_2D,0,GL_RGBA,image.width,image.height,0,GL_RGBA,GL_UNSIGNED_BYTE,byteImage)
+
+    return texIndex
+
+
+#step 4.1 optional get the image from BufferFrame
+def saveImageFromFBO(width, height, output_img_path):
+    glReadBuffer(GL_COLOR_ATTACHMENT0)
+    glPixelStorei(GL_PACK_ALIGNMENT, 1)
+    data = glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE)
+    image = Image.new("RGB", (width, height), (0, 0, 0))
+    image.frombytes(data)
+    image = image.transpose(Image.FLIP_TOP_BOTTOM)
+    image.save(output_img_path)
+
+
+#step 4 optional if need to OSR generate a BufferFrame
+def setupSelfDefineFBO(program, image, data, output_img_path):
+    fbWidth, fbHeight = image.width, image.height
+
+    # Setup framebuffer
+    framebuffer = glGenFramebuffers(1)
+    glBindFramebuffer(GL_FRAMEBUFFER, framebuffer)
+
+    # # Setup colorbuffer
+    # colorbuffer = glGenRenderbuffers(1)
+    # glBindRenderbuffer(GL_RENDERBUFFER, colorbuffer)
+    # glRenderbufferStorage(GL_RENDERBUFFER, GL_RGBA, fbWidth, fbHeight)
+    # glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_RENDERBUFFER, colorbuffer)
+
+    # Setup depthbuffer
+    depthbuffer = glGenRenderbuffers (1)
+    glBindRenderbuffer (GL_RENDERBUFFER,depthbuffer)
+    glRenderbufferStorage (GL_RENDERBUFFER, GL_DEPTH_COMPONENT, image.width, image.height)
+    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, depthbuffer)
+
+    #first init VBO, then other parameters
+    buffer = glGenBuffers(1)     # Request a buffer slot from GPU
+    glBindBuffer(GL_ARRAY_BUFFER, buffer)     # Make this buffer the default one
+    glBufferData(GL_ARRAY_BUFFER, data.nbytes, data, GL_DYNAMIC_DRAW)    # Upload data
+
+    # Create texture to render to
+    # glBufferData(GL_FRAMEBUFFER, data.nbytes, data, GL_DYNAMIC_DRAW)
+    loc = glGetAttribLocation(program, "position")   #get the index of  the attribute in program
+    glEnableVertexAttribArray(loc)  #allow this attribute decide by index can be use
+    stride = data.strides[0]   #define how to read buffer
+    offset = ctypes.c_void_p(0) #define the offset where the data begin in buffer
+    glVertexAttribPointer(loc, 2, GL_FLOAT, False, stride, offset)
+    offset = ctypes.c_void_p(data.dtype["position"].itemsize)
+    loc = glGetAttribLocation(program, "color")
+    glEnableVertexAttribArray(loc)
+    glVertexAttribPointer(loc, 4, GL_FLOAT, False, stride, offset)
+
+    #setup other parameters
+    loc = glGetUniformLocation(program, "scale")
+    glUniform1f(loc, 1.0)
+
+    # originPosition = glGetUniformLocation(program, "originPosition")
+    # glUniform2f(originPosition, 0.5, 0.5)
+    #
+    # targetPosition = glGetUniformLocation(program, "targetPosition")
+    # glUniform2f(targetPosition, 0.47, 0.47)
+    # # glUniform2f(targetPosition, 0.5, 0.5)
+
+    # following code to bind uniform texture if needed
+    aTexture = getTextureFromFile(image)
+
+    glViewport(0, 0, fbWidth, fbHeight)
+
+    glActiveTexture(GL_TEXTURE0)
+    glBindTexture(GL_TEXTURE_2D, aTexture)
+    loc = glGetUniformLocation(program, "Texture")
+    glUniform1i(loc, 0)
+
+    loc = glGetAttribLocation(program, "TexCoordIn")
+    glEnableVertexAttribArray(loc)
+    offset=ctypes.c_void_p(data.dtype["color"].itemsize+8)
+    glVertexAttribPointer(loc, 2, GL_FLOAT, False, stride, offset)
+
+    status = glCheckFramebufferStatus (GL_FRAMEBUFFER)
+    if status != GL_FRAMEBUFFER_COMPLETE:
+        print( "Error in framebuffer activation")
+
+    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)
+    saveImageFromFBO(fbWidth, fbHeight, output_img_path)
+
+    glBindFramebuffer(GL_FRAMEBUFFER, GL_NONE)
+    glDeleteTextures([aTexture])
+    glDeleteFramebuffers(1, [framebuffer])
+
+    print('save image from FBO success')
+
+def pt_in_img(pt, img_w, img_h):
+
+    if pt[0] < 0 or pt[1] < 0 or pt[0] >= img_w or pt[1] >= img_h:
+        return False
+    else:
+        return True
+
+def demo_test():
+    input_img_path = "../test_data/pics/female003.jpg"
+    output_img_path = "../tmp.png"
+
+    ###### define vertex and color array
+
+    # data = np.zeros(4, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)])
+    # data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)]
+    # # data['position'] = [(-1, -1),   (-1, 1),   (1, -1),   (1, 1)]
+    # data['position'] = [(-1, -1), (-1, 1), (0, -1), (0, 1)]
+    # # data['textureCoord'] = [(0, 0), (0, 1), (1, 0), (1, 1)]
+    # data['textureCoord'] = [(0, 0),  (0, 1),  (1, 0),  (0.5, 1)]
+
+    data = np.zeros(8, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)])
+    data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)]
+    data['position'] = [(-1, -1), (-1, +1), (0, -1), (0, +1), (0, -1), (0, +1), (+1, -1), (+1, +1)]
+    # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0), (0.5, 1), (0.5, 0), (0.5, 1), (1, 0), (1, 1)]
+    mid_x = 0.6
+    data['textureCoord'] = [(0, 0), (0, 1), (mid_x, 0), (mid_x, 1), (mid_x, 0), (mid_x, 1), (1, 0), (1, 1)]
+
+    # data = np.zeros(3, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)])
+    # data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)]
+    # data['position'] = [(-1, -1), (-1, +1), (1, -1)]
+    # # data['position'] = [(-1, -1), (1, 1), (1, -1)]
+    # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0)]
+    # # data['textureCoord'] = [(0, 0), (1, 1), (1, 0)]
+
+    # data = np.zeros(6, dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)])
+    # data['color'] = [(1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1), (1, 1, 0, 1)]
+    # data['position'] = [(-1, -1), (-1, +1), (0, -1), (0, +1), (1, -1), (1, 1)]
+    # # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0), (0.5, 1), (1, 0), (1, 1)]
+    # data['textureCoord'] = [(0, 0), (0, 1), (0.5, 0), (0.5, 1), (1, 0), (1, 1)]
+    # mid_x = 0.4
+    # data['textureCoord'] = [(0, 0), (0, 1), (mid_x, 0), (mid_x, 1), (1, 0), (1, 1)]
+
+
+    image = init(input_img_path)
+    program = initShaderProgram()
+    setupSelfDefineFBO(program, image, data, output_img_path)
+
+def demo_warp_pt137():
+    input_img_path = "../test_data/pics/female003.jpg"
+    output_img_path = "../tmp.png"
+
+    import pickle
+    with open("../test_render.pkl", "rb") as fp:
+        info = pickle.load(fp)
+    vertice = info["vertice"]
+    dst_vertice = info["dst_vertice"]
+    faces = info["faces"]
+    color_list = []
+    position_list = []
+    textureCoord_list = []
+    img = cv2.imread(input_img_path)
+    img_h, img_w, _ = img.shape
+    for i in range(len(faces)):
+        # print("index: ", faces[i])
+        pt1 = vertice[faces[i][0]]
+        pt2 = vertice[faces[i][1]]
+        pt3 = vertice[faces[i][2]]
+        dst_pt1 = dst_vertice[faces[i][0]]
+        dst_pt2 = dst_vertice[faces[i][1]]
+        dst_pt3 = dst_vertice[faces[i][2]]
+        # if pt1[0] != dst_pt1[0]:
+        #     print("use warp!")
+        if pt_in_img(pt1, img_w, img_h) and pt_in_img(pt2, img_w, img_h) and pt_in_img(pt3, img_w, img_h):
+            color_list.append((1, 1, 0, 1))
+            color_list.append((1, 1, 0, 1))
+            color_list.append((1, 1, 0, 1))
+            position_list.append((dst_pt1[0] * 2 / img_w - 1, dst_pt1[1] * 2 / img_h - 1))
+            position_list.append((dst_pt2[0] * 2 / img_w - 1, dst_pt2[1] * 2 / img_h - 1))
+            position_list.append((dst_pt3[0] * 2 / img_w - 1, dst_pt3[1] * 2 / img_h - 1))
+            textureCoord_list.append((pt1[0] / img_w, pt1[1] / img_h))
+            textureCoord_list.append((pt2[0] / img_w, pt2[1] / img_h))
+            textureCoord_list.append((pt3[0] / img_w, pt3[1] / img_h))
+
+    data = np.zeros(len(color_list),
+                    dtype=[("position", np.float32, 2), ("color", np.float32, 4), ("textureCoord", np.float32, 2)])
+    data['color'] = color_list
+    data['position'] = position_list
+    data['textureCoord'] = textureCoord_list
+
+    image = init(input_img_path)
+    program = initShaderProgram()
+    setupSelfDefineFBO(program, image, data, output_img_path)
+
+if __name__ == '__main__':
+    demo_test()
+    # demo_warp_pt137()
@@ -0,0 +1,202 @@
+import math
+import os
+import time
+from copy import deepcopy
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+
+def init_seeds(seed=0):
+    torch.manual_seed(seed)
+
+    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
+    if seed == 0:  # slower, more reproducible
+        cudnn.deterministic = True
+        cudnn.benchmark = False
+    else:  # faster, less reproducible
+        cudnn.deterministic = False
+        cudnn.benchmark = True
+
+
+def select_device(device='', apex=False, batch_size=None):
+    # device = 'cpu' or '0' or '0,1,2,3'
+    cpu_request = device.lower() == 'cpu'
+    if device and not cpu_request:  # if device requested other than 'cpu'
+        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
+        assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device  # check availablity
+
+    cuda = False if cpu_request else torch.cuda.is_available()
+    if cuda:
+        c = 1024 ** 2  # bytes to MB
+        ng = torch.cuda.device_count()
+        if ng > 1 and batch_size:  # check that batch_size is compatible with device_count
+            assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)
+        x = [torch.cuda.get_device_properties(i) for i in range(ng)]
+        s = 'Using CUDA ' + ('Apex ' if apex else '')  # apex for mixed precision https://github.com/NVIDIA/apex
+        for i in range(0, ng):
+            if i == 1:
+                s = ' ' * len(s)
+            print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
+                  (s, i, x[i].name, x[i].total_memory / c))
+    else:
+        print('Using CPU')
+
+    print('')  # skip a line
+    return torch.device('cuda:0' if cuda else 'cpu')
+
+
+def time_synchronized():
+    torch.cuda.synchronize() if torch.cuda.is_available() else None
+    return time.time()
+
+
+def initialize_weights(model):
+    for m in model.modules():
+        t = type(m)
+        if t is nn.Conv2d:
+            pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        elif t is nn.BatchNorm2d:
+            m.eps = 1e-4
+            m.momentum = 0.03
+        elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
+            m.inplace = True
+
+
+def find_modules(model, mclass=nn.Conv2d):
+    # finds layer indices matching module class 'mclass'
+    return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
+
+
+def fuse_conv_and_bn(conv, bn):
+    # https://tehnokv.com/posts/fusing-batchnorm-and-conv/
+    with torch.no_grad():
+        # init
+        fusedconv = torch.nn.Conv2d(conv.in_channels,
+                                    conv.out_channels,
+                                    kernel_size=conv.kernel_size,
+                                    stride=conv.stride,
+                                    padding=conv.padding,
+                                    bias=True)
+
+        # prepare filters
+        w_conv = conv.weight.clone().view(conv.out_channels, -1)
+        w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
+        fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
+
+        # prepare spatial bias
+        if conv.bias is not None:
+            b_conv = conv.bias
+        else:
+            b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device)
+        b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
+        fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
+
+        return fusedconv
+
+
+def model_info(model, verbose=False):
+    # Plots a line-by-line description of a PyTorch model
+    n_p = sum(x.numel() for x in model.parameters())  # number parameters
+    n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients
+    if verbose:
+        print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
+        for i, (name, p) in enumerate(model.named_parameters()):
+            name = name.replace('module_list.', '')
+            print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
+                  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
+
+    try:  # FLOPS
+        from thop import profile
+        macs, _ = profile(model, inputs=(torch.zeros(1, 3, 480, 640),), verbose=False)
+        fs = ', %.1f GFLOPS' % (macs / 1E9 * 2)
+    except:
+        fs = ''
+
+    print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs))
+
+
+def load_classifier(name='resnet101', n=2):
+    # Loads a pretrained model reshaped to n-class output
+    model = models.__dict__[name](pretrained=True)
+
+    # Display model properties
+    input_size = [3, 224, 224]
+    input_space = 'RGB'
+    input_range = [0, 1]
+    mean = [0.485, 0.456, 0.406]
+    std = [0.229, 0.224, 0.225]
+    for x in [input_size, input_space, input_range, mean, std]:
+        print(x + ' =', eval(x))
+
+    # Reshape output to n classes
+    filters = model.fc.weight.shape[1]
+    model.fc.bias = torch.nn.Parameter(torch.zeros(n), requires_grad=True)
+    model.fc.weight = torch.nn.Parameter(torch.zeros(n, filters), requires_grad=True)
+    model.fc.out_features = n
+    return model
+
+
+def scale_img(img, ratio=1.0, same_shape=False):  # img(16,3,256,416), r=ratio
+    # scales img(bs,3,y,x) by ratio
+    h, w = img.shape[2:]
+    s = (int(h * ratio), int(w * ratio))  # new size
+    img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
+    if not same_shape:  # pad/crop img
+        gs = 32  # (pixels) grid size
+        h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
+    return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
+
+
+class ModelEMA:
+    """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
+    Keep a moving average of everything in the model state_dict (parameters and buffers).
+    This is intended to allow functionality like
+    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
+    A smoothed version of the weights is necessary for some training schemes to perform well.
+    E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use
+    RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA
+    smoothing of weights to match results. Pay attention to the decay constant you are using
+    relative to your update count per epoch.
+    To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but
+    disable validation of the EMA weights. Validation will have to be done manually in a separate
+    process, or after the training stops converging.
+    This class is sensitive where it is initialized in the sequence of model init,
+    GPU assignment and distributed training wrappers.
+    I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and single-GPU.
+    """
+
+    def __init__(self, model, decay=0.9999, device=''):
+        # make a copy of the model for accumulating moving average of weights
+        self.ema = deepcopy(model)
+        self.ema.eval()
+        self.updates = 0  # number of EMA updates
+        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)
+        self.device = device  # perform ema on different device from model if set
+        if device:
+            self.ema.to(device=device)
+        for p in self.ema.parameters():
+            p.requires_grad_(False)
+
+    def update(self, model):
+        self.updates += 1
+        d = self.decay(self.updates)
+        with torch.no_grad():
+            if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel):
+                msd, esd = model.module.state_dict(), self.ema.module.state_dict()
+            else:
+                msd, esd = model.state_dict(), self.ema.state_dict()
+
+            for k, v in esd.items():
+                if v.dtype.is_floating_point:
+                    v *= d
+                    v += (1. - d) * msd[k].detach()
+
+    def update_attr(self, model):
+        # Assign attributes (which may change during training)
+        for k in model.__dict__.keys():
+            if not k.startswith('_'):
+                setattr(self.ema, k, getattr(model, k))
@@ -0,0 +1,71 @@
+import numpy as np
+
+def umeyama(src, dst, estimate_scale):
+    """Estimate N-D similarity transformation with or without scaling.
+    Parameters
+    ----------
+    src : (M, N) array
+        Source coordinates.
+    dst : (M, N) array
+        Destination coordinates.
+    estimate_scale : bool
+        Whether to estimate scaling factor.
+    Returns
+    -------
+    T : (N + 1, N + 1)
+        The homogeneous similarity transformation matrix. The matrix contains
+        NaN values only if the problem is not well-conditioned.
+    References
+    ----------
+    .. [1] "Least-squares estimation of transformation parameters between two
+            point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573
+    """
+
+    num = src.shape[0]
+    dim = src.shape[1]
+
+    # Compute mean of src and dst.
+    src_mean = src.mean(axis=0)
+    dst_mean = dst.mean(axis=0)
+
+    # Subtract mean from src and dst.
+    src_demean = src - src_mean
+    dst_demean = dst - dst_mean
+
+    # Eq. (38).
+    A = np.dot(dst_demean.T, src_demean) / num
+
+    # Eq. (39).
+    d = np.ones((dim,), dtype=np.double)
+    if np.linalg.det(A) < 0:
+        d[dim - 1] = -1
+
+    T = np.eye(dim + 1, dtype=np.double)
+
+    U, S, V = np.linalg.svd(A)
+
+    # Eq. (40) and (43).
+    rank = np.linalg.matrix_rank(A)
+    if rank == 0:
+        return np.nan * T
+    elif rank == dim - 1:
+        if np.linalg.det(U) * np.linalg.det(V) > 0:
+            T[:dim, :dim] = np.dot(U, V)
+        else:
+            s = d[dim - 1]
+            d[dim - 1] = -1
+            T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V))
+            d[dim - 1] = s
+    else:
+        T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T))
+
+    if estimate_scale:
+        # Eq. (41) and (42).
+        scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d)
+    else:
+        scale = 1.0
+
+    T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T)
+    T[:dim, :dim] *= scale
+
+    return T
@@ -0,0 +1,287 @@
+import os
+import cv2
+import torch
+import logging
+import numpy as np
+# from utils.config import CONFIG
+# import torch.distributed as dist
+
+def mkdirs(paths):
+    """create empty directories if they don't exist
+    Parameters:
+        paths (str list) -- a list of directory paths
+    """
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            os.makedirs(path)
+    else:
+        os.makedirs(paths)
+
+def make_dir(target_dir):
+    """
+    Create dir if not exists
+    """
+    if not os.path.exists(target_dir):
+        os.makedirs(target_dir)
+
+
+def print_network(model, name):
+    """
+    Print out the network information
+    """
+    logger = logging.getLogger("Logger")
+    num_params = 0
+    for p in model.parameters():
+        num_params += p.numel()
+
+    logger.info(model)
+    logger.info(name)
+    logger.info("Number of parameters: {}".format(num_params))
+
+
+def update_lr(lr, optimizer):
+    """
+    update learning rates
+    """
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def warmup_lr(init_lr, step, iter_num):
+    """
+    Warm up learning rate
+    """
+    return step/iter_num*init_lr
+
+
+def add_prefix_state_dict(state_dict, prefix="module"):
+    """
+    add prefix from the key of pretrained state dict for Data-Parallel
+    """
+    new_state_dict = {}
+    first_state_name = list(state_dict.keys())[0]
+    if not first_state_name.startswith(prefix):
+        for key, value in state_dict.items():
+            new_state_dict[prefix+"."+key] = state_dict[key].float()
+    else:
+        for key, value in state_dict.items():
+            new_state_dict[key] = state_dict[key].float()
+    return new_state_dict
+
+
+def remove_prefix_state_dict(state_dict, prefix="module"):
+    """
+    remove prefix from the key of pretrained state dict for Data-Parallel
+    """
+    new_state_dict = {}
+    first_state_name = list(state_dict.keys())[0]
+    if not first_state_name.startswith(prefix):
+        for key, value in state_dict.items():
+            new_state_dict[key] = state_dict[key].float()
+    else:
+        for key, value in state_dict.items():
+            new_state_dict[key[len(prefix)+1:]] = state_dict[key].float()
+    return new_state_dict
+
+#
+# def load_imagenet_pretrain(model, checkpoint_file):
+#     """
+#     Load imagenet pretrained resnet
+#     Add zeros channel to the first convolution layer
+#     Since we have the spectral normalization, we need to do a little more
+#     """
+#     checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu))
+#     state_dict = remove_prefix_state_dict(checkpoint['state_dict'])
+#     for key, value in state_dict.items():
+#         state_dict[key] = state_dict[key].float()
+#
+#     logger = logging.getLogger("Logger")
+#     logger.debug("Imagenet pretrained keys:")
+#     logger.debug(state_dict.keys())
+#     logger.debug("Generator keys:")
+#     logger.debug(model.module.encoder.state_dict().keys())
+#     logger.debug("Intersection  keys:")
+#     logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys()))
+#
+#     weight_u = state_dict["conv1.module.weight_u"]
+#     weight_v = state_dict["conv1.module.weight_v"]
+#     weight_bar = state_dict["conv1.module.weight_bar"]
+#
+#     logger.debug("weight_v: {}".format(weight_v))
+#     logger.debug("weight_bar: {}".format(weight_bar.view(32, -1)))
+#     logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v))))
+#
+#     new_weight_v = torch.zeros(6, 3, 3).cuda()
+#     new_weight_bar = torch.zeros(32, 6, 3, 3).cuda()
+#
+#     new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3))
+#     new_weight_bar[:, :3, :, :].copy_(weight_bar)
+#
+#     logger.debug("new weight_v: {}".format(new_weight_v.view(-1)))
+#     logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1)))
+#     logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1)))))
+#
+#     state_dict["conv1.module.weight_v"] = new_weight_v.view(-1)
+#     state_dict["conv1.module.weight_bar"] = new_weight_bar
+#
+#     model.module.encoder.load_state_dict(state_dict, strict=False)
+
+
+def load_VGG_pretrain(model, checkpoint_file):
+    """
+    Load imagenet pretrained resnet
+    Add zeros channel to the first convolution layer
+    Since we have the spectral normalization, we need to do a little more
+    """
+    checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda())
+    backbone_state_dict = remove_prefix_state_dict(checkpoint['state_dict'])
+
+    model.module.encoder.load_state_dict(backbone_state_dict, strict=False)
+
+
+def get_unknown_tensor(trimap):
+    """
+    get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor
+    """
+    # if CONFIG.model.trimap_channel == 3:
+    weight = trimap[:, 1:2, :, :].float()
+    # else:
+    #     weight = trimap.eq(1).float()
+    return weight
+
+
+def get_gaborfilter(angles):
+    """
+    generate gabor filter as the conv kernel
+    :param angles: number of different angles
+    """
+    gabor_filter = []
+    for angle in range(angles):
+        gabor_filter.append(cv2.getGaborKernel(ksize=(5,5), sigma=0.5, theta=angle*np.pi/8, lambd=5, gamma=0.5))
+    gabor_filter = np.array(gabor_filter)
+    gabor_filter = np.expand_dims(gabor_filter, axis=1)
+    return gabor_filter.astype(np.float32)
+
+
+def get_gradfilter():
+    """
+    generate gradient filter as the conv kernel
+    """
+    grad_filter = []
+    grad_filter.append([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
+    grad_filter.append([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
+    grad_filter = np.array(grad_filter)
+    grad_filter = np.expand_dims(grad_filter, axis=1)
+    return grad_filter.astype(np.float32)
+
+
+# def reduce_tensor_dict(tensor_dict, mode='mean'):
+#     """
+#     average tensor dict over different GPUs
+#     """
+#     for key, tensor in tensor_dict.items():
+#         if tensor is not None:
+#             tensor_dict[key] = reduce_tensor(tensor, mode)
+#     return tensor_dict
+#
+#
+# def reduce_tensor(tensor, mode='mean'):
+#     """
+#     average tensor over different GPUs
+#     """
+#     rt = tensor.clone()
+#     dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+#     if mode == 'mean':
+#         rt /= CONFIG.world_size
+#     elif mode == 'sum':
+#         pass
+#     else:
+#         raise NotImplementedError("reduce mode can only be 'mean' or 'sum'")
+#     return rt
+
+def make_color_wheel():
+    # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
+    RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6)
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros([ncols, 3])
+    col = 0
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
+    col += RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
+    colorwheel[col:col+YG, 1] = 255
+    col += YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
+    col += GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
+    colorwheel[col:col+CB, 2] = 255
+    col += CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
+    col += + BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+
+COLORWHEEL = make_color_wheel()
+
+def compute_color(u,v):
+    # from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
+    h, w = u.shape
+    img = np.zeros([h, w, 3])
+    nanIdx = np.isnan(u) | np.isnan(v)
+    u[nanIdx] = 0
+    v[nanIdx] = 0
+    colorwheel = COLORWHEEL
+    # colorwheel = make_color_wheel()
+    ncols = np.size(colorwheel, 0)
+    rad = np.sqrt(u**2+v**2)
+    a = np.arctan2(-v, -u) / np.pi
+    fk = (a+1) / 2 * (ncols - 1) + 1
+    k0 = np.floor(fk).astype(int)
+    k1 = k0 + 1
+    k1[k1 == ncols+1] = 1
+    f = fk - k0
+    for i in range(np.size(colorwheel,1)):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0-1] / 255
+        col1 = tmp[k1-1] / 255
+        col = (1-f) * col0 + f * col1
+        idx = rad <= 1
+        col[idx] = 1-rad[idx]*(1-col[idx])
+        notidx = np.logical_not(idx)
+        col[notidx] *= 0.75
+        img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))
+    return img
+
+def flow_to_image(flow):
+    # part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
+    maxrad = -1
+    u = flow[0, :, :]
+    v = flow[1, :, :]
+    rad = np.sqrt(u ** 2 + v ** 2)
+    maxrad = max(maxrad, np.max(rad))
+    u = u/(maxrad + np.finfo(float).eps)
+    v = v/(maxrad + np.finfo(float).eps)
+    img = compute_color(u, v)
+
+    return img
+
+
+if __name__ == "__main__":
+    import networks
+    logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
+                        datefmt='%m-%d %H:%M:%S')
+    G = networks.get_generator().cuda()
+    # load_imagenet_pretrain(G, CONFIG.model.imagenet_pretrain_path)
+    x = torch.randn(4,3,512,512).cuda()
+    y = torch.randn(4,3,512,512).cuda()
+    z = G(x, y)
@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+
+import torch.nn as nn
+
+
+def c2_xavier_fill(module: nn.Module):
+    """
+    Initialize `module.weight` using the "XavierFill" implemented in Caffe2.
+    Also initializes `module.bias` to 0.
+
+    Args:
+        module (torch.nn.Module): module to initialize.
+    """
+    # Caffe2 implementation of XavierFill in fact
+    # corresponds to kaiming_uniform_ in PyTorch
+    nn.init.kaiming_uniform_(module.weight, a=1)
+    if module.bias is not None:
+        nn.init.constant_(module.bias, 0)
+
+
+def c2_msra_fill(module: nn.Module):
+    """
+    Initialize `module.weight` using the "MSRAFill" implemented in Caffe2.
+    Also initializes `module.bias` to 0.
+
+    Args:
+        module (torch.nn.Module): module to initialize.
+    """
+    nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+    if module.bias is not None:
+        nn.init.constant_(module.bias, 0)