m53297601
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

new add centerface network.

Browse Source
pull/7961/head
linqingke 4 years ago
parent f0e99e1099
commit 11ba97acde
33 changed files with 5886 additions and 0 deletions
Show Stats Download Patch File Download Diff File

502
model_zoo/official/cv/centerface/README.MD
View File

File diff suppressed because it is too large Load Diff

89
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/datasets/dataset/coco_hp.py
Unescape View File

@ -0,0 +1,89 @@
"""
MIT License
Copyright (c) 2019 Xingyi Zhou
All rights reserved.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import numpy as np
import pycocotools.coco as coco
import cv2
class CenterfaceDataset():
"""
Centerface dataset definition.
"""
def __init__(self, config, split='train'):
self.split = split
self.config = config
self.max_objs = config.max_objs
self.img_dir = self.config.img_dir
self.annot_path = self.config.annot_path
print('==> getting centerface key point {} data.'.format(split))
self.coco = coco.COCO(self.annot_path)
image_ids = self.coco.getImgIds()
if split == 'train':
self.images = []
for img_id in image_ids:
idxs = self.coco.getAnnIds(imgIds=[img_id])
if idxs:
self.images.append(img_id)
else:
self.images = image_ids
self.num_samples = len(self.images)
print('Loaded {} {} samples'.format(split, self.num_samples)) # Loaded train 12671 samples
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
(image, target) (tuple): target is index of the target class.
"""
img_id = self.images[index]
file_name = self.coco.loadImgs(ids=[img_id])[0]['file_name']
img_path = os.path.join(self.img_dir, file_name)
ann_ids = self.coco.getAnnIds(imgIds=[img_id])
anns = self.coco.loadAnns(ids=ann_ids)
num_objs = len(anns)
if num_objs > self.max_objs:
num_objs = self.max_objs
anns = np.random.choice(anns, num_objs)
# dataType ERROR —— to_list
target = []
for ann in anns:
tmp = []
tmp.extend(ann['bbox'])
tmp.extend(ann['keypoints'])
target.append(tmp)
img = cv2.imread(img_path)
return img, target
def __len__(self):
return self.num_samples

217
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/datasets/sample/multi_pose.py
Unescape View File

@ -0,0 +1,217 @@
"""
MIT License
Copyright (c) 2019 Xingyi Zhou
All rights reserved.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import numpy as np
import cv2
from dependency.centernet.src.lib.utils.image import color_aug
from dependency.centernet.src.lib.utils.image import get_affine_transform, affine_transform
from dependency.centernet.src.lib.utils.image import gaussian_radius, draw_umich_gaussian
from dependency.extd.utils.augmentations import anchor_crop_image_sampling
def get_border(border, size):
"""
Get border
"""
i = 1
while size - border // i <= border // i: # size > 2 * (border // i)
i *= 2
return border // i
def coco_box_to_bbox(box):
"""
(x1, y1, w, h) -> (x1, y1, x2, y2)
"""
bbox = np.array([box[0], box[1], box[0] + box[2], box[1] + box[3]], dtype=np.float32)
return bbox
def preprocess_train(image, target, config):
"""
Preprocess training data
"""
data_rng = np.random.RandomState(123)
eig_val = np.array([0.2141788, 0.01817699, 0.00341571], dtype=np.float32)
eig_vec = np.array([
[-0.58752847, -0.69563484, 0.41340352],
[-0.5832747, 0.00994535, -0.81221408],
[-0.56089297, 0.71832671, 0.41158938]
], dtype=np.float32)
mean = np.array([0.40789654, 0.44719302, 0.47026115], dtype=np.float32).reshape((1, 1, 3))
std = np.array([0.28863828, 0.27408164, 0.27809835], dtype=np.float32).reshape((1, 1, 3))
num_objs = len(target)
anns = []
for each in target:
ann = {}
ann['bbox'] = each[0:4]
ann['keypoints'] = each[4:]
anns.append(ann)
cv2.setNumThreads(0)
img, anns = anchor_crop_image_sampling(image, anns)
_, width = img.shape[0], img.shape[1]
c = np.array([img.shape[1] / 2., img.shape[0] / 2.], dtype=np.float32)
s = max(img.shape[0], img.shape[1]) * 1.0
rot = 0
flipped = False
if config.rand_crop:
#s = s * np.random.choice(np.arange(0.8, 1.3, 0.05)) # for 768*768 or 800* 800
s = s * np.random.choice(np.arange(0.6, 1.0, 0.05)) # for 512 * 512
border = s * np.random.choice([0.1, 0.2, 0.25])
w_border = get_border(border, img.shape[1]) # w > 2 * w_border
h_border = get_border(border, img.shape[0]) # h > 2 * h_border
c[0] = np.random.randint(low=w_border, high=img.shape[1] - w_border)
c[1] = np.random.randint(low=h_border, high=img.shape[0] - h_border)
else:
sf = config.scale
cf = config.shift
c[0] += s * np.clip(np.random.randn() * cf, -2 * cf, 2 * cf)
c[1] += s * np.clip(np.random.randn() * cf, -2 * cf, 2 * cf)
s = s * np.clip(np.random.randn() * sf + 1, 1 - sf, 1 + sf)
if np.random.random() < config.rotate:
rf = config.rotate
rot = np.clip(np.random.randn() * rf, -rf * 2, rf * 2)
if np.random.random() < config.flip: # opt.flip = 0.5
flipped = True
img = img[:, ::-1, :]
c[0] = width - c[0] - 1
trans_input = get_affine_transform(c, s, rot, [config.input_res, config.input_res])
inp = cv2.warpAffine(img, trans_input, (config.input_res, config.input_res), flags=cv2.INTER_LINEAR)
inp = (inp.astype(np.float32) / 255.)
if config.color_aug:
color_aug(data_rng, inp, eig_val, eig_vec)
inp = (inp - mean) / std
inp = inp.transpose(2, 0, 1)
output_res = config.output_res
num_joints = config.num_joints
max_objs = config.max_objs
trans_output_rot = get_affine_transform(c, s, rot, [output_res, output_res])
trans_output = get_affine_transform(c, s, 0, [output_res, output_res])
# map
hm = np.zeros((config.num_classes, output_res, output_res), dtype=np.float32)
hm_hp = np.zeros((num_joints, output_res, output_res), dtype=np.float32)
wh = np.zeros((output_res, output_res, 2), dtype=np.float32)
reg = np.zeros((output_res, output_res, 2), dtype=np.float32)
ind = np.zeros((output_res, output_res), dtype=np.float32) # as float32, need no data_type change later
reg_mask = np.zeros((max_objs), dtype=np.uint8)
wight_mask = np.zeros((output_res, output_res, 2), dtype=np.float32)
kps = np.zeros((output_res, output_res, num_joints * 2), dtype=np.float32)
kps_mask = np.zeros((output_res, output_res, num_joints * 2), dtype=np.float32)
#
hp_offset = np.zeros((max_objs * num_joints, 2), dtype=np.float32)
hp_ind = np.zeros((max_objs * num_joints), dtype=np.int64)
hp_mask = np.zeros((max_objs * num_joints), dtype=np.int64)
draw_gaussian = draw_umich_gaussian
gt_det = []
for k in range(num_objs):
ann = anns[k]
bbox = coco_box_to_bbox(ann['bbox']) # [x,y,w,h]--[x1,y1,x2,y2]
cls_id = 0 #int(ann['category_id']) - 1
pts = np.array(ann['keypoints'], np.float32).reshape(num_joints, 3) # (x,y,0/1)
if flipped:
bbox[[0, 2]] = width - bbox[[2, 0]] - 1
pts[:, 0] = width - pts[:, 0] - 1
for e in config.flip_idx: # flip_idx = [[0, 1], [3, 4]]
pts[e[0]], pts[e[1]] = pts[e[1]].copy(), pts[e[0]].copy()
bbox[:2] = affine_transform(bbox[:2], trans_output) # [0, 1] -- (x1, y1)
bbox[2:] = affine_transform(bbox[2:], trans_output) # [2, 3] -- (x2, y2)
bbox = np.clip(bbox, 0, output_res - 1)
h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
if (h > 0 and w > 0) or (rot != 0):
radius = gaussian_radius((math.ceil(h), math.ceil(w)))
radius = max(0, int(radius))
ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
ct_int = ct.astype(np.int32)
ind[ct_int[1], ct_int[0]] = 1.0
wh[ct_int[1], ct_int[0], :] = np.log(1. * w / 4), np.log(1. * h / 4)
reg[ct_int[1], ct_int[0], :] = ct[0] - ct_int[0], ct[1] - ct_int[1]
reg_mask[k] = 1.0
wight_mask[ct_int[1], ct_int[0], 0] = 1
wight_mask[ct_int[1], ct_int[0], 1] = 1
# if w*h <= 20: # can get what we want sometime, but unstable
# wight_mask[k] = 15
if w*h <= 40:
wight_mask[ct_int[1], ct_int[0], 0] = 5
wight_mask[ct_int[1], ct_int[0], 1] = 5
if w*h <= 20:
wight_mask[ct_int[1], ct_int[0], 0] = 10
wight_mask[ct_int[1], ct_int[0], 1] = 10
if w*h <= 10:
wight_mask[ct_int[1], ct_int[0], 0] = 15
wight_mask[ct_int[1], ct_int[0], 1] = 15
if w*h <= 4:
wight_mask[ct_int[1], ct_int[0], 0] = 0.1
wight_mask[ct_int[1], ct_int[0], 1] = 0.1
num_kpts = pts[:, 2].sum()
if num_kpts == 0:
hm[cls_id, ct_int[1], ct_int[0]] = 0.9999
hp_radius = gaussian_radius((math.ceil(h), math.ceil(w)))
hp_radius = max(0, int(hp_radius))
for j in range(num_joints):
if pts[j, 2] > 0:
pts[j, :2] = affine_transform(pts[j, :2], trans_output_rot)
if pts[j, 0] >= 0 and pts[j, 0] < output_res and pts[j, 1] >= 0 and pts[j, 1] < output_res:
kps[ct_int[1], ct_int[0], j * 2 : j * 2 + 2] = pts[j, :2] - ct_int
kps[ct_int[1], ct_int[0], j * 2 : j * 2 + 1] = kps[ct_int[1], ct_int[0], j * 2 : j * 2 + 1] / w
kps[ct_int[1], ct_int[0], j * 2 + 1: j * 2 + 2] = kps[ct_int[1], ct_int[0],
j * 2 + 1 : j * 2 + 2] / h
kps_mask[ct_int[1], ct_int[0], j * 2 : j * 2 + 2] = 1.0
pt_int = pts[j, :2].astype(np.int32)
hp_offset[k * num_joints + j] = pts[j, :2] - pt_int
hp_ind[k * num_joints + j] = pt_int[1] * output_res + pt_int[0]
hp_mask[k * num_joints + j] = 1
draw_gaussian(hm_hp[j], pt_int, hp_radius)
kps_mask[ct_int[1], ct_int[0], j * 2 : j * 2 + 2] = \
0.0 if ann['bbox'][2] * ann['bbox'][3] <= 8.0 else 1.0
draw_gaussian(hm[cls_id], ct_int, radius)
gt_det.append([ct[0] - w / 2, ct[1] - h / 2,
ct[0] + w / 2, ct[1] + h / 2, 1] +
pts[:, :2].reshape(num_joints * 2).tolist() + [cls_id])
return inp, hm, reg_mask, ind, wh, wight_mask, reg, kps_mask, kps

216
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/detectors/base_detector.py
Unescape View File

@ -0,0 +1,216 @@
###modified based on centernet###
#MIT License
#Copyright (c) 2019 Xingyi Zhou
#All rights reserved.
"""Basic definition of detector"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import cv2
from mindspore import Tensor
from dependency.centernet.src.lib.external.nms import soft_nms
from dependency.centernet.src.lib.utils.image import get_affine_transform, affine_transform
def transform_preds(coords, center, scale, output_size):
"""
Transform target coords
"""
target_coords = np.zeros(coords.shape)
trans = get_affine_transform(center, scale, 0, output_size, inv=1)
for p in range(coords.shape[0]):
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
return target_coords
def multi_pose_post_process(dets, c, s, h, w):
"""
Multi pose post process
dets_result: 4 + score:1 + kpoints:10 + class:1 = 16
dets: batch x max_dets x 40
return list of 39 in image coord
"""
ret = []
for i in range(dets.shape[0]):
bbox = transform_preds(dets[i, :, :4].reshape(-1, 2), c[i], s[i], (w, h))
pts = transform_preds(dets[i, :, 5:15].reshape(-1, 2), c[i], s[i], (w, h))
top_preds = np.concatenate([bbox.reshape(-1, 4), dets[i, :, 4:5], pts.reshape(-1, 10)],
axis=1).astype(np.float32).tolist()
ret.append({np.ones(1, dtype=np.int32)[0]: top_preds})
return ret
class CenterFaceDetector():
"""
Centerface detector
"""
def __init__(self, opt, model):
self.flip_idx = opt.flip_idx
print('Creating model...')
self.model = model
self.mean = np.array(opt.mean, dtype=np.float32).reshape((1, 1, 3))
self.std = np.array(opt.std, dtype=np.float32).reshape((1, 1, 3))
self.max_per_image = 100
self.num_classes = opt.num_classes
self.scales = opt.test_scales
self.opt = opt
self.pause = False
def pre_process(self, image, scale, meta=None):
"""
Preprocess method
"""
height, width = image.shape[0:2]
new_height = int(height * scale)
new_width = int(width * scale)
if self.opt.fix_res: # True
inp_height, inp_width = self.opt.input_h, self.opt.input_w
c = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
s = max(height, width) * 1.0
else:
inp_height = int(np.ceil(new_height / 32) * 32)
inp_width = int(np.ceil(new_width / 32) * 32)
c = np.array([new_width // 2, new_height // 2], dtype=np.float32)
s = np.array([inp_width, inp_height], dtype=np.float32)
trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
resized_image = cv2.resize(image, (new_width, new_height))
inp_image = cv2.warpAffine(
resized_image, trans_input, (inp_width, inp_height),
flags=cv2.INTER_LINEAR)
inp_image = ((inp_image / 255. - self.mean) / self.std).astype(np.float32)
images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)
if self.opt.flip_test:
images = np.concatenate((images, images[:, :, :, ::-1]), axis=0)
meta = {'c': c, 's': s, 'out_height': inp_height // self.opt.down_ratio,
'out_width': inp_width // self.opt.down_ratio}
return images, meta
def process(self, images):
"""
Process method
"""
images = Tensor(images)
# test with mindspore model
output_hm, output_wh, output_off, output_kps, topk_inds = self.model(images)
# Tensor to numpy
output_hm = output_hm.asnumpy().astype(np.float32)
output_wh = output_wh.asnumpy().astype(np.float32)
output_off = output_off.asnumpy().astype(np.float32)
output_kps = output_kps.asnumpy().astype(np.float32)
topk_inds = topk_inds.asnumpy().astype(np.long)
reg = output_off if self.opt.reg_offset else None
dets = self.centerface_decode(output_hm, output_wh, output_kps, reg=reg, opt_k=self.opt.K, topk_inds=topk_inds)
return dets
def post_process(self, dets, meta, scale=1):
"""
Post process process
"""
dets = dets.reshape(1, -1, dets.shape[2])
dets = multi_pose_post_process(
dets.copy(), [meta['c']], [meta['s']],
meta['out_height'], meta['out_width'])
for j in range(1, self.num_classes + 1):
dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 15)
# import pdb; pdb.set_trace()
dets[0][j][:, :4] /= scale
dets[0][j][:, 5:] /= scale
return dets[0]
def merge_outputs(self, detections):
"""
Merge detection outputs
"""
results = {}
results[1] = np.concatenate([detection[1] for detection in detections], axis=0).astype(np.float32)
if self.opt.nms or len(self.opt.test_scales) > 1:
soft_nms(results[1], Nt=0.5, method=2)
results[1] = results[1].tolist()
return results
def run(self, image_or_path_or_tensor, meta=None):
"""
Run method
"""
pre_processed = False
if isinstance(image_or_path_or_tensor, np.ndarray):
image = image_or_path_or_tensor
elif isinstance(image_or_path_or_tensor, str):
image = cv2.imread(image_or_path_or_tensor)
else:
image = image_or_path_or_tensor['image'][0].numpy()
pre_processed_images = image_or_path_or_tensor
pre_processed = True
detections = []
for scale in self.scales: # [1]
if not pre_processed:
images, meta = self.pre_process(image, scale, meta) # --1: pre_process
else:
# import pdb; pdb.set_trace()
images = pre_processed_images['images'][scale][0]
meta = pre_processed_images['meta'][scale]
meta = {k: v.numpy()[0] for k, v in meta.items()}
dets = self.process(images) # --2: process
dets = self.post_process(dets, meta, scale) # box:4+score:1+kpoints:10+class:1=16 ## --3: post_process
detections.append(dets)
results = self.merge_outputs(detections) # --4: merge_outputs
return {'results': results}
def centerface_decode(self, heat, wh, kps, reg=None, opt_k=100, topk_inds=None):
"""
Decode detection bbox
"""
batch, _, _, width = wh.shape
num_joints = kps.shape[1] // 2
scores = heat
inds = topk_inds
ys_int = (topk_inds / width).astype(np.int32).astype(np.float32)
xs_int = (topk_inds % width).astype(np.int32).astype(np.float32)
reg = reg.reshape(batch, 2, -1)
reg_tmp = np.zeros((batch, 2, opt_k), dtype=np.float32)
for i in range(batch):
reg_tmp[i, 0, :] = reg[i, 0, inds[i]]
reg_tmp[i, 1, :] = reg[i, 1, inds[i]]
reg = reg_tmp.transpose(0, 2, 1)
if reg is not None:
xs = xs_int.reshape(batch, opt_k, 1) + reg[:, :, 0:1]
ys = ys_int.reshape(batch, opt_k, 1) + reg[:, :, 1:2]
else:
xs = xs_int.reshape(batch, opt_k, 1) + 0.5
ys = ys_int.reshape(batch, opt_k, 1) + 0.5
wh = wh.reshape(batch, 2, -1)
wh_tmp = np.zeros((batch, 2, opt_k), dtype=np.float32)
for i in range(batch):
wh_tmp[i, 0, :] = wh[i, 0, inds[i]]
wh_tmp[i, 1, :] = wh[i, 1, inds[i]]
wh = wh_tmp.transpose(0, 2, 1)
wh = np.exp(wh) * 4.
scores = scores.reshape(batch, opt_k, 1)
bboxes = np.concatenate([xs - wh[..., 0:1] / 2, ys - wh[..., 1:2] / 2, xs + wh[..., 0:1] / 2,
ys + wh[..., 1:2] / 2], axis=2)
clses = np.zeros((batch, opt_k, 1), dtype=np.float32)
kps = np.zeros((batch, opt_k, num_joints * 2), dtype=np.float32)
detections = np.concatenate([bboxes, scores, kps, clses], axis=2) # box:4+score:1+kpoints:10+class:1=16
return detections

0
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/external/__init__.py vendored
Unescape View File

391
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/external/nms.pyx vendored
View File

File diff suppressed because it is too large Load Diff

42
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/external/setup.py vendored
Unescape View File

@ -0,0 +1,42 @@
"""
MIT License
Copyright (c) 2019 Xingyi Zhou
All rights reserved.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
from distutils.core import setup
from distutils.extension import Extension
import numpy
from Cython.Build import cythonize
extensions = [
Extension(
"nms",
["nms.pyx"],
extra_compile_args=["-Wno-cpp", "-Wno-unused-function"]
)
]
setup(
name="coco",
ext_modules=cythonize(extensions),
include_dirs=[numpy.get_include()]
)

170
model_zoo/official/cv/centerface/denpendency/centernet/src/lib/utils/image.py
Unescape View File

@ -0,0 +1,170 @@
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# Modified by Xingyi Zhou
# ------------------------------------------------------------------------------
"""Image process"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import random
import numpy as np
import cv2
def get_3rd_point(a, b):
"""
Get 3rd point
"""
direct = a - b
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
def get_dir(src_point, rot_rad):
"""
Get dir
"""
sn, cs = np.sin(rot_rad), np.cos(rot_rad) # (0, 1)
src_result = [0, 0]
src_result[0] = src_point[0] * cs - src_point[1] * sn
src_result[1] = src_point[0] * sn + src_point[1] * cs
return src_result
def get_affine_transform(center,
scale,
rot,
output_size,
shift=np.array([0, 0], dtype=np.float32),
inv=0):
"""
Get affine transform
"""
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
scale = np.array([scale, scale], dtype=np.float32)
scale_tmp = scale
src_w = scale_tmp[0]
dst_w = output_size[0]
dst_h = output_size[1]
rot_rad = np.pi * rot / 180
src_dir = get_dir([0, src_w * -0.5], rot_rad)
dst_dir = np.array([0, dst_w * -0.5], np.float32)
src = np.zeros((3, 2), dtype=np.float32)
dst = np.zeros((3, 2), dtype=np.float32)
src[0, :] = center + scale_tmp * shift
src[1, :] = center + src_dir + scale_tmp * shift
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
if inv:
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
else:
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
return trans
def affine_transform(pt, t):
"""
Affine transform
"""
new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T
new_pt = np.dot(t, new_pt)
return new_pt[:2]
def grayscale(image):
return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
def lighting_(data_rng, image, alphastd, eigval, eigvec):
alpha = data_rng.normal(scale=alphastd, size=(3,))
image += np.dot(eigvec, eigval * alpha)
def blend_(alpha, image1, image2):
image1 *= alpha
image2 *= (1 - alpha)
image1 += image2
def saturation_(data_rng, image, gs, gs_mean, var):
gs_mean = gs_mean
alpha = 1. + data_rng.uniform(low=-var, high=var)
blend_(alpha, image, gs[:, :, None])
def brightness_(data_rng, image, gs, gs_mean, var):
gs = gs
gs_mean = gs_mean
alpha = 1. + data_rng.uniform(low=-var, high=var)
image *= alpha
def contrast_(data_rng, image, gs, gs_mean, var):
gs = gs
alpha = 1. + data_rng.uniform(low=-var, high=var)
blend_(alpha, image, gs_mean)
def color_aug(data_rng, image, eig_val, eig_vec):
functions = [brightness_, contrast_, saturation_]
random.shuffle(functions)
gs = grayscale(image)
gs_mean = gs.mean()
for f in functions:
f(data_rng, image, gs, gs_mean, 0.4)
lighting_(data_rng, image, 0.1, eig_val, eig_vec)
def gaussian_radius(det_size, min_overlap=0.7):
"""
Gaussian radius
"""
height, width = det_size
a1 = 1
b1 = (height + width)
c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
r1 = (b1 + sq1) / 2
a2 = 4
b2 = 2 * (height + width)
c2 = (1 - min_overlap) * width * height
sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
r2 = (b2 + sq2) / 2
a3 = 4 * min_overlap
b3 = -2 * min_overlap * (height + width)
c3 = (min_overlap - 1) * width * height
sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
r3 = (b3 + sq3) / 2
return min(r1, r2, r3)
def gaussian2d(shape, sigma=1):
m, n = [(ss - 1.) / 2. for ss in shape]
y, x = np.ogrid[-m:m+1, -n:n+1]
h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
h[h < np.finfo(h.dtype).eps * h.max()] = 0
return h
def draw_umich_gaussian(heatmap, center, radius, k=1):
"""
Draw umich gaussian
"""
diameter = 2 * radius + 1
gaussian = gaussian2d((diameter, diameter), sigma=diameter / 6)
x, y = int(center[0]), int(center[1])
height, width = heatmap.shape[0:2]
left, right = min(x, radius), min(width - x, radius + 1)
top, bottom = min(y, radius), min(height - y, radius + 1)
masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:
np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
return heatmap

55
model_zoo/official/cv/centerface/denpendency/evaluate/box_overlaps.pyx
Unescape View File

@ -0,0 +1,55 @@
# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Sergey Karayev
# --------------------------------------------------------
cimport cython
import numpy as np
cimport numpy as np
DTYPE = np.float
ctypedef np.float_t DTYPE_t
def bbox_overlaps(
np.ndarray[DTYPE_t, ndim=2] boxes,
np.ndarray[DTYPE_t, ndim=2] query_boxes):
"""
Parameters
----------
boxes: (N, 4) ndarray of float
query_boxes: (K, 4) ndarray of float
Returns
-------
overlaps: (N, K) ndarray of overlap between boxes and query_boxes
"""
cdef unsigned int N = boxes.shape[0]
cdef unsigned int K = query_boxes.shape[0]
cdef np.ndarray[DTYPE_t, ndim=2] overlaps = np.zeros((N, K), dtype=DTYPE)
cdef DTYPE_t iw, ih, box_area
cdef DTYPE_t ua
cdef unsigned int k, n
for k in range(K):
box_area = (
(query_boxes[k, 2] - query_boxes[k, 0] + 1) *
(query_boxes[k, 3] - query_boxes[k, 1] + 1)
)
for n in range(N):
iw = (
min(boxes[n, 2], query_boxes[k, 2]) -
max(boxes[n, 0], query_boxes[k, 0]) + 1
)
if iw > 0:
ih = (
min(boxes[n, 3], query_boxes[k, 3]) -
max(boxes[n, 1], query_boxes[k, 1]) + 1
)
if ih > 0:
ua = float(
(boxes[n, 2] - boxes[n, 0] + 1) *
(boxes[n, 3] - boxes[n, 1] + 1) +
box_area - iw * ih
)
overlaps[n, k] = iw * ih / ua
return overlaps

316
model_zoo/official/cv/centerface/denpendency/evaluate/eval.py
View File

File diff suppressed because it is too large Load Diff

13
model_zoo/official/cv/centerface/denpendency/evaluate/setup.py
Unescape View File

@ -0,0 +1,13 @@
"""
WiderFace evaluation code
author: wondervictor
mail: tianhengcheng@gmail.com
copyright@wondervictor
"""
from distutils.core import setup, Extension
import numpy
from Cython.Build import cythonize
package = Extension('bbox', ['box_overlaps.pyx'], include_dirs=[numpy.get_include()])
setup(ext_modules=cythonize([package]))

78
model_zoo/official/cv/centerface/denpendency/extd/utils/augmentations.py
Unescape View File

@ -0,0 +1,78 @@
#EXTD: Extremely Tiny Face Detector via Iterative Filter Reuse
# MIT license
# Copyright (c) 2019-present NAVER Corp.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE
"""Augmentations"""
import random
import numpy as np
import cv2
def anchor_crop_image_sampling(image, anns):
"""
Crop anchors.
"""
max_size = 12000
inf_distance = 9999999
boxes = []
for ann in anns:
boxes.append([ann['bbox'][0], ann['bbox'][1], ann['bbox'][0] + ann['bbox'][2], ann['bbox'][1] + ann['bbox'][3]])
boxes = np.asarray(boxes, dtype=np.float32)
height, width, _ = image.shape
box_area = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
rand_idx = random.randint(0, len(box_area) - 1)
rand_side = box_area[rand_idx] ** 0.5
anchors = [16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 128, 256, 512]
distance = inf_distance
anchor_idx = 5
for i, anchor in enumerate(anchors):
if abs(anchor - rand_side) < distance:
distance = abs(anchor - rand_side)
anchor_idx = i
target_anchor = random.choice(anchors[0:min(anchor_idx + 1, 11)])
ratio = float(target_anchor) / rand_side
ratio = ratio * (2 ** random.uniform(-1, 1))
if int(height * ratio * width * ratio) > max_size * max_size:
ratio = (max_size * max_size / (height * width)) ** 0.5
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
interp_method = random.choice(interp_methods)
image = cv2.resize(image, None, None, fx=ratio, fy=ratio, interpolation=interp_method)
boxes[:, 0] *= ratio
boxes[:, 1] *= ratio
boxes[:, 2] *= ratio
boxes[:, 3] *= ratio
boxes = boxes.tolist()
for i, _ in enumerate(anns):
anns[i]['bbox'] = [boxes[i][0], boxes[i][1], boxes[i][2] - boxes[i][0], boxes[i][3] - boxes[i][1]]
for j in range(5):
anns[i]['keypoints'][j * 3] *= ratio
anns[i]['keypoints'][j * 3 + 1] *= ratio
return image, anns

62
model_zoo/official/cv/centerface/export.py
Unescape View File

@ -0,0 +1,62 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Convert ckpt to air."""
import os
import argparse
import numpy as np
from mindspore import context
from mindspore import Tensor
from mindspore.train.serialization import export, load_checkpoint, load_param_into_net
from src.centerface import CenterfaceMobilev2
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=False)
def save_air():
"""Save air file"""
print('============= centerface start save air ==================')
parser = argparse.ArgumentParser(description='Convert ckpt to air')
parser.add_argument('--pretrained', type=str, default='', help='pretrained model to load')
parser.add_argument('--batch_size', type=int, default=8, help='batch size')
args = parser.parse_args()
network = CenterfaceMobilev2()
if os.path.isfile(args.pretrained):
param_dict = load_checkpoint(args.pretrained)
param_dict_new = {}
for key, values in param_dict.items():
if key.startswith('moments.') or key.startswith('moment1.') or key.startswith('moment2.'):
continue
elif key.startswith('centerface_network.'):
param_dict_new[key[19:]] = values
else:
param_dict_new[key] = values
load_param_into_net(network, param_dict_new)
print('load model {} success'.format(args.pretrained))
input_data = np.random.uniform(low=0, high=1.0, size=(args.batch_size, 3, 832, 832)).astype(np.float32)
tensor_input_data = Tensor(input_data)
export(network, tensor_input_data,
file_name=args.pretrained.replace('.ckpt', '_' + str(args.batch_size) + 'b.air'), file_format='AIR')
print("export model success.")
if __name__ == "__main__":
save_air()

22
model_zoo/official/cv/centerface/scripts/eval.sh
Unescape View File

@ -0,0 +1,22 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
root=$PWD
save_path=$root/output/centerface/
ground_truth_path=$root/dataset/centerface/ground_truth
echo "start eval"
python ../dependency/evaluate/eval.py --pred=$save_path --gt=$ground_truth_path
echo "end eval"

26
model_zoo/official/cv/centerface/scripts/eval_all.sh
Unescape View File

@ -0,0 +1,26 @@
#!/bin/sh
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
root=$PWD
save_path=$root/output/centerface/
ground_truth_path=$root/dataset/centerface/ground_truth
#for i in $(seq start_epoch end_epoch+1)
for i in $(seq 89 200)
do
python ../dependency/evaluate/eval.py --pred=$save_path$i --gt=$ground_truth_path &
sleep 10
done
wait

131
model_zoo/official/cv/centerface/scripts/test.sh
Unescape View File

@ -0,0 +1,131 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [ $# -gt 6 ]
then
echo "Usage: sh test.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_ID] [CKPT]"
echo " or: sh test.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_ID]"
echo " or: sh test.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH]"
echo " or: sh test.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT]"
echo " or: sh test.sh [MODEL_PATH] [DATASET]"
echo " or: sh test.sh [MODEL_PATH]"
echo " or: sh test.sh "
exit 1
fi
get_real_path(){
if [ "${1:0:1}" == "/" ]; then
echo "$1"
else
echo "$(realpath -m $PWD/$1)"
fi
}
current_exec_path=$(pwd)
echo ${current_exec_path}
dirname_path=$(dirname "$(pwd)")
echo ${dirname_path}
SCRIPT_NAME='test.py'
ulimit -c unlimited
root=${current_exec_path} # your script path
model_path=$root/model/
dataset_root=$root/dataset
dataset_path=$dataset_root/centerface/images/val/images/
ground_truth_mat=$dataset_root/centerface/ground_truth/val.mat
save_path=$root/output/centerface/
device_id=0
ckpt="0.ckpt" # the model saved for epoch=125
if [ $# == 1 ]
then
model_path=$(get_real_path $1)
if [ ! -f $model_path ]
then
echo "error: model_path=$model_path is not a file"
exit 1
fi
fi
if [ $# == 2 ]
then
dataset_path=$(get_real_path $2)
if [ ! -f $dataset_path ]
then
echo "error: dataset_path=$dataset_path is not a file"
exit 1
fi
fi
if [ $# == 3 ]
then
ground_truth_mat=$(get_real_path $3)
if [ ! -f $ground_truth_mat ]
then
echo "error: ground_truth_mat=$ground_truth_mat is not a file"
exit 1
fi
fi
if [ $# == 4 ]
then
save_path=$(get_real_path $4)
if [ ! -f $save_path ]
then
echo "error: save_path=$save_path is not a file"
exit 1
fi
fi
if [ $# == 5 ]
then
device_id=$5
fi
if [ $# == 6 ]
then
ckpt=$6
fi
echo $model_path
echo $dataset_path
echo $ground_truth_mat
echo $save_path
export PYTHONPATH=${dirname_path}:$PYTHONPATH
export RANK_SIZE=1
echo 'start testing'
rm -rf ${current_exec_path}/device_test$device_id
echo 'start rank '$device_id
mkdir ${current_exec_path}/device_test$device_id
cd ${current_exec_path}/device_test$device_id || exit
export RANK_ID=0
dev=`expr $device_id + 0`
export DEVICE_ID=$dev
python ${dirname_path}/${SCRIPT_NAME} \
--is_distributed=0 \
--data_dir=$dataset_path \
--test_model=$model_path \
--ground_truth_mat=$ground_truth_mat \
--save_dir=$save_path \
--rank=$device_id \
--ckpt_name=$ckpt > test.log 2>&1 &
echo 'running'

146
model_zoo/official/cv/centerface/scripts/test_and_eval.sh
Unescape View File

@ -0,0 +1,146 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [ $# -gt 6 ]
then
echo "Usage: sh test_and_eval.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_ID] [CKPT] [GROUND_TRUTH_PATH]"
echo " or: sh test_and_eval.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_ID] [CKPT]"
echo " or: sh test_and_eval.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_ID]"
echo " or: sh test_and_eval.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH]"
echo " or: sh test_and_eval.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT]"
echo " or: sh test_and_eval.sh [MODEL_PATH] [DATASET]"
echo " or: sh test_and_eval.sh [MODEL_PATH]"
echo " or: sh test_and_eval.sh "
exit 1
fi
get_real_path(){
if [ "${1:0:1}" == "/" ]; then
echo "$1"
else
echo "$(realpath -m $PWD/$1)"
fi
}
current_exec_path=$(pwd)
echo ${current_exec_path}
dirname_path=$(dirname "$(pwd)")
echo ${dirname_path}
SCRIPT_NAME='test.py'
ulimit -c unlimited
root=${current_exec_path} # your script path
model_path=$root/model/
dataset_root=$root/dataset
dataset_path=$dataset_root/centerface/images/val/images/
ground_truth_mat=$dataset_root/centerface/ground_truth/val.mat
save_path=$root/output/centerface/999
device_id=0
ckpt="0-125_24750.ckpt" # the model saved for epoch=125
ground_truth_path=$root/dataset/centerface/ground_truth
if [ $# == 1 ]
then
model_path=$(get_real_path $1)
if [ ! -f $model_path ]
then
echo "error: model_path=$model_path is not a file"
exit 1
fi
fi
if [ $# == 2 ]
then
dataset_path=$(get_real_path $2)
if [ ! -f $dataset_path ]
then
echo "error: dataset_path=$dataset_path is not a file"
exit 1
fi
fi
if [ $# == 3 ]
then
ground_truth_mat=$(get_real_path $3)
if [ ! -f $ground_truth_mat ]
then
echo "error: ground_truth_mat=$ground_truth_mat is not a file"
exit 1
fi
fi
if [ $# == 4 ]
then
save_path=$(get_real_path $4)
if [ ! -f $save_path ]
then
echo "error: save_path=$save_path is not a file"
exit 1
fi
fi
if [ $# == 5 ]
then
device_id=$5
fi
if [ $# == 6 ]
then
ckpt=$6
fi
if [ $# == 7 ]
then
ground_truth_path=$(get_real_path $7)
if [ ! -f $ground_truth_path ]
then
echo "error: ground_truth_path=$ground_truth_path is not a file"
exit 1
fi
fi
echo $model_path
echo $dataset_path
echo $ground_truth_mat
echo $save_path
echo $ground_truth_path
export PYTHONPATH=${dirname_path}:$PYTHONPATH
export RANK_SIZE=1
echo 'start testing'
rm -rf ${current_exec_path}/device_test$device_id
echo 'start rank '$device_id
mkdir ${current_exec_path}/device_test$device_id
cd ${current_exec_path}/device_test$device_id || exit
export RANK_ID=0
dev=`expr $device_id + 0`
export DEVICE_ID=$dev
python ${dirname_path}/${SCRIPT_NAME} \
--is_distributed=0 \
--data_dir=$dataset_path \
--test_model=$model_path \
--ground_truth_mat=$ground_truth_mat \
--save_dir=$save_path \
--rank=$device_id \
--ckpt_name=$ckpt \
--eval=1 \
--ground_truth_path=$ground_truth_path > test.log 2>&1 &
echo 'running'

157
model_zoo/official/cv/centerface/scripts/test_distribute.sh
Unescape View File

@ -0,0 +1,157 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [ $# -gt 8 ]
then
echo "Usage: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_NUM] [STEPS_PER_EPOCH] [START] [END]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_NUM] [STEPS_PER_EPOCH] [START]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_NUM] [STEPS_PER_EPOCH]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_NUM]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH] [DEVICE_NUM]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT] [SAVE_PATH]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET] [GROUND_TRUTH_MAT]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET]"
echo " or: sh test_distribute.sh [MODEL_PATH] [DATASET]"
echo " or: sh test_distribute.sh [MODEL_PATH]"
echo " or: sh test_distribute.sh "
exit 1
fi
get_real_path(){
if [ "${1:0:1}" == "/" ]; then
echo "$1"
else
echo "$(realpath -m $PWD/$1)"
fi
}
current_exec_path=$(pwd)
echo ${current_exec_path}
dirname_path=$(dirname "$(pwd)")
echo ${dirname_path}
SCRIPT_NAME='test.py'
ulimit -c unlimited
root=${current_exec_path} # your script path
model_path=$root/model/
dataset_root=$root/dataset
dataset_path=$dataset_root/centerface/images/val/images/
ground_truth_mat=$dataset_root/centerface/ground_truth/val.mat
save_path=$root/output/centerface/
# blow are used for calculate model name
# model/ckpt name is "0-" + str(ckpt_num) + "_" + str(198*ckpt_num) + ".ckpt";
# ckpt_num is epoch number, can be calculated by device_num
# detail can be found in "test.py"
device_num=8
steps_per_epoch=198 #198 for 8P; 1583 for 1p
start=11 # start epoch number = start * device_num + min(device_phy_id) + 1
end=18 # end epoch number = end * device_num + max(device_phy_id) + 1
if [ $# == 1 ]
then
model_path=$(get_real_path $1)
if [ ! -f $model_path ]
then
echo "error: model_path=$model_path is not a file"
exit 1
fi
fi
if [ $# == 2 ]
then
dataset_path=$(get_real_path $2)
if [ ! -f $dataset_path ]
then
echo "error: dataset_path=$dataset_path is not a file"
exit 1
fi
fi
if [ $# == 3 ]
then
ground_truth_mat=$(get_real_path $3)
if [ ! -f $ground_truth_mat ]
then
echo "error: ground_truth_mat=$ground_truth_mat is not a file"
exit 1
fi
fi
if [ $# == 4 ]
then
save_path=$(get_real_path $4)
if [ ! -f $save_path ]
then
echo "error: save_path=$save_path is not a file"
exit 1
fi
fi
if [ $# == 5 ]
then
device_num=$5
fi
if [ $# == 6 ]
then
steps_per_epoch=$6
fi
if [ $# == 7 ]
then
start=$7
fi
if [ $# == 8 ]
then
end=$8
fi
echo $model_path
echo $dataset_path
echo $ground_truth_mat
echo $save_path
export PYTHONPATH=${dirname_path}:$PYTHONPATH
export RANK_SIZE=1
echo 'start testing'
rm -rf ${current_exec_path}/device_test*
for((i=0;i<=$device_num-1;i++));
do
echo 'start rank '$i
mkdir ${current_exec_path}/device_test$i
cd ${current_exec_path}/device_test$i || exit
export RANK_ID=0
dev=`expr $i + 0`
export DEVICE_ID=$dev
python ${dirname_path}/${SCRIPT_NAME} \
--is_distributed=0 \
--data_dir=$dataset_path \
--test_model=$model_path \
--ground_truth_mat=$ground_truth_mat \
--save_dir=$save_path \
--rank=$i \
--device_num=$device_num \
--steps_per_epoch=$steps_per_epoch \
--start=$start \
--end=$end > test.log 2>&1 &
done
echo 'running'

141
model_zoo/official/cv/centerface/scripts/train_distribute.sh
Unescape View File

@ -0,0 +1,141 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [ $# != 0 ] && [ $# != 1 ] && [ $# != 2 ] && [ $# != 3 ] && [ $# != 4 ] && [ $# != 5 ]
then
echo "Usage: sh train_distribute.sh [RANK_TABLE] [PRETRAINED_BACKBONE] [DATASET] [ANNOTATIONS] [IMAGES]"
echo " or: sh train_distribute.sh [RANK_TABLE] [PRETRAINED_BACKBONE] [DATASET] [ANNOTATIONS]"
echo " or: sh train_distribute.sh [RANK_TABLE] [PRETRAINED_BACKBONE] [DATASET]"
echo " or: sh train_distribute.sh [RANK_TABLE] [PRETRAINED_BACKBONE]"
echo " or: sh train_distribute.sh [RANK_TABLE]"
echo " or: sh train_distribute.sh "
exit 1
fi
get_real_path(){
if [ "${1:0:1}" == "/" ]; then
echo "$1"
else
echo "$(realpath -m $PWD/$1)"
fi
}
current_exec_path=$(pwd)
echo ${current_exec_path}
dirname_path=$(dirname "$(pwd)")
echo ${dirname_path}
rm -rf ${current_exec_path}/device*
SCRIPT_NAME='train.py'
ulimit -c unlimited
root=${current_exec_path} # your script path
pretrained_backbone=${dirname_path}/mobilenet_v2.ckpt # or mobilenet_v2-b0353104.ckpt
dataset_path=$root/dataset/centerface
annot_path=$dataset_path/annotations/train.json
img_dir=$dataset_path/images/train/images
rank_table=$root/rank_table_8p.json
if [ $# == 1 ]
then
rank_table=$(get_real_path $1)
if [ ! -f $rank_table ]
then
echo "error: rank_table=$rank_table is not a file"
exit 1
fi
fi
if [ $# == 2 ]
then
pretrained_backbone=$(get_real_path $2)
if [ ! -f $pretrained_backbone ]
then
echo "error: pretrained_backbone=$pretrained_backbone is not a file"
exit 1
fi
fi
if [ $# == 3 ]
then
dataset_path=$(get_real_path $3)
if [ ! -f $dataset_path ]
then
echo "error: dataset_path=$dataset_path is not a file"
exit 1
fi
fi
if [ $# == 4 ]
then
annot_path=$(get_real_path $4)
if [ ! -f $annot_path ]
then
echo "error: annot_path=$annot_path is not a file"
exit 1
fi
fi
if [ $# == 5 ]
then
img_dir=$(get_real_path $5)
if [ ! -f $img_dir ]
then
echo "error: img_dir=$img_dir is not a file"
exit 1
fi
fi
echo $rank_table
echo $pretrained_backbone
echo $dataset_path
echo $annot_path
echo $img_dir
export PYTHONPATH=${dirname_path}:$PYTHONPATH
export RANK_TABLE_FILE=$rank_table
export RANK_SIZE=8
task_set_core=24 # for taskset, task_set_core=total cpu number/RANK_SIZE
echo 'start training'
for((i=0;i<=$RANK_SIZE-1;i++));
do
echo 'start rank '$i
mkdir ${current_exec_path}/device$i
cd ${current_exec_path}/device$i || exit
export RANK_ID=$i
dev=`expr $i + 0`
export DEVICE_ID=$dev
taskset -c $((i*task_set_core))-$(((i+1)*task_set_core-1)) python ${dirname_path}/${SCRIPT_NAME} \
--lr=4e-3 \
--per_batch_size=8 \
--is_distributed=1 \
--t_max=140 \
--max_epoch=140 \
--warmup_epochs=0 \
--lr_scheduler=multistep \
--lr_epochs=90,120 \
--weight_decay=0.0000 \
--loss_scale=1024 \
--pretrained_backbone=$pretrained_backbone \
--data_dir=$dataset_path \
--annot_path=$annot_path \
--img_dir=$img_dir > train.log 2>&1 &
done
echo 'running'

131
model_zoo/official/cv/centerface/scripts/train_standalone.sh
Unescape View File

@ -0,0 +1,131 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [ $# != 0 ] && [ $# != 1 ] && [ $# != 2 ] && [ $# != 3 ] && [ $# != 4 ] && [ $# != 5 ]
then
echo "Usage: sh train_standalone.sh [USE_DEVICE_ID] [PRETRAINED_BACKBONE] [DATASET] [ANNOTATIONS] [IMAGES]"
echo " or: sh train_standalone.sh [USE_DEVICE_ID] [PRETRAINED_BACKBONE] [DATASET] [ANNOTATIONS]"
echo " or: sh train_standalone.sh [USE_DEVICE_ID] [PRETRAINED_BACKBONE] [DATASET]"
echo " or: sh train_standalone.sh [USE_DEVICE_ID] [PRETRAINED_BACKBONE]"
echo " or: sh train_standalone.sh [USE_DEVICE_ID]"
echo " or: sh train_standalone.sh "
exit 1
fi
get_real_path(){
if [ "${1:0:1}" == "/" ]; then
echo "$1"
else
echo "$(realpath -m $PWD/$1)"
fi
}
current_exec_path=$(pwd)
echo ${current_exec_path}
dirname_path=$(dirname "$(pwd)")
echo ${dirname_path}
SCRIPT_NAME='train.py'
ulimit -c unlimited
root=${current_exec_path} # your script path
pretrained_backbone=${dirname_path}/mobilenet_v2.ckpt # or mobilenet_v2-b0353104.ckpt
dataset_path=$root/dataset/centerface
annot_path=$dataset_path/annotations/train.json
img_dir=$dataset_path/images/train/images
use_device_id=0
if [ $# == 1 ]
then
use_device_id=$1
fi
if [ $# == 2 ]
then
pretrained_backbone=$(get_real_path $2)
if [ ! -f $pretrained_backbone ]
then
echo "error: pretrained_backbone=$pretrained_backbone is not a file"
exit 1
fi
fi
if [ $# == 3 ]
then
dataset_path=$(get_real_path $3)
if [ ! -f $dataset_path ]
then
echo "error: dataset_path=$dataset_path is not a file"
exit 1
fi
fi
if [ $# == 4 ]
then
annot_path=$(get_real_path $4)
if [ ! -f $annot_path ]
then
echo "error: annot_path=$annot_path is not a file"
exit 1
fi
fi
if [ $# == 5 ]
then
img_dir=$(get_real_path $5)
if [ ! -f $img_dir ]
then
echo "error: img_dir=$img_dir is not a file"
exit 1
fi
fi
echo $use_device_id
echo $pretrained_backbone
echo $dataset_path
echo $annot_path
echo $img_dir
export PYTHONPATH=${dirname_path}:$PYTHONPATH
export RANK_SIZE=1
echo 'start training'
echo 'start rank '$use_device_id
rm -rf ${current_exec_path}/device$use_device_id
mkdir ${current_exec_path}/device$use_device_id
cd ${current_exec_path}/device$use_device_id || exit
export RANK_ID=0
dev=`expr $use_device_id + 0`
export DEVICE_ID=$dev
python ${dirname_path}/${SCRIPT_NAME} \
--lr=5e-4 \
--per_batch_size=8 \
--is_distributed=0 \
--t_max=140 \
--max_epoch=140 \
--warmup_epochs=0 \
--lr_scheduler=multistep \
--lr_epochs=90,120 \
--weight_decay=0.0000 \
--loss_scale=1024 \
--pretrained_backbone=$pretrained_backbone \
--data_dir=$dataset_path \
--annot_path=$annot_path \
--img_dir=$img_dir > train.log 2>&1 &
echo 'running'

0
model_zoo/official/cv/centerface/src/__init__.py
Unescape View File

324
model_zoo/official/cv/centerface/src/centerface.py
View File

File diff suppressed because it is too large Load Diff

64
model_zoo/official/cv/centerface/src/config.py
Unescape View File

@ -0,0 +1,64 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""centerface unique configs"""
class ConfigCenterface():
"""
Config setup
"""
flip_idx = [[0, 1], [3, 4]]
default_resolution = [512, 512]
heads = {'hm': 1, 'wh': 2, 'hm_offset': 2, 'landmarks': 5 * 2}
head_conv = 64
max_objs = 64
rand_crop = True
scale = 0.4
shift = 0.1
aug_rot = 0
color_aug = True
flip = 0.5
input_res = 512 #768 #800
output_res = 128 #192 #200
num_classes = 1
num_joints = 5
reg_offset = True
hm_hp = True
reg_hp_offset = True
dense_hp = False
hm_weight = 1.0
wh_weight = 0.1
off_weight = 1.0
lm_weight = 0.1
rotate = 0
# for test
mean = [0.408, 0.447, 0.470]
std = [0.289, 0.274, 0.278]
test_scales = [0.999,]
nms = 1
flip_test = 0
fix_res = True
input_h = 832 #800
input_w = 832 #800
K = 200
down_ratio = 4
test_batch_size = 1
seed = 317
master_batch_size = 8
num_workers = 8
not_rand_crop = False
no_color_aug = False

181
model_zoo/official/cv/centerface/src/convert_weight.py
Unescape View File

@ -0,0 +1,181 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Centerface model transform
"""
import os
import argparse
import torch
from mindspore.train.serialization import load_checkpoint, save_checkpoint
from mindspore import Tensor
parser = argparse.ArgumentParser(description='')
parser.add_argument('--ckpt_fn', type=str, default='/model_path/centerface.ckpt',
help='ckpt for user to get cell/module name')
parser.add_argument('--pt_fn', type=str, default='/model_path/centerface.pth', help='checkpoint filename to convert')
parser.add_argument('--out_fn', type=str, default='/model_path/centerface_out.ckpt',
help='convert output ckpt/pth path')
parser.add_argument('--pt2ckpt', type=int, default=1, help='1 : pt2ckpt; 0 : ckpt2pt')
args = parser.parse_args()
def load_model(model_path):
"""
Load model
"""
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch']))
state_dict_ = checkpoint['state_dict']
state_dict = {}
# convert data_parallal to model
for k in state_dict_:
if k.find("num_batches_tracked") != -1:
continue
elif k.startswith('module') and not k.startswith('module_list'):
state_dict[k[7:]] = state_dict_[k]
else:
state_dict[k] = state_dict_[k]
return state_dict
def save_model(path, epoch=0, model=None, optimizer=None, state_dict=None):
"""
Sace model file
"""
if state_dict is None:
if isinstance(model, torch.nn.DataParallel):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
data = {'epoch': epoch,
'state_dict': state_dict}
if not optimizer is None:
data['optimizer'] = optimizer.state_dict()
torch.save(data, path)
def load_model_ms(model_path):
"""
Load mindspore model
"""
state_dict_useless = ['global_step', 'learning_rate',
'beta1_power', 'beta2_power']
if os.path.isfile(model_path):
param_dict = load_checkpoint(model_path)
param_dict_new = {}
for key, values in param_dict.items():
if key in state_dict_useless or key.startswith('moments.') \
or key.startswith('moment1.') or key.startswith('moment2.'):
continue
elif key.startswith('centerface_network.'):
param_dict_new[key[19:]] = values
else:
param_dict_new[key] = values
else:
assert FileNotFoundError('{} not exists or not a pre-trained file'.format(model_path))
exit(1)
return param_dict_new
def name_map(ckpt):
"""
Name map
"""
out = {}
for name in ckpt:
# conv + bn
pt_name = name
# backbone
pt_name = pt_name.replace('need_fp1', 'feature_1')
pt_name = pt_name.replace('need_fp2', 'feature_2')
pt_name = pt_name.replace('need_fp3', 'feature_4')
pt_name = pt_name.replace('need_fp4', 'feature_6')
pt_name = pt_name.replace('.features', '')
pt_name = pt_name.replace('.moving_mean', '.running_mean')
pt_name = pt_name.replace('.moving_variance', '.running_var')
pt_name = pt_name.replace('.gamma', '.weight')
pt_name = pt_name.replace('.beta', '.bias')
# fpn
pt_name = pt_name.replace('.up1', '.up_0')
pt_name = pt_name.replace('.up2', '.up_1')
pt_name = pt_name.replace('.up3', '.up_2')
# heads
pt_name = pt_name.replace('hm_head.0.', 'hm.')
pt_name = pt_name.replace('wh_head.', 'wh.')
pt_name = pt_name.replace('off_head.', 'hm_offset.')
pt_name = pt_name.replace('kps_head.', 'landmarks.')
out[pt_name] = name
return out
def pt_to_ckpt(pt, ckpt, out_path):
"""
Pt convert to ckpt file
"""
state_dict_torch = load_model(pt)
state_dict_ms = load_model_ms(ckpt)
name_relate = name_map(state_dict_ms)
new_params_list = []
for key in state_dict_torch:
param_dict = {}
parameter = state_dict_torch[key]
parameter = parameter.numpy()
# depwise conv pytorch[cout, 1, k , k] -> ms[1, cin, k , k], cin = cout
if state_dict_ms[name_relate[key]].data.shape != parameter.shape:
parameter = parameter.transpose(1, 0, 2, 3)
print('ms=', state_dict_ms[name_relate[key]].data.shape, 'pytorch=', parameter.shape, 'name=', key)
param_dict['name'] = name_relate[key]
param_dict['data'] = Tensor(parameter)
new_params_list.append(param_dict)
save_checkpoint(new_params_list, out_path)
return state_dict_ms
def ckpt_to_pt(pt, ckpt, out_path):
"""
Ckpt convert to pt file
"""
state_dict_torch = load_model(pt)
state_dict_ms = load_model_ms(ckpt)
name_relate = name_map(state_dict_ms)
state_dict = {}
for key in state_dict_torch:
name = name_relate[key]
parameter = state_dict_ms[name].data
parameter = parameter.asnumpy()
if state_dict_ms[name_relate[key]].data.shape != state_dict_torch[key].numpy().shape:
print('before ms=', state_dict_ms[name_relate[key]].data.shape, 'pytorch=',
state_dict_torch[key].numpy().shape, 'name=', key)
parameter = parameter.transpose(1, 0, 2, 3)
print('after ms=', state_dict_ms[name_relate[key]].data.shape, 'pytorch=',
state_dict_torch[key].numpy().shape, 'name=', key)
state_dict[key] = torch.from_numpy(parameter)
save_model(out_path, epoch=0, model=None, optimizer=None, state_dict=state_dict)
return state_dict
if __name__ == "__main__":
if args.pt2ckpt == 1:
pt_to_ckpt(args.pt_fn, args.ckpt_fn, args.out_fn)
elif args.pt2ckpt == 0:
ckpt_to_pt(args.pt_fn, args.ckpt_fn, args.out_fn)
else:
# user defined functions
pass

138
model_zoo/official/cv/centerface/src/convert_weight_mobilenetv2.py
Unescape View File

@ -0,0 +1,138 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Mobilenet model transform: torch => mindspore
"""
import os
import argparse
import torch
from mindspore.train.serialization import load_checkpoint, save_checkpoint
from mindspore import Tensor
parser = argparse.ArgumentParser(description='')
parser.add_argument('--ckpt_fn', type=str, default='/model_path/mobilenet_v2_key.ckpt',
help='ckpt for user to get cell/module name')
parser.add_argument('--pt_fn', type=str, default='/model_path/mobilenet_v2-b0353104.pth',
help='checkpoint filename to convert')
parser.add_argument('--out_ckpt_fn', type=str, default='/model_path/mobilenet_v2-b0353104.ckpt',
help='convert output ckpt path')
args = parser.parse_args()
def load_model(model_path):
"""
Load model
"""
state_dict_ = torch.load(model_path, map_location=torch.device('cpu'))
state_dict = {}
# convert data_parallal to model
for k in state_dict_:
if k.find("num_batches_tracked") != -1:
continue
elif k.startswith('module') and not k.startswith('module_list'):
state_dict[k[7:]] = state_dict_[k]
else:
state_dict[k] = state_dict_[k]
return state_dict
def load_model_ms(model_path):
"""
Load mindspore model
"""
state_dict_useless = ['global_step', 'learning_rate',
'beta1_power', 'beta2_power']
if os.path.isfile(model_path):
param_dict = load_checkpoint(model_path)
param_dict_new = {}
for key, values in param_dict.items():
if key in state_dict_useless or key.startswith('moments.') \
or key.startswith('moment1.') or key.startswith('moment2.'):
continue
elif key.startswith('centerface_network.'): #useless, since the start name is "network.backbone."
param_dict_new[key[19:]] = values
else:
param_dict_new[key] = values
else:
assert FileNotFoundError('{} not exists or not a pre-trained file'.format(model_path))
exit(1)
return param_dict_new
def name_map(ckpt):
"""
Name map
"""
out = {}
for name in ckpt:
# conv + bn
pt_name = name
pt_name = pt_name.replace('network.backbone.', '')
# backbone
pt_name = pt_name.replace('need_fp1', 'feature_1')
pt_name = pt_name.replace('need_fp2', 'feature_2')
pt_name = pt_name.replace('need_fp3', 'feature_4')
pt_name = pt_name.replace('need_fp4', 'feature_6')
pt_name = pt_name.replace('.features', '')
pt_name = pt_name.replace('.moving_mean', '.running_mean')
pt_name = pt_name.replace('.moving_variance', '.running_var')
pt_name = pt_name.replace('.gamma', '.weight')
pt_name = pt_name.replace('.beta', '.bias')
# fpn
pt_name = pt_name.replace('.up1', '.up_0')
pt_name = pt_name.replace('.up2', '.up_1')
pt_name = pt_name.replace('.up3', '.up_2')
# heads
pt_name = pt_name.replace('hm_head.0.', 'hm.')
pt_name = pt_name.replace('wh_head.', 'wh.')
pt_name = pt_name.replace('off_head.', 'hm_offset.')
pt_name = pt_name.replace('kps_head.', 'landmarks.')
pt_name = pt_name.replace('network.head.fc.', 'classifier.1.')
out[pt_name] = name
return out
def pt_to_ckpt(pt, ckpt, out_ckpt):
"""
Pt convert to ckpt file
"""
state_dict_torch = load_model(pt)
state_dict_ms = load_model_ms(ckpt)
name_relate = name_map(state_dict_ms)
new_params_list = []
for key in state_dict_torch:
param_dict = {}
parameter = state_dict_torch[key]
parameter = parameter.numpy()
# depwise conv pytorch[cout, 1, k , k] -> ms[1, cin, k , k], cin = cout
if state_dict_ms[name_relate[key]].data.shape != parameter.shape:
parameter = parameter.transpose(1, 0, 2, 3)
print('ms=', state_dict_ms[name_relate[key]].data.shape, 'pytorch=', parameter.shape, 'name=', key)
param_dict['name'] = name_relate[key]
param_dict['data'] = Tensor(parameter)
new_params_list.append(param_dict)
save_checkpoint(new_params_list, out_ckpt)
return state_dict_ms
if __name__ == "__main__":
# beta <=> bias, gamma <=> weight
pt_to_ckpt(args.pt_fn, args.ckpt_fn, args.out_ckpt_fn)

Some files were not shown because too many files have changed in this diff Show More

Write Preview
Loading…
Cancel
Save