From 051991c845a626e59081e385b7d60fd58741147c Mon Sep 17 00:00:00 2001 From: Gaurvi Goyal Date: Fri, 13 Jan 2023 15:15:24 +0100 Subject: [PATCH 1/4] Initial move of movenet libraries to moveenet. To be tested/ --- example/movenet/Dockerfile | 2 +- example/movenet/evaluate.py | 4 +- example/movenet/lib/__init__.py | 30 - example/movenet/movenet_online.py | 4 +- example/movenet/predict.py | 4 +- example/op_detector_example_module/Dockerfile | 4 +- pycore/__init__.py | 0 pycore/moveenet/__init__.py | 30 + pycore/moveenet/data/__init__.py | 0 .../moveenet}/data/center_weight_origin.npy | Bin .../lib => pycore/moveenet}/data/data.py | 4 +- .../moveenet}/data/data_augment.py | 1636 ++++++++--------- .../moveenet}/data/data_tools.py | 4 +- pycore/moveenet/loss/__init__.py | 0 .../moveenet}/loss/movenet_loss.py | 2 +- pycore/moveenet/models/__init__.py | 0 .../moveenet}/models/movenet_mobilenetv2.py | 748 ++++---- pycore/moveenet/task/__init__.py | 0 .../lib => pycore/moveenet}/task/task.py | 12 +- .../moveenet}/task/task_tools.py | 4 +- pycore/moveenet/utils/__init__.py | 0 .../lib => pycore/moveenet}/utils/metrics.py | 0 .../lib => pycore/moveenet}/utils/utils.py | 538 +++--- 23 files changed, 1513 insertions(+), 1513 deletions(-) delete mode 100644 example/movenet/lib/__init__.py create mode 100644 pycore/__init__.py create mode 100644 pycore/moveenet/__init__.py create mode 100644 pycore/moveenet/data/__init__.py rename {example/movenet/lib => pycore/moveenet}/data/center_weight_origin.npy (100%) rename {example/movenet/lib => pycore/moveenet}/data/data.py (95%) rename {example/movenet/lib => pycore/moveenet}/data/data_augment.py (97%) rename {example/movenet/lib => pycore/moveenet}/data/data_tools.py (99%) create mode 100644 pycore/moveenet/loss/__init__.py rename {example/movenet/lib => pycore/moveenet}/loss/movenet_loss.py (99%) create mode 100644 pycore/moveenet/models/__init__.py rename {example/movenet/lib => pycore/moveenet}/models/movenet_mobilenetv2.py (96%) create mode 100644 pycore/moveenet/task/__init__.py rename {example/movenet/lib => pycore/moveenet}/task/task.py (97%) rename {example/movenet/lib => pycore/moveenet}/task/task_tools.py (98%) create mode 100644 pycore/moveenet/utils/__init__.py rename {example/movenet/lib => pycore/moveenet}/utils/metrics.py (100%) rename {example/movenet/lib => pycore/moveenet}/utils/utils.py (96%) diff --git a/example/movenet/Dockerfile b/example/movenet/Dockerfile index 6608ce0e..87d1232f 100644 --- a/example/movenet/Dockerfile +++ b/example/movenet/Dockerfile @@ -92,7 +92,7 @@ RUN apt install -y \ cmake \ && apt-get autoremove \ && apt-get clean \ - && rm -rf /tmp/* /var/lib/apt/lists/* /var/tmp/* + && rm -rf /tmp/* /var/moveenet/apt/lists/* /var/tmp/* RUN pip3 install \ scikit-build diff --git a/example/movenet/evaluate.py b/example/movenet/evaluate.py index 13ae831e..923d4699 100644 --- a/example/movenet/evaluate.py +++ b/example/movenet/evaluate.py @@ -7,10 +7,10 @@ import sys sys.path.append('.') -from lib import init, Data, MoveNet, Task +from pycore.moveenet import init, Data, MoveNet, Task from config import cfg -from lib.utils.utils import arg_parser +from pycore.moveenet.utils.utils import arg_parser def main(cfg): diff --git a/example/movenet/lib/__init__.py b/example/movenet/lib/__init__.py deleted file mode 100644 index f9f7389c..00000000 --- a/example/movenet/lib/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -""" -@Fire -https://github.com/fire717 -""" -import os - -from lib.data.data import Data -from lib.models.movenet_mobilenetv2 import MoveNet -from lib.task.task import Task - - -from lib.utils.utils import setRandomSeed, printDash - - -def init(cfg): - - if cfg["cfg_verbose"]: - printDash() - print(cfg) - printDash() - - - - os.environ["CUDA_VISIBLE_DEVICES"] = cfg['GPU_ID'] - setRandomSeed(cfg['random_seed']) - - - - if not os.path.exists(cfg['save_dir']): - os.makedirs(cfg['save_dir']) \ No newline at end of file diff --git a/example/movenet/movenet_online.py b/example/movenet/movenet_online.py index 56bc816d..14ad0541 100644 --- a/example/movenet/movenet_online.py +++ b/example/movenet/movenet_online.py @@ -1,8 +1,8 @@ from lib import init, MoveNet, Task from config import cfg -from lib.utils.utils import arg_parser -from lib.task.task_tools import image_show, write_output, superimpose +from pycore.moveenet.utils.utils import arg_parser +from pycore.moveenet.task.task_tools import image_show, write_output, superimpose import sys import yarp diff --git a/example/movenet/predict.py b/example/movenet/predict.py index 9dd7dd76..99840a3f 100644 --- a/example/movenet/predict.py +++ b/example/movenet/predict.py @@ -3,10 +3,10 @@ https://github.com/fire717 """ -from lib import init, Data, MoveNet, Task +from pycore.moveenet import init, Data, MoveNet, Task from config import cfg -from lib.utils.utils import arg_parser +from pycore.moveenet.utils.utils import arg_parser # Script to create and save as images all the various outputs of the model diff --git a/example/op_detector_example_module/Dockerfile b/example/op_detector_example_module/Dockerfile index 0f5c2942..c000f480 100644 --- a/example/op_detector_example_module/Dockerfile +++ b/example/op_detector_example_module/Dockerfile @@ -235,7 +235,7 @@ RUN python3 -m pip install --upgrade pip && \ python3 -m pip install torch==1.9.1 torchvision==0.10.1 pandas scipy opencv-python #ENV PYTHONHOME=/root/anaconda/bin -#ENV PYTHONPATH=/root/anaconda/lib:/root/anaconda/lib/python3.8:/root/anaconda/lib/python3.8/site-packages +#ENV PYTHONPATH=/root/anaconda/moveenet:/root/anaconda/moveenet/python3.8:/root/anaconda/moveenet/python3.8/site-packages #ENV PATH=$PYTHONHOME:$PATH # download hpe-core @@ -244,7 +244,7 @@ RUN cd $SOURCE_FOLDER && \ cd hpe-core && \ git checkout $HPE_BRANCH -# add /usr/local/lib to the library path, so that libcaffe.so compiled with openpose will be used +# add /usr/local/moveenet to the library path, so that libcaffe.so compiled with openpose will be used # instead of the one provided by the nvidia/cuda docker image ENV LD_LIBRARY_PATH /usr/local/lib:$LD_LIBRARY_PATH diff --git a/pycore/__init__.py b/pycore/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pycore/moveenet/__init__.py b/pycore/moveenet/__init__.py new file mode 100644 index 00000000..0038ee7d --- /dev/null +++ b/pycore/moveenet/__init__.py @@ -0,0 +1,30 @@ +""" +@Fire +https://github.com/fire717 +""" +import os + +from pycore.moveenet.data.data import Data +from pycore.moveenet.models.movenet_mobilenetv2 import MoveNet +from pycore.moveenet.task.task import Task + + +from pycore.moveenet.utils.utils import setRandomSeed, printDash + + +def init(cfg): + + if cfg["cfg_verbose"]: + printDash() + print(cfg) + printDash() + + + + os.environ["CUDA_VISIBLE_DEVICES"] = cfg['GPU_ID'] + setRandomSeed(cfg['random_seed']) + + + # + # if not os.path.exists(cfg['save_dir']): + # os.makedirs(cfg['save_dir']) \ No newline at end of file diff --git a/pycore/moveenet/data/__init__.py b/pycore/moveenet/data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example/movenet/lib/data/center_weight_origin.npy b/pycore/moveenet/data/center_weight_origin.npy similarity index 100% rename from example/movenet/lib/data/center_weight_origin.npy rename to pycore/moveenet/data/center_weight_origin.npy diff --git a/example/movenet/lib/data/data.py b/pycore/moveenet/data/data.py similarity index 95% rename from example/movenet/lib/data/data.py rename to pycore/moveenet/data/data.py index 5ed718ad..dc2e5596 100644 --- a/example/movenet/lib/data/data.py +++ b/pycore/moveenet/data/data.py @@ -8,8 +8,8 @@ import json import copy -from lib.data.data_tools import getDataLoader, getFileNames -from lib.task.task_tools import movenetDecode +from pycore.moveenet.data.data_tools import getDataLoader, getFileNames +from pycore.moveenet.task.task_tools import movenetDecode class Data(): diff --git a/example/movenet/lib/data/data_augment.py b/pycore/moveenet/data/data_augment.py similarity index 97% rename from example/movenet/lib/data/data_augment.py rename to pycore/moveenet/data/data_augment.py index af52c982..9734b6ad 100644 --- a/example/movenet/lib/data/data_augment.py +++ b/pycore/moveenet/data/data_augment.py @@ -1,818 +1,818 @@ -""" -@Fire -https://github.com/fire717 -""" -from PIL import Image -import numpy as np -import pandas as pd -import os -import torch -from torch.utils.data.dataset import Dataset -import torchvision.transforms as transforms -import torchvision.transforms.functional as F -import random -import cv2 -import albumentations as A -import json -import platform -from copy import deepcopy - -###### tools -def Mirror(src,label=None): - """ - item = { - "img_name":save_name, - "keypoints":save_keypoints, relative position - "center":save_center, - "other_centers":other_centers, - "other_keypoints":other_keypoints, - } - # mirror 后左右手顺序就变了！ - """ - keypoints = label['keypoints'] - center = label['center'] - other_centers = label['other_centers'] - other_keypoints = label['other_keypoints'] - - img = cv2.flip(src, 1) - if label is None: - return img,label - - - for i in range(len(keypoints)): - if i%3==0: - keypoints[i] = 1-keypoints[i] - try: - keypoints = [ #17 keypoint skeleton mirror - keypoints[0],keypoints[1],keypoints[2], - keypoints[6],keypoints[7],keypoints[8], - keypoints[3],keypoints[4],keypoints[5], - keypoints[12],keypoints[13],keypoints[14], - keypoints[9],keypoints[10],keypoints[11], - keypoints[18],keypoints[19],keypoints[20], - keypoints[15],keypoints[16],keypoints[17], - keypoints[24],keypoints[25],keypoints[26], - keypoints[21],keypoints[22],keypoints[23], - keypoints[30],keypoints[31],keypoints[32], - keypoints[27],keypoints[28],keypoints[29], - keypoints[36],keypoints[37],keypoints[38], - keypoints[33],keypoints[34],keypoints[35], - keypoints[42],keypoints[43],keypoints[44], - keypoints[39],keypoints[40],keypoints[41], - keypoints[48],keypoints[49],keypoints[50], - keypoints[45],keypoints[46],keypoints[47]] - except IndexError: - keypoints = [ #13 keypoint skeleton mirror - keypoints[0], keypoints[1], keypoints[2], - keypoints[6], keypoints[7], keypoints[8], - keypoints[3], keypoints[4], keypoints[5], - keypoints[12], keypoints[13], keypoints[14], - keypoints[9], keypoints[10], keypoints[11], - keypoints[18], keypoints[19], keypoints[20], - keypoints[15], keypoints[16], keypoints[17], - keypoints[24], keypoints[25], keypoints[26], - keypoints[21], keypoints[22], keypoints[23], - keypoints[30], keypoints[31], keypoints[32], - keypoints[27], keypoints[28], keypoints[29], - keypoints[36], keypoints[37], keypoints[38], - keypoints[33], keypoints[34], keypoints[35]] - #print(center, other_centers, other_keypoints) - center[0] = 1-center[0] - - for i in range(len(other_centers)): - other_centers[i][0] = 1 - other_centers[i][0] - - for i in range(len(other_keypoints)): - for j in range(len(other_keypoints[i])): - other_keypoints[i][j][0] = 1 - other_keypoints[i][j][0] - other_keypoints = other_keypoints[::-1] - - label["keypoints"]=keypoints - label["center"]=center - label["other_centers"]=other_centers - label["other_keypoints"]=other_keypoints - - - - return img,label - - -def Padding(src, label, pad_color, max_pad_ratio=0.12): - """ - item = { - "img_name":save_name, - "keypoints":save_keypoints, relative position - "center":save_center, - "other_centers":other_centers, - "other_keypoints":other_keypoints, - } - """ - h,w = src.shape[:2] - max_size = max(h,w) - - pad_ratio = (random.random()+1)*max_pad_ratio/2 - pad_size = int(max_size*pad_ratio) - new_size = max_size+pad_size*2 - - new_img = np.ones((new_size,new_size,3))*pad_color - new_x0 = (new_size-w)//2 - new_y0 = (new_size-h)//2 - new_img[new_y0:new_y0+h,new_x0:new_x0+w,:] = src - - new_img = cv2.resize(new_img, (w, h)) - - if label is None: - return new_img.astype(np.uint8) - - - keypoints = label['keypoints'] - center = label['center'] - other_centers = label['other_centers'] - other_keypoints = label['other_keypoints'] - - - for i in range(len(keypoints)): - if i%3==0: - keypoints[i] = (keypoints[i]*w+new_x0)/new_size - elif i%3==1: - keypoints[i] = (keypoints[i]*h+new_y0)/new_size - - - center[0] = (center[0]*w+new_x0)/new_size - center[1] = (center[1]*h+new_y0)/new_size - - for i in range(len(other_centers)): - other_centers[i][0] = (other_centers[i][0]*w+new_x0)/new_size - other_centers[i][1] = (other_centers[i][1]*h+new_y0)/new_size - - - for i in range(len(other_keypoints)): - for j in range(len(other_keypoints[i])): - other_keypoints[i][j][0] = (other_keypoints[i][j][0]*w+new_x0)/new_size - other_keypoints[i][j][1] = (other_keypoints[i][j][1]*h+new_y0)/new_size - - - label["keypoints"]=keypoints - label["center"]=center - label["other_centers"]=other_centers - label["other_keypoints"]=other_keypoints - - - return new_img.astype(np.uint8), label - - -def Crop(src, label, pad_color, max_pad_ratio=0.3): - """ - item = { - "img_name":save_name, - "keypoints":save_keypoints, relative position - "center":save_center, - "other_centers":other_centers, - "other_keypoints":other_keypoints, - } - """ - h,w = src.shape[:2] - - keypoints = label['keypoints'] - center = label['center'] - other_centers = label['other_centers'] - other_keypoints = label['other_keypoints'] - - pad_ratio = random.uniform(max_pad_ratio/3, max_pad_ratio) - if len(other_centers)>0: - crop_x = int(w*pad_ratio) - crop_y = int(h*pad_ratio) - else: - crop_x = random.randint(int(w*pad_ratio)//2,int(w*pad_ratio)) - crop_y = random.randint(int(w*pad_ratio)//2,int(h*pad_ratio)) - - new_w = int(w-w*pad_ratio*2) - new_h = int(h-h*pad_ratio*2) - - new_img = src[crop_y:crop_y+new_h,crop_x:crop_x+new_w] - new_img = cv2.resize(new_img, (w, h)) - ## - - if label is None: - return new_img.astype(np.uint8) - - - - for i in range(len(keypoints)): - if i%3==0: - keypoints[i] = (keypoints[i]*w-crop_x)/new_w - elif i%3==1: - keypoints[i] = (keypoints[i]*h-crop_y)/new_h - for i in range(len(keypoints)//3): - if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: - keypoints[i*3+2] = 0 - - center[0] = min(max(0,(center[0]*w-crop_x)/new_w),1) - center[1] = min(max(0,(center[1]*h-crop_y)/new_h),1) - - for i in range(len(other_centers)): - other_centers[i][0] = (other_centers[i][0]*w-crop_x)/new_w - other_centers[i][1] = (other_centers[i][1]*h-crop_y)/new_h - other_centers_new = [] - for item in other_centers: - if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: - continue - else: - other_centers_new.append(item) - other_centers = other_centers_new - - for i in range(len(other_keypoints)): - for j in range(len(other_keypoints[i])): - other_keypoints[i][j][0] = (other_keypoints[i][j][0]*w-crop_x)/new_w - other_keypoints[i][j][1] = (other_keypoints[i][j][1]*h-crop_y)/new_h - other_keypoints_res = [] - for item in other_keypoints[i]: - if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: - continue - else: - other_keypoints_res.append(item) - other_keypoints[i] = other_keypoints_res - - label["keypoints"]=keypoints - label["center"]=center - label["other_centers"]=other_centers - label["other_keypoints"]=other_keypoints - - return new_img.astype(np.uint8), label - - -def Move(src, label, pad_color, max_move_ratio=0.2): - """ - item = { - "img_name":save_name, - "keypoints":save_keypoints, relative position - "center":save_center, - "other_centers":other_centers, - "other_keypoints":other_keypoints, - } - """ - h,w = src.shape[:2] - - keypoints = label['keypoints'] - center = label['center'] - - pad_ratio = random.uniform(max_move_ratio/2, max_move_ratio) - - move_x = random.randint(int(w*pad_ratio)//2,int(w*pad_ratio))*random.choice([-1,1]) - move_y = random.randint(int(w*pad_ratio)//2,int(h*pad_ratio))*random.choice([-1,1]) - - - M = np.float32([[1, 0, move_x], [0, 1, move_y]]) - src = cv2.warpAffine(src, M, (src.shape[1], src.shape[0])) - - - # new_img = src[move_y:move_y+new_h,move_x:move_x+new_w] - - - for i in range(len(keypoints)): - if i%3==0: - keypoints[i] = keypoints[i]+move_x/w - elif i%3==1: - keypoints[i] = keypoints[i]+move_y/h - for i in range(len(keypoints)//3): - if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: - keypoints[i*3+2] = 0 - - center[0] = min(max(0,(center[0]+move_x/w)),1) - center[1] = min(max(0,(center[1]+move_y/h)),1) - - - label["keypoints"]=keypoints - label["center"]=center - - - return src.astype(np.uint8), label - -def Rotate(src,angle,pad_color,label=None,center=None,scale=1.0): - ''' - :param src: src image - :param label: dict - :param angle: - :param center: - :param scale: - :return: the rotated image and the points - ''' - image=src - (h, w) = image.shape[:2] - - - # 若未指定旋转中心，则将图像中心设为旋转中心 - if center is None: - center = (w / 2, h / 2) - # 执行旋转 - M = cv2.getRotationMatrix2D(center, angle, scale) - if label is None: - for i in range(image.shape[2]): - image[:,:,i] = cv2.warpAffine(image[:,:,i], M, (w, h), - flags=cv2.INTER_CUBIC, - borderMode=cv2.BORDER_CONSTANT, - borderValue=pad_color) - return image,None - else: - - keypoints = label['keypoints'] - center = label['center'] - other_centers = label['other_centers'] - other_keypoints = label['other_keypoints'] - - - ####make it as a 3x3 RT matrix - full_M=np.row_stack((M,np.asarray([0,0,1]))) - img_rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, - borderMode=cv2.BORDER_CONSTANT, borderValue=pad_color) - ###make the keypoints as 3xN matrix - keypoints = np.reshape(keypoints, (-1,3)) - keypoints_z = keypoints[:,2].reshape((-1,1)) - keypoints = keypoints[:,:2] - keypoints[:,0] *= w - keypoints[:,1] *= h - keypoints = keypoints.astype(np.int32) - - keypoints = keypoints.T - full_keypoints = np.row_stack((keypoints, np.ones(shape=(1,keypoints.shape[1])))) - keypoints_rotated=np.dot(full_M,full_keypoints) - keypoints_rotated=keypoints_rotated[0:2,:] - #keypoints_rotated = keypoints_rotated.astype(np.int32) - keypoints_rotated=keypoints_rotated.T - - keypoints_rotated = keypoints_rotated.astype(np.float32) - keypoints_rotated[:,0] /=w - keypoints_rotated[:,1] /=h - # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) - keypoints = np.concatenate([keypoints_rotated,keypoints_z],-1).reshape((-1)).tolist() - for i in range(len(keypoints)//3): - if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: - keypoints[i*3+2] = 0 - - center = np.reshape(center, (-1,2)) - center[:,0] *= w - center[:,1] *= h - center = center.astype(np.int32) - - center = center.T - full_center = np.row_stack((center, np.ones(shape=(1,center.shape[1])))) - center_rotated=np.dot(full_M,full_center) - center_rotated=center_rotated[0:2,:] - #keypoints_rotated = keypoints_rotated.astype(np.int32) - center_rotated=center_rotated.T - - center_rotated = center_rotated.astype(np.float32) - center_rotated[:,0] /=w - center_rotated[:,1] /=h - # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) - center = center_rotated.reshape((-1)).tolist() - - - - other_centers = np.reshape(other_centers, (-1,2)) - other_centers[:,0] *= w - other_centers[:,1] *= h - other_centers = other_centers.astype(np.int32) - - other_centers = other_centers.T - full_center = np.row_stack((other_centers, np.ones(shape=(1,other_centers.shape[1])))) - center_rotated=np.dot(full_M,full_center) - center_rotated=center_rotated[0:2,:] - #keypoints_rotated = keypoints_rotated.astype(np.int32) - center_rotated=center_rotated.T - - center_rotated = center_rotated.astype(np.float32) - center_rotated[:,0] /= w - center_rotated[:,1] /= h - # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) - other_centers_raw = center_rotated.reshape((-1,2)).tolist() - other_centers = [] - for item in other_centers_raw: - if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: - continue - else: - other_centers.append(item) - - - for i in range(len(other_keypoints)): - if len(other_keypoints[i])>0: - other_keypointsi = np.reshape(other_keypoints[i], (-1,2)) - other_keypointsi[:,0] *= w - other_keypointsi[:,1] *= h - other_keypointsi = other_keypointsi.astype(np.int32) - - other_keypointsi = other_keypointsi.T - full_center = np.row_stack((other_keypointsi, np.ones(shape=(1,other_keypointsi.shape[1])))) - center_rotated=np.dot(full_M,full_center) - center_rotated=center_rotated[0:2,:] - #keypoints_rotated = keypoints_rotated.astype(np.int32) - center_rotated=center_rotated.T - - center_rotated = center_rotated.astype(np.float32) - center_rotated[:,0] /=w - center_rotated[:,1] /=h - # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) - other_keypoints_i = center_rotated.reshape((-1,2)).tolist() - other_keypoints_res = [] - for item in other_keypoints_i: - if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: - continue - else: - other_keypoints_res.append(item) - other_keypoints[i] = other_keypoints_res - - label["keypoints"]= keypoints - label["center"]=center - label["other_centers"]=other_centers - label["other_keypoints"]=other_keypoints - - - return img_rotated, label - - -def Affine(src,strength,pad_color,label=None): - image = src - (h, w) = image.shape[:2] - - - keypoints = label['keypoints'] - center = label['center'] - other_centers = label['other_centers'] - other_keypoints = label['other_keypoints'] - - - pts_base = np.float32([[10,100],[200,50],[100,250]]) - pts1 = np.random.rand(3, 2) * -40 + pts_base - pts1 = pts1.astype(np.float32) - M = cv2.getAffineTransform(pts1, pts_base) - trans_img = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]) , - borderMode=cv2.BORDER_CONSTANT, - borderValue=pad_color) - if label is None: - return trans_img, label - - keypoints = np.reshape(keypoints, (-1,3)) - keypoints_z = keypoints[:,2].reshape((-1,1)) - keypoints = keypoints[:,:2] - keypoints[:,0] *= w - keypoints[:,1] *= h - keypoints = keypoints.astype(np.int32) - - keypoints=keypoints.T - full_data = np.row_stack((keypoints, np.ones(shape=(1, keypoints.shape[1])))) - data_rotated = np.dot(M, full_data) - #label_rotated = label_rotated.astype(np.int32) - data_rotated=data_rotated.T - - data_rotated = data_rotated.astype(np.float32) - data_rotated[:,0] /=w - data_rotated[:,1] /=h - keypoints = np.concatenate([data_rotated,keypoints_z],-1).reshape((-1)).tolist() - for i in range(len(keypoints)//3): - if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: - keypoints[i*3+2] = 0 - - - center = np.reshape(center, (-1,2)) - center[:,0] *= w - center[:,1] *= h - center = center.astype(np.int32) - - center=center.T - full_data = np.row_stack((center, np.ones(shape=(1, center.shape[1])))) - data_rotated = np.dot(M, full_data) - #label_rotated = label_rotated.astype(np.int32) - data_rotated=data_rotated.T - - data_rotated = data_rotated.astype(np.float32) - data_rotated[:,0] /=w - data_rotated[:,1] /=h - center = data_rotated.reshape((-1)).tolist() - - - - other_centers = np.reshape(other_centers, (-1,2)) - other_centers[:,0] *= w - other_centers[:,1] *= h - other_centers = other_centers.astype(np.int32) - - other_centers = other_centers.T - full_data = np.row_stack((other_centers, np.ones(shape=(1, other_centers.shape[1])))) - data_rotated = np.dot(M, full_data) - #label_rotated = label_rotated.astype(np.int32) - data_rotated=data_rotated.T - - data_rotated = data_rotated.astype(np.float32) - data_rotated[:,0] /=w - data_rotated[:,1] /=h - other_centers_raw = data_rotated.reshape((-1,2)).tolist() - other_centers = [] - for item in other_centers_raw: - if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: - continue - else: - other_centers.append(item) - - - for i in range(len(other_keypoints)): - if len(other_keypoints[i])>0: - other_keypointsi = np.reshape(other_keypoints[i], (-1,2)) - other_keypointsi[:,0] *= w - other_keypointsi[:,1] *= h - other_keypointsi = other_keypointsi.astype(np.int32) - - other_keypointsi = other_keypointsi.T - full_data = np.row_stack((other_keypointsi, np.ones(shape=(1, other_keypointsi.shape[1])))) - data_rotated = np.dot(M, full_data) - #label_rotated = label_rotated.astype(np.int32) - data_rotated=data_rotated.T - - data_rotated = data_rotated.astype(np.float32) - data_rotated[:,0] /=w - data_rotated[:,1] /=h - # other_keypoints[i] = data_rotated.reshape((-1,2)).tolist() - other_keypoints_i = data_rotated.reshape((-1,2)).tolist() - other_keypoints_res = [] - for item in other_keypoints_i: - if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: - continue - else: - other_keypoints_res.append(item) - other_keypoints[i] = other_keypoints_res - - - label["keypoints"]=keypoints - label["center"]=center - label["other_centers"]=other_centers - label["other_keypoints"]=other_keypoints - - return trans_img, label - - -def AID(img, label): - h,w = img.shape[:2] - - half_size = int(random.uniform(3/192,6/192)*(h+w)/2) - - - keypoints = np.array(label["keypoints"]).reshape((-1,3)) - - keypoints = keypoints[keypoints[:,2]>0] - dropout_id = random.randint(0,keypoints.shape[0]-1) - - - cx = keypoints[dropout_id][0]*w - cy = keypoints[dropout_id][1]*h - - x0 = int(max(0,cx-half_size)) - y0 = int(max(0,cy-half_size)) - x1 = int(min(w-1,cx+half_size)) - y1 = int(min(h-1,cy+half_size)) - - color = random.randint(0,255) - img[y0:y1, x0:x1] = (color,color,color) - - return img - -def AID2(img, label): - h,w = img.shape[:2] - - half_size = int(random.uniform(3/192,6/192)*(h+w)/2) - - - keypoints = np.array(label["keypoints"]).reshape((-1,3)) - - keypoints = keypoints[keypoints[:,2]>0] - dropout_id = random.randint(0,keypoints.shape[0]-1) - - - cx = keypoints[dropout_id][0]*w - cy = keypoints[dropout_id][1]*h - - x0 = int(max(0,cx-half_size)) - y0 = int(max(0,cy-half_size)) - x1 = int(min(w-1,cx+half_size)) - y1 = int(min(h-1,cy+half_size)) - # print(x0,y0,x1,y1) - # b - color = random.randint(0,255) - dist = half_size**2 - for i in range(y1-y0+1): - for j in range(x1-x0+1): - px = x0+j - py = y0+i - if ((py-cy)**2+(px-cx)**2)<=dist: - img[py, px] = (color,color,color) - - return img - -def dropout(src,pad_color,max_pattern_ratio=0.05): - width_ratio = random.uniform(0, max_pattern_ratio) - height_ratio = random.uniform(0, max_pattern_ratio) - width=src.shape[1] - height=src.shape[0] - block_width=width*width_ratio - block_height=height*height_ratio - width_start=int(random.uniform(0,width-block_width)) - width_end=int(width_start+block_width) - height_start=int(random.uniform(0,height-block_height)) - height_end=int(height_start+block_height) - src[height_start:height_end,width_start:width_end,:]=np.array(pad_color,dtype=src.dtype) - - return src - -def pixel_jitter(src,p=0.5,max_=5.): - - src=src.astype(np.float32) - - pattern=(np.random.rand(src.shape[0], src.shape[1],src.shape[2])-0.5)*2*max_ - img = src + pattern - - img[img<0]=0 - img[img >255] = 255 - - img = img.astype(np.uint8) - - return img - -def _clip(image): - """ - Clip and convert an image to np.uint8. - Args - image: Image to clip. - """ - return np.clip(image, 0, 255).astype(np.uint8) - -def adjust_contrast(image, factor): - """ Adjust contrast of an image. - Args - image: Image to adjust. - factor: A factor for adjusting contrast. - """ - mean = image.mean(axis=0).mean(axis=0) - return _clip((image - mean) * factor + mean) - -def adjust_brightness(image, delta): - """ Adjust brightness of an image - Args - image: Image to adjust. - delta: Brightness offset between -1 and 1 added to the pixel values. - """ - return _clip(image + delta * 255) - -def adjust_hue(image, delta): - """ Adjust hue of an image. - Args - image: Image to adjust. - delta: An interval between -1 and 1 for the amount added to the hue channel. - The values are rotated if they exceed 180. - """ - image[..., 0] = np.mod(image[..., 0] + delta * 180, 180) - return image - -def adjust_saturation(image, factor): - """ Adjust saturation of an image. - Args - image: Image to adjust. - factor: An interval for the factor multiplying the saturation values of each pixel. - """ - image[..., 1] = np.clip(image[..., 1] * factor, 0, 255) - return image - -class ColorDistort(): - def __init__( - self, - contrast_range=(0.8, 1.2), - brightness_range=(-.2, .2), - hue_range=(-0.1, 0.1), - saturation_range=(0.8, 1.2) - ): - self.contrast_range = contrast_range - self.brightness_range = brightness_range - self.hue_range = hue_range - self.saturation_range = saturation_range - - - def _uniform(self,val_range): - """ Uniformly sample from the given range. - Args - val_range: A pair of lower and upper bound. - """ - return np.random.uniform(val_range[0], val_range[1]) - - def __call__(self, image): - - - if self.contrast_range is not None: - contrast_factor = self._uniform(self.contrast_range) - image = adjust_contrast(image,contrast_factor) - if self.brightness_range is not None: - brightness_delta = self._uniform(self.brightness_range) - image = adjust_brightness(image, brightness_delta) - - if self.hue_range is not None or self.saturation_range is not None: - - image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) - - if self.hue_range is not None: - hue_delta = self._uniform(self.hue_range) - image = adjust_hue(image, hue_delta) - - if self.saturation_range is not None: - saturation_factor = self._uniform(self.saturation_range) - image = adjust_saturation(image, saturation_factor) - - image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) - - return image - - - -###### Data Aug api -class DataAug: - def __init__(self, img_size): - self.h = img_size - self.w = img_size - self.color_augmentor = ColorDistort() - - def __call__(self, img, label): - """ - img: opencv img, BGR - item = { - "img_name":save_name, - "keypoints":save_keypoints, relative position - "center":save_center, - "other_centers":other_centers, - "other_keypoints":other_keypoints, - } - return: same type as img and label - """ - pad_color = random.randint(0,255) - - - new_img = deepcopy(img) - new_label = deepcopy(label) - - - if random.random() < 0.5: - new_img, new_label = Mirror(new_img, label=new_label) - - - rd = random.random() - if rd <0.5: - if random.random() < 0.2: - new_img, new_label = Padding(new_img, new_label, (0,0,0)) - else: - new_img, new_label = Crop(new_img, new_label, (0,0,0)) - elif rd < 0.65: - if random.random() < 0.7: - if len(new_label["other_centers"])==0: - strength = random.uniform(20, 50) - new_img, new_label = Affine(new_img, strength=strength, label=new_label, pad_color=(0,0,0)) - - else: - angle = random.uniform(-24, 24) - new_img, new_label = Rotate(new_img, label=new_label, angle=angle, pad_color=(0,0,0)) - - elif rd < 0.85: - - if len(new_label["other_centers"])==0 and (len(new_label["other_keypoints"][0]) + len(new_label["other_keypoints"][3]))==0: - if random.random() < 0.5: - new_img, new_label = Move(new_img, new_label, (0,0,0)) - else: - pass - ### in case of no points - count_zero = 0 - count_zero_new = 0 - for i in range(len(new_label["keypoints"])): - if i%3==2: - if label["keypoints"][i]==0: - count_zero+=1 - if new_label["keypoints"][i]==0: - count_zero_new+=1 - if count_zero_new-count_zero > 3 or count_zero_new>5: - new_img = img - new_label = label - - - - - if random.random() < 0.4: - new_img = self.color_augmentor(new_img) - - if random.random() < 0.4: - new_img = pixel_jitter(new_img,15) - - # if random.random() < 0.3: - # new_img = dropout(new_img, pad_color,0.1) - - if random.random() < 0.3: - - new_img = AID(new_img, new_label) - - - return new_img, new_label - - - - +""" +@Fire +https://github.com/fire717 +""" +from PIL import Image +import numpy as np +import pandas as pd +import os +import torch +from torch.utils.data.dataset import Dataset +import torchvision.transforms as transforms +import torchvision.transforms.functional as F +import random +import cv2 +import albumentations as A +import json +import platform +from copy import deepcopy + +###### tools +def Mirror(src,label=None): + """ + item = { + "img_name":save_name, + "keypoints":save_keypoints, relative position + "center":save_center, + "other_centers":other_centers, + "other_keypoints":other_keypoints, + } + # mirror 后左右手顺序就变了！ + """ + keypoints = label['keypoints'] + center = label['center'] + other_centers = label['other_centers'] + other_keypoints = label['other_keypoints'] + + img = cv2.flip(src, 1) + if label is None: + return img,label + + + for i in range(len(keypoints)): + if i%3==0: + keypoints[i] = 1-keypoints[i] + try: + keypoints = [ #17 keypoint skeleton mirror + keypoints[0],keypoints[1],keypoints[2], + keypoints[6],keypoints[7],keypoints[8], + keypoints[3],keypoints[4],keypoints[5], + keypoints[12],keypoints[13],keypoints[14], + keypoints[9],keypoints[10],keypoints[11], + keypoints[18],keypoints[19],keypoints[20], + keypoints[15],keypoints[16],keypoints[17], + keypoints[24],keypoints[25],keypoints[26], + keypoints[21],keypoints[22],keypoints[23], + keypoints[30],keypoints[31],keypoints[32], + keypoints[27],keypoints[28],keypoints[29], + keypoints[36],keypoints[37],keypoints[38], + keypoints[33],keypoints[34],keypoints[35], + keypoints[42],keypoints[43],keypoints[44], + keypoints[39],keypoints[40],keypoints[41], + keypoints[48],keypoints[49],keypoints[50], + keypoints[45],keypoints[46],keypoints[47]] + except IndexError: + keypoints = [ #13 keypoint skeleton mirror + keypoints[0], keypoints[1], keypoints[2], + keypoints[6], keypoints[7], keypoints[8], + keypoints[3], keypoints[4], keypoints[5], + keypoints[12], keypoints[13], keypoints[14], + keypoints[9], keypoints[10], keypoints[11], + keypoints[18], keypoints[19], keypoints[20], + keypoints[15], keypoints[16], keypoints[17], + keypoints[24], keypoints[25], keypoints[26], + keypoints[21], keypoints[22], keypoints[23], + keypoints[30], keypoints[31], keypoints[32], + keypoints[27], keypoints[28], keypoints[29], + keypoints[36], keypoints[37], keypoints[38], + keypoints[33], keypoints[34], keypoints[35]] + #print(center, other_centers, other_keypoints) + center[0] = 1-center[0] + + for i in range(len(other_centers)): + other_centers[i][0] = 1 - other_centers[i][0] + + for i in range(len(other_keypoints)): + for j in range(len(other_keypoints[i])): + other_keypoints[i][j][0] = 1 - other_keypoints[i][j][0] + other_keypoints = other_keypoints[::-1] + + label["keypoints"]=keypoints + label["center"]=center + label["other_centers"]=other_centers + label["other_keypoints"]=other_keypoints + + + + return img,label + + +def Padding(src, label, pad_color, max_pad_ratio=0.12): + """ + item = { + "img_name":save_name, + "keypoints":save_keypoints, relative position + "center":save_center, + "other_centers":other_centers, + "other_keypoints":other_keypoints, + } + """ + h,w = src.shape[:2] + max_size = max(h,w) + + pad_ratio = (random.random()+1)*max_pad_ratio/2 + pad_size = int(max_size*pad_ratio) + new_size = max_size+pad_size*2 + + new_img = np.ones((new_size,new_size,3))*pad_color + new_x0 = (new_size-w)//2 + new_y0 = (new_size-h)//2 + new_img[new_y0:new_y0+h,new_x0:new_x0+w,:] = src + + new_img = cv2.resize(new_img, (w, h)) + + if label is None: + return new_img.astype(np.uint8) + + + keypoints = label['keypoints'] + center = label['center'] + other_centers = label['other_centers'] + other_keypoints = label['other_keypoints'] + + + for i in range(len(keypoints)): + if i%3==0: + keypoints[i] = (keypoints[i]*w+new_x0)/new_size + elif i%3==1: + keypoints[i] = (keypoints[i]*h+new_y0)/new_size + + + center[0] = (center[0]*w+new_x0)/new_size + center[1] = (center[1]*h+new_y0)/new_size + + for i in range(len(other_centers)): + other_centers[i][0] = (other_centers[i][0]*w+new_x0)/new_size + other_centers[i][1] = (other_centers[i][1]*h+new_y0)/new_size + + + for i in range(len(other_keypoints)): + for j in range(len(other_keypoints[i])): + other_keypoints[i][j][0] = (other_keypoints[i][j][0]*w+new_x0)/new_size + other_keypoints[i][j][1] = (other_keypoints[i][j][1]*h+new_y0)/new_size + + + label["keypoints"]=keypoints + label["center"]=center + label["other_centers"]=other_centers + label["other_keypoints"]=other_keypoints + + + return new_img.astype(np.uint8), label + + +def Crop(src, label, pad_color, max_pad_ratio=0.3): + """ + item = { + "img_name":save_name, + "keypoints":save_keypoints, relative position + "center":save_center, + "other_centers":other_centers, + "other_keypoints":other_keypoints, + } + """ + h,w = src.shape[:2] + + keypoints = label['keypoints'] + center = label['center'] + other_centers = label['other_centers'] + other_keypoints = label['other_keypoints'] + + pad_ratio = random.uniform(max_pad_ratio/3, max_pad_ratio) + if len(other_centers)>0: + crop_x = int(w*pad_ratio) + crop_y = int(h*pad_ratio) + else: + crop_x = random.randint(int(w*pad_ratio)//2,int(w*pad_ratio)) + crop_y = random.randint(int(w*pad_ratio)//2,int(h*pad_ratio)) + + new_w = int(w-w*pad_ratio*2) + new_h = int(h-h*pad_ratio*2) + + new_img = src[crop_y:crop_y+new_h,crop_x:crop_x+new_w] + new_img = cv2.resize(new_img, (w, h)) + ## + + if label is None: + return new_img.astype(np.uint8) + + + + for i in range(len(keypoints)): + if i%3==0: + keypoints[i] = (keypoints[i]*w-crop_x)/new_w + elif i%3==1: + keypoints[i] = (keypoints[i]*h-crop_y)/new_h + for i in range(len(keypoints)//3): + if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: + keypoints[i*3+2] = 0 + + center[0] = min(max(0,(center[0]*w-crop_x)/new_w),1) + center[1] = min(max(0,(center[1]*h-crop_y)/new_h),1) + + for i in range(len(other_centers)): + other_centers[i][0] = (other_centers[i][0]*w-crop_x)/new_w + other_centers[i][1] = (other_centers[i][1]*h-crop_y)/new_h + other_centers_new = [] + for item in other_centers: + if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: + continue + else: + other_centers_new.append(item) + other_centers = other_centers_new + + for i in range(len(other_keypoints)): + for j in range(len(other_keypoints[i])): + other_keypoints[i][j][0] = (other_keypoints[i][j][0]*w-crop_x)/new_w + other_keypoints[i][j][1] = (other_keypoints[i][j][1]*h-crop_y)/new_h + other_keypoints_res = [] + for item in other_keypoints[i]: + if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: + continue + else: + other_keypoints_res.append(item) + other_keypoints[i] = other_keypoints_res + + label["keypoints"]=keypoints + label["center"]=center + label["other_centers"]=other_centers + label["other_keypoints"]=other_keypoints + + return new_img.astype(np.uint8), label + + +def Move(src, label, pad_color, max_move_ratio=0.2): + """ + item = { + "img_name":save_name, + "keypoints":save_keypoints, relative position + "center":save_center, + "other_centers":other_centers, + "other_keypoints":other_keypoints, + } + """ + h,w = src.shape[:2] + + keypoints = label['keypoints'] + center = label['center'] + + pad_ratio = random.uniform(max_move_ratio/2, max_move_ratio) + + move_x = random.randint(int(w*pad_ratio)//2,int(w*pad_ratio))*random.choice([-1,1]) + move_y = random.randint(int(w*pad_ratio)//2,int(h*pad_ratio))*random.choice([-1,1]) + + + M = np.float32([[1, 0, move_x], [0, 1, move_y]]) + src = cv2.warpAffine(src, M, (src.shape[1], src.shape[0])) + + + # new_img = src[move_y:move_y+new_h,move_x:move_x+new_w] + + + for i in range(len(keypoints)): + if i%3==0: + keypoints[i] = keypoints[i]+move_x/w + elif i%3==1: + keypoints[i] = keypoints[i]+move_y/h + for i in range(len(keypoints)//3): + if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: + keypoints[i*3+2] = 0 + + center[0] = min(max(0,(center[0]+move_x/w)),1) + center[1] = min(max(0,(center[1]+move_y/h)),1) + + + label["keypoints"]=keypoints + label["center"]=center + + + return src.astype(np.uint8), label + +def Rotate(src,angle,pad_color,label=None,center=None,scale=1.0): + ''' + :param src: src image + :param label: dict + :param angle: + :param center: + :param scale: + :return: the rotated image and the points + ''' + image=src + (h, w) = image.shape[:2] + + + # 若未指定旋转中心，则将图像中心设为旋转中心 + if center is None: + center = (w / 2, h / 2) + # 执行旋转 + M = cv2.getRotationMatrix2D(center, angle, scale) + if label is None: + for i in range(image.shape[2]): + image[:,:,i] = cv2.warpAffine(image[:,:,i], M, (w, h), + flags=cv2.INTER_CUBIC, + borderMode=cv2.BORDER_CONSTANT, + borderValue=pad_color) + return image,None + else: + + keypoints = label['keypoints'] + center = label['center'] + other_centers = label['other_centers'] + other_keypoints = label['other_keypoints'] + + + ####make it as a 3x3 RT matrix + full_M=np.row_stack((M,np.asarray([0,0,1]))) + img_rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, + borderMode=cv2.BORDER_CONSTANT, borderValue=pad_color) + ###make the keypoints as 3xN matrix + keypoints = np.reshape(keypoints, (-1,3)) + keypoints_z = keypoints[:,2].reshape((-1,1)) + keypoints = keypoints[:,:2] + keypoints[:,0] *= w + keypoints[:,1] *= h + keypoints = keypoints.astype(np.int32) + + keypoints = keypoints.T + full_keypoints = np.row_stack((keypoints, np.ones(shape=(1,keypoints.shape[1])))) + keypoints_rotated=np.dot(full_M,full_keypoints) + keypoints_rotated=keypoints_rotated[0:2,:] + #keypoints_rotated = keypoints_rotated.astype(np.int32) + keypoints_rotated=keypoints_rotated.T + + keypoints_rotated = keypoints_rotated.astype(np.float32) + keypoints_rotated[:,0] /=w + keypoints_rotated[:,1] /=h + # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) + keypoints = np.concatenate([keypoints_rotated,keypoints_z],-1).reshape((-1)).tolist() + for i in range(len(keypoints)//3): + if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: + keypoints[i*3+2] = 0 + + center = np.reshape(center, (-1,2)) + center[:,0] *= w + center[:,1] *= h + center = center.astype(np.int32) + + center = center.T + full_center = np.row_stack((center, np.ones(shape=(1,center.shape[1])))) + center_rotated=np.dot(full_M,full_center) + center_rotated=center_rotated[0:2,:] + #keypoints_rotated = keypoints_rotated.astype(np.int32) + center_rotated=center_rotated.T + + center_rotated = center_rotated.astype(np.float32) + center_rotated[:,0] /=w + center_rotated[:,1] /=h + # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) + center = center_rotated.reshape((-1)).tolist() + + + + other_centers = np.reshape(other_centers, (-1,2)) + other_centers[:,0] *= w + other_centers[:,1] *= h + other_centers = other_centers.astype(np.int32) + + other_centers = other_centers.T + full_center = np.row_stack((other_centers, np.ones(shape=(1,other_centers.shape[1])))) + center_rotated=np.dot(full_M,full_center) + center_rotated=center_rotated[0:2,:] + #keypoints_rotated = keypoints_rotated.astype(np.int32) + center_rotated=center_rotated.T + + center_rotated = center_rotated.astype(np.float32) + center_rotated[:,0] /= w + center_rotated[:,1] /= h + # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) + other_centers_raw = center_rotated.reshape((-1,2)).tolist() + other_centers = [] + for item in other_centers_raw: + if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: + continue + else: + other_centers.append(item) + + + for i in range(len(other_keypoints)): + if len(other_keypoints[i])>0: + other_keypointsi = np.reshape(other_keypoints[i], (-1,2)) + other_keypointsi[:,0] *= w + other_keypointsi[:,1] *= h + other_keypointsi = other_keypointsi.astype(np.int32) + + other_keypointsi = other_keypointsi.T + full_center = np.row_stack((other_keypointsi, np.ones(shape=(1,other_keypointsi.shape[1])))) + center_rotated=np.dot(full_M,full_center) + center_rotated=center_rotated[0:2,:] + #keypoints_rotated = keypoints_rotated.astype(np.int32) + center_rotated=center_rotated.T + + center_rotated = center_rotated.astype(np.float32) + center_rotated[:,0] /=w + center_rotated[:,1] /=h + # keypoints_rotated = np.reshape(keypoints_rotated, (-1)) + other_keypoints_i = center_rotated.reshape((-1,2)).tolist() + other_keypoints_res = [] + for item in other_keypoints_i: + if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: + continue + else: + other_keypoints_res.append(item) + other_keypoints[i] = other_keypoints_res + + label["keypoints"]= keypoints + label["center"]=center + label["other_centers"]=other_centers + label["other_keypoints"]=other_keypoints + + + return img_rotated, label + + +def Affine(src,strength,pad_color,label=None): + image = src + (h, w) = image.shape[:2] + + + keypoints = label['keypoints'] + center = label['center'] + other_centers = label['other_centers'] + other_keypoints = label['other_keypoints'] + + + pts_base = np.float32([[10,100],[200,50],[100,250]]) + pts1 = np.random.rand(3, 2) * -40 + pts_base + pts1 = pts1.astype(np.float32) + M = cv2.getAffineTransform(pts1, pts_base) + trans_img = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]) , + borderMode=cv2.BORDER_CONSTANT, + borderValue=pad_color) + if label is None: + return trans_img, label + + keypoints = np.reshape(keypoints, (-1,3)) + keypoints_z = keypoints[:,2].reshape((-1,1)) + keypoints = keypoints[:,:2] + keypoints[:,0] *= w + keypoints[:,1] *= h + keypoints = keypoints.astype(np.int32) + + keypoints=keypoints.T + full_data = np.row_stack((keypoints, np.ones(shape=(1, keypoints.shape[1])))) + data_rotated = np.dot(M, full_data) + #label_rotated = label_rotated.astype(np.int32) + data_rotated=data_rotated.T + + data_rotated = data_rotated.astype(np.float32) + data_rotated[:,0] /=w + data_rotated[:,1] /=h + keypoints = np.concatenate([data_rotated,keypoints_z],-1).reshape((-1)).tolist() + for i in range(len(keypoints)//3): + if keypoints[i*3]<0 or keypoints[i*3]>=1 or keypoints[i*3+1]<0 or keypoints[i*3+1]>=1: + keypoints[i*3+2] = 0 + + + center = np.reshape(center, (-1,2)) + center[:,0] *= w + center[:,1] *= h + center = center.astype(np.int32) + + center=center.T + full_data = np.row_stack((center, np.ones(shape=(1, center.shape[1])))) + data_rotated = np.dot(M, full_data) + #label_rotated = label_rotated.astype(np.int32) + data_rotated=data_rotated.T + + data_rotated = data_rotated.astype(np.float32) + data_rotated[:,0] /=w + data_rotated[:,1] /=h + center = data_rotated.reshape((-1)).tolist() + + + + other_centers = np.reshape(other_centers, (-1,2)) + other_centers[:,0] *= w + other_centers[:,1] *= h + other_centers = other_centers.astype(np.int32) + + other_centers = other_centers.T + full_data = np.row_stack((other_centers, np.ones(shape=(1, other_centers.shape[1])))) + data_rotated = np.dot(M, full_data) + #label_rotated = label_rotated.astype(np.int32) + data_rotated=data_rotated.T + + data_rotated = data_rotated.astype(np.float32) + data_rotated[:,0] /=w + data_rotated[:,1] /=h + other_centers_raw = data_rotated.reshape((-1,2)).tolist() + other_centers = [] + for item in other_centers_raw: + if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: + continue + else: + other_centers.append(item) + + + for i in range(len(other_keypoints)): + if len(other_keypoints[i])>0: + other_keypointsi = np.reshape(other_keypoints[i], (-1,2)) + other_keypointsi[:,0] *= w + other_keypointsi[:,1] *= h + other_keypointsi = other_keypointsi.astype(np.int32) + + other_keypointsi = other_keypointsi.T + full_data = np.row_stack((other_keypointsi, np.ones(shape=(1, other_keypointsi.shape[1])))) + data_rotated = np.dot(M, full_data) + #label_rotated = label_rotated.astype(np.int32) + data_rotated=data_rotated.T + + data_rotated = data_rotated.astype(np.float32) + data_rotated[:,0] /=w + data_rotated[:,1] /=h + # other_keypoints[i] = data_rotated.reshape((-1,2)).tolist() + other_keypoints_i = data_rotated.reshape((-1,2)).tolist() + other_keypoints_res = [] + for item in other_keypoints_i: + if item[0]<0 or item[0]>=1 or item[1]<0 or item[1]>=1: + continue + else: + other_keypoints_res.append(item) + other_keypoints[i] = other_keypoints_res + + + label["keypoints"]=keypoints + label["center"]=center + label["other_centers"]=other_centers + label["other_keypoints"]=other_keypoints + + return trans_img, label + + +def AID(img, label): + h,w = img.shape[:2] + + half_size = int(random.uniform(3/192,6/192)*(h+w)/2) + + + keypoints = np.array(label["keypoints"]).reshape((-1,3)) + + keypoints = keypoints[keypoints[:,2]>0] + dropout_id = random.randint(0,keypoints.shape[0]-1) + + + cx = keypoints[dropout_id][0]*w + cy = keypoints[dropout_id][1]*h + + x0 = int(max(0,cx-half_size)) + y0 = int(max(0,cy-half_size)) + x1 = int(min(w-1,cx+half_size)) + y1 = int(min(h-1,cy+half_size)) + + color = random.randint(0,255) + img[y0:y1, x0:x1] = (color,color,color) + + return img + +def AID2(img, label): + h,w = img.shape[:2] + + half_size = int(random.uniform(3/192,6/192)*(h+w)/2) + + + keypoints = np.array(label["keypoints"]).reshape((-1,3)) + + keypoints = keypoints[keypoints[:,2]>0] + dropout_id = random.randint(0,keypoints.shape[0]-1) + + + cx = keypoints[dropout_id][0]*w + cy = keypoints[dropout_id][1]*h + + x0 = int(max(0,cx-half_size)) + y0 = int(max(0,cy-half_size)) + x1 = int(min(w-1,cx+half_size)) + y1 = int(min(h-1,cy+half_size)) + # print(x0,y0,x1,y1) + # b + color = random.randint(0,255) + dist = half_size**2 + for i in range(y1-y0+1): + for j in range(x1-x0+1): + px = x0+j + py = y0+i + if ((py-cy)**2+(px-cx)**2)<=dist: + img[py, px] = (color,color,color) + + return img + +def dropout(src,pad_color,max_pattern_ratio=0.05): + width_ratio = random.uniform(0, max_pattern_ratio) + height_ratio = random.uniform(0, max_pattern_ratio) + width=src.shape[1] + height=src.shape[0] + block_width=width*width_ratio + block_height=height*height_ratio + width_start=int(random.uniform(0,width-block_width)) + width_end=int(width_start+block_width) + height_start=int(random.uniform(0,height-block_height)) + height_end=int(height_start+block_height) + src[height_start:height_end,width_start:width_end,:]=np.array(pad_color,dtype=src.dtype) + + return src + +def pixel_jitter(src,p=0.5,max_=5.): + + src=src.astype(np.float32) + + pattern=(np.random.rand(src.shape[0], src.shape[1],src.shape[2])-0.5)*2*max_ + img = src + pattern + + img[img<0]=0 + img[img >255] = 255 + + img = img.astype(np.uint8) + + return img + +def _clip(image): + """ + Clip and convert an image to np.uint8. + Args + image: Image to clip. + """ + return np.clip(image, 0, 255).astype(np.uint8) + +def adjust_contrast(image, factor): + """ Adjust contrast of an image. + Args + image: Image to adjust. + factor: A factor for adjusting contrast. + """ + mean = image.mean(axis=0).mean(axis=0) + return _clip((image - mean) * factor + mean) + +def adjust_brightness(image, delta): + """ Adjust brightness of an image + Args + image: Image to adjust. + delta: Brightness offset between -1 and 1 added to the pixel values. + """ + return _clip(image + delta * 255) + +def adjust_hue(image, delta): + """ Adjust hue of an image. + Args + image: Image to adjust. + delta: An interval between -1 and 1 for the amount added to the hue channel. + The values are rotated if they exceed 180. + """ + image[..., 0] = np.mod(image[..., 0] + delta * 180, 180) + return image + +def adjust_saturation(image, factor): + """ Adjust saturation of an image. + Args + image: Image to adjust. + factor: An interval for the factor multiplying the saturation values of each pixel. + """ + image[..., 1] = np.clip(image[..., 1] * factor, 0, 255) + return image + +class ColorDistort(): + def __init__( + self, + contrast_range=(0.8, 1.2), + brightness_range=(-.2, .2), + hue_range=(-0.1, 0.1), + saturation_range=(0.8, 1.2) + ): + self.contrast_range = contrast_range + self.brightness_range = brightness_range + self.hue_range = hue_range + self.saturation_range = saturation_range + + + def _uniform(self,val_range): + """ Uniformly sample from the given range. + Args + val_range: A pair of lower and upper bound. + """ + return np.random.uniform(val_range[0], val_range[1]) + + def __call__(self, image): + + + if self.contrast_range is not None: + contrast_factor = self._uniform(self.contrast_range) + image = adjust_contrast(image,contrast_factor) + if self.brightness_range is not None: + brightness_delta = self._uniform(self.brightness_range) + image = adjust_brightness(image, brightness_delta) + + if self.hue_range is not None or self.saturation_range is not None: + + image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + + if self.hue_range is not None: + hue_delta = self._uniform(self.hue_range) + image = adjust_hue(image, hue_delta) + + if self.saturation_range is not None: + saturation_factor = self._uniform(self.saturation_range) + image = adjust_saturation(image, saturation_factor) + + image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) + + return image + + + +###### Data Aug api +class DataAug: + def __init__(self, img_size): + self.h = img_size + self.w = img_size + self.color_augmentor = ColorDistort() + + def __call__(self, img, label): + """ + img: opencv img, BGR + item = { + "img_name":save_name, + "keypoints":save_keypoints, relative position + "center":save_center, + "other_centers":other_centers, + "other_keypoints":other_keypoints, + } + return: same type as img and label + """ + pad_color = random.randint(0,255) + + + new_img = deepcopy(img) + new_label = deepcopy(label) + + + if random.random() < 0.5: + new_img, new_label = Mirror(new_img, label=new_label) + + + rd = random.random() + if rd <0.5: + if random.random() < 0.2: + new_img, new_label = Padding(new_img, new_label, (0,0,0)) + else: + new_img, new_label = Crop(new_img, new_label, (0,0,0)) + elif rd < 0.65: + if random.random() < 0.7: + if len(new_label["other_centers"])==0: + strength = random.uniform(20, 50) + new_img, new_label = Affine(new_img, strength=strength, label=new_label, pad_color=(0,0,0)) + + else: + angle = random.uniform(-24, 24) + new_img, new_label = Rotate(new_img, label=new_label, angle=angle, pad_color=(0,0,0)) + + elif rd < 0.85: + + if len(new_label["other_centers"])==0 and (len(new_label["other_keypoints"][0]) + len(new_label["other_keypoints"][3]))==0: + if random.random() < 0.5: + new_img, new_label = Move(new_img, new_label, (0,0,0)) + else: + pass + ### in case of no points + count_zero = 0 + count_zero_new = 0 + for i in range(len(new_label["keypoints"])): + if i%3==2: + if label["keypoints"][i]==0: + count_zero+=1 + if new_label["keypoints"][i]==0: + count_zero_new+=1 + if count_zero_new-count_zero > 3 or count_zero_new>5: + new_img = img + new_label = label + + + + + if random.random() < 0.4: + new_img = self.color_augmentor(new_img) + + if random.random() < 0.4: + new_img = pixel_jitter(new_img,15) + + # if random.random() < 0.3: + # new_img = dropout(new_img, pad_color,0.1) + + if random.random() < 0.3: + + new_img = AID(new_img, new_label) + + + return new_img, new_label + + + + diff --git a/example/movenet/lib/data/data_tools.py b/pycore/moveenet/data/data_tools.py similarity index 99% rename from example/movenet/lib/data/data_tools.py rename to pycore/moveenet/data/data_tools.py index e671d1e1..a72aef0a 100644 --- a/example/movenet/lib/data/data_tools.py +++ b/pycore/moveenet/data/data_tools.py @@ -16,8 +16,8 @@ import platform import math -from lib.data.data_augment import DataAug -from lib.utils.utils import maxPoint, extract_keypoints +from pycore.moveenet.data.data_augment import DataAug +from pycore.moveenet.utils.utils import maxPoint, extract_keypoints def getFileNames(file_dir, tail_list=['.png', '.jpg', '.JPG', '.PNG']): diff --git a/pycore/moveenet/loss/__init__.py b/pycore/moveenet/loss/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example/movenet/lib/loss/movenet_loss.py b/pycore/moveenet/loss/movenet_loss.py similarity index 99% rename from example/movenet/lib/loss/movenet_loss.py rename to pycore/moveenet/loss/movenet_loss.py index 038a081f..6677d246 100644 --- a/example/movenet/lib/loss/movenet_loss.py +++ b/pycore/moveenet/loss/movenet_loss.py @@ -14,7 +14,7 @@ _img_size = 192 _feature_map_size = _img_size // 4 -_center_weight_path = '/usr/local/src/hpe-core/example/movenet/lib/data/center_weight_origin.npy' +_center_weight_path = '../data/center_weight_origin.npy' class JointBoneLoss(torch.nn.Module): diff --git a/pycore/moveenet/models/__init__.py b/pycore/moveenet/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example/movenet/lib/models/movenet_mobilenetv2.py b/pycore/moveenet/models/movenet_mobilenetv2.py similarity index 96% rename from example/movenet/lib/models/movenet_mobilenetv2.py rename to pycore/moveenet/models/movenet_mobilenetv2.py index ded41ad5..9a0f3594 100644 --- a/example/movenet/lib/models/movenet_mobilenetv2.py +++ b/pycore/moveenet/models/movenet_mobilenetv2.py @@ -1,375 +1,375 @@ -""" -@Fire -https://github.com/fire717 -""" -import torch -import torch.nn as nn -import math - - - - -""" -nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, - padding=0, dilation=1, groups=1, bias=True)) -""" - -def conv_3x3_act(inp, oup, stride): - return nn.Sequential( - nn.Conv2d(inp, oup, 3, stride, 1, bias=False), - nn.BatchNorm2d(oup), - nn.ReLU(inplace=True) - ) - - -def conv_1x1_act(inp, oup): - return nn.Sequential( - nn.Conv2d(inp, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup), - nn.ReLU(inplace=True) - ) - -def conv_1x1_act2(inp, oup): - return nn.Sequential( - nn.Conv2d(inp, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup), - nn.ReLU(inplace=True) - ) - - - -def dw_conv(inp, oup, stride=1): - return nn.Sequential( - nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), - nn.BatchNorm2d(inp), - nn.ReLU(inplace=True), - nn.Conv2d(inp, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup) - ) - -def dw_conv2(inp, oup, stride=1): - return nn.Sequential( - nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), - nn.BatchNorm2d(inp), - nn.Conv2d(inp, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup), - nn.ReLU(inplace=True), - ) - -def dw_conv3(inp, oup, stride=1): - return nn.Sequential( - nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), - nn.BatchNorm2d(inp), - nn.Conv2d(inp, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup), - nn.ReLU(inplace=True), - ) - - - -def upsample(inp, oup, scale=2): - return nn.Sequential( - nn.Conv2d(inp, inp, 3, 1, 1, groups=inp), - nn.ReLU(inplace=True), - conv_1x1_act2(inp,oup), - nn.Upsample(scale_factor=scale, mode='bilinear', align_corners=False)) - - -class InvertedResidual(nn.Module): - def __init__(self, inp, oup, stride, expand_ratio, n): - super(InvertedResidual, self).__init__() - assert stride in [1, 2] - - hidden_dim = round(inp * expand_ratio) - self.n = n - - - self.conv1 = nn.Sequential( - # pw - nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False), - nn.BatchNorm2d(hidden_dim), - nn.ReLU(inplace=True), - # dw - nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=True), - nn.BatchNorm2d(hidden_dim), - nn.ReLU(inplace=True), - # pw-linear - nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup), - ) - - self.conv2 = nn.Sequential( - # pw - nn.Conv2d(oup, hidden_dim, 1, 1, 0, bias=False), - nn.BatchNorm2d(hidden_dim), - nn.ReLU(inplace=False), - # dw - nn.Conv2d(hidden_dim, hidden_dim, 3, 1, 1, groups=hidden_dim, bias=False), - nn.BatchNorm2d(hidden_dim), - nn.ReLU(inplace=False), - # pw-linear - nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), - nn.BatchNorm2d(oup), - ) - - def forward(self, x): - x = self.conv1(x) - - for _ in range(self.n): - x = x + self.conv2(x) - - return x - - - -# class MulReshapeArgMax(nn.Module): -# def __init__(self): -# super(MulReshapeArgMax, self).__init__() - - -# self.weight = torch.reshape(weight, (1,48,48)) - -# def forward(self, x): -# # x n,1,48,48 -# #print(self.weight.shape) -# x = x.mul_(self.weight) -# x = torch.reshape(x, (-1,2304)) -# x = torch.argmax(x, dim=1, keepdim=True) - -# return x - - - - -class Backbone(nn.Module): - def __init__(self): - super(Backbone, self).__init__() - #mobilenet v2 - - - input_channel = 32 - - self.features1 = nn.Sequential(*[ - conv_3x3_act(3, input_channel, 2), - dw_conv(input_channel, 16, 1), - InvertedResidual(16, 24, 2, 6, 1) - ]) - - self.features2 = InvertedResidual(24, 32, 2, 6, 2) - self.features3 = InvertedResidual(32, 64, 2, 6, 3) - - self.features4 = nn.Sequential(*[ - InvertedResidual(64, 96, 1, 6, 2), - InvertedResidual(96, 160, 2, 6, 2), - InvertedResidual(160, 320, 1, 6, 0), - conv_1x1_act(320,1280), - nn.Conv2d(1280, 64, 1, 1, 0, bias=False), - nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) - ]) - - - self.upsample2 = upsample(64, 32) - self.upsample1 = upsample(32, 24) - - self.conv3 = nn.Conv2d(64, 64, 1, 1, 0) - self.conv2 = nn.Conv2d(32, 32, 1, 1, 0) - self.conv1 = nn.Conv2d(24, 24, 1, 1, 0) - - self.conv4 = dw_conv3(24, 24, 1) - - - - def forward(self, x): - x = x/127.5-1 - - - f1 = self.features1(x) - #print(f1.shape)#1, 24, 48, 48] - - f2 = self.features2(f1) - #print(f2.shape)#1, 32, 24, 24] - - f3 = self.features3(f2) - #print(f3.shape)#[1, 64, 12, 12] - - f4 = self.features4(f3) - f3 = self.conv3(f3) - #print(f4.shape)#[1, 64, 12, 12] - f4 += f3 - - f4 = self.upsample2(f4) - f2 = self.conv2(f2) - f4 += f2 - - f4 = self.upsample1(f4) - f1 = self.conv1(f1) - f4 += f1 - - f4 = self.conv4(f4) - - return f4 - - - -class Header(nn.Module): - def __init__(self, num_classes, mode='train'): - super(Header, self).__init__() - - self.mode = mode - - #heatmaps, centers, regs, offsets - #Person keypoint heatmap - self.header_heatmaps = nn.Sequential(*[ - dw_conv3(24, 96), - nn.Conv2d(96, num_classes, 1, 1, 0, bias=True), - nn.Sigmoid() - ]) - - #Person center heatmap - self.header_centers = nn.Sequential(*[ - dw_conv3(24, 96), - nn.Conv2d(96, 1, 1, 1, 0, bias=True), - nn.Sigmoid(), - # MulReshapeArgMax() - ]) - - #Keypoint regression field: - self.header_regs = nn.Sequential(*[ - dw_conv3(24, 96), - nn.Conv2d(96, num_classes*2, 1, 1, 0, bias=True), - ]) - - #2D per-keypoint offset field - self.header_offsets = nn.Sequential(*[ - dw_conv3(24, 96), - nn.Conv2d(96, num_classes*2, 1, 1, 0, bias=True), - ]) - - def argmax2loc(self, x, h=48, w=48): - ## n,1 - y0 = torch.div(x,w).long() - x0 = torch.sub(x, y0*w).long() - return x0,y0 - - - def forward(self, x): - - res = [] - if self.mode=='train': - h1 = self.header_heatmaps(x) - h2 = self.header_centers(x) - h3 = self.header_regs(x) - h4 = self.header_offsets(x) - res = [h1,h2,h3,h4] - - - elif self.mode=='test': - pass - - elif self.mode=='all': - #print("Header: ", x.shape) - pass - # h1 = self.head1(x) - # # print(h1.shape) # n,24,48,48 - - # h2 = self.head2(x) # n,1 - # # print(h2.shape,h2) - # x0,y0 = self.argmax2loc(h2) #n,1 - # # print(x0,y0,x0.shape,y0.shape) - - # # h2_1 = torch.cat([y0,x0], -1).long().unsqueeze(1) - # #torch.Tensor([[0] for _ in range(x.shape[0])]), - # # print(h2_1.shape, h2_1) # n,2 - - # h3 = self.head3(x) # n,34,48,48 - # h2_1 = [] - # for i in range(x.shape[0]): - # #print(h3[i,:,y0[i],x0[i]].shape) - # h2_1.append(h3[i,:,y0[i],x0[i]]) - # h2_1 = torch.cat(h2_1, dim=1).transpose(1,0).reshape((-1,17,2)) - # #print(h2_1.shape) # n,17,2 - - # h2_2 = h2_1[:,:,0].add_(y0) #n,17 - # #print(h2_2.shape) - - # h2_2_tmp = [] - # # print(self.range_tensor.shape) #48,48,17 - # for i in range(x.shape[0]): - # h2_2_t = self.range_tensor.sub_(h2_2[i]) - # print(h2_2_t.shape) - # h2_2_tmp.append(h2_2_t) - # h2_2 = torch.cat(h2_2_tmp, dim=0) - # print(h2_2.shape) - - # h2_1 = h2_1[:,:,1].add_(x0) - - - # # h2_1 = h3[h2_1] - # # h2_1 = torch.index_select(h3, 0, h2_1) - # print(h2_1.shape) #n,17 - - # h4 = self.head4(x) # n,34,48,48 - # # print(h4.shape) - # b - else: - print("[ERROR] wrong mode.") - - - - return res - - -class MoveNet(nn.Module): - def __init__(self, num_classes=17, width_mult=1.,mode='train'): - super(MoveNet, self).__init__() - - self.backbone = Backbone() - - self.header = Header(num_classes, mode) - - - self._initialize_weights() - - - def forward(self, x): - x = self.backbone(x) # n,24,48,48 - # print(x.shape) - - x = self.header(x) - # print([x0.shape for x0 in x]) - - return x - - - def _initialize_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - # m.weight.data.normal_(0, 0.01) - torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity='relu') - # torch.nn.init.xavier_normal_(m.weight.data) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - # elif isinstance(m, nn.Linear): - # m.weight.data.normal_(0, 0.01) - # m.bias.data.zero_() - - - - -if __name__ == "__main__": - from torchsummary import summary - - model = MoveNet().cuda() - print(summary(model, (3, 192, 192))) - - dummy_input1 = torch.randn(1, 3, 192, 192).cuda() - input_names = [ "input1"] #自己命名 - output_names = [ "output1" ] - - torch.onnx.export(model, dummy_input1, "pose.onnx", - verbose=True, input_names=input_names, output_names=output_names, +""" +@Fire +https://github.com/fire717 +""" +import torch +import torch.nn as nn +import math + + + + +""" +nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, + padding=0, dilation=1, groups=1, bias=True)) +""" + +def conv_3x3_act(inp, oup, stride): + return nn.Sequential( + nn.Conv2d(inp, oup, 3, stride, 1, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True) + ) + + +def conv_1x1_act(inp, oup): + return nn.Sequential( + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True) + ) + +def conv_1x1_act2(inp, oup): + return nn.Sequential( + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True) + ) + + + +def dw_conv(inp, oup, stride=1): + return nn.Sequential( + nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), + nn.BatchNorm2d(inp), + nn.ReLU(inplace=True), + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup) + ) + +def dw_conv2(inp, oup, stride=1): + return nn.Sequential( + nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), + nn.BatchNorm2d(inp), + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True), + ) + +def dw_conv3(inp, oup, stride=1): + return nn.Sequential( + nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), + nn.BatchNorm2d(inp), + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True), + ) + + + +def upsample(inp, oup, scale=2): + return nn.Sequential( + nn.Conv2d(inp, inp, 3, 1, 1, groups=inp), + nn.ReLU(inplace=True), + conv_1x1_act2(inp,oup), + nn.Upsample(scale_factor=scale, mode='bilinear', align_corners=False)) + + +class InvertedResidual(nn.Module): + def __init__(self, inp, oup, stride, expand_ratio, n): + super(InvertedResidual, self).__init__() + assert stride in [1, 2] + + hidden_dim = round(inp * expand_ratio) + self.n = n + + + self.conv1 = nn.Sequential( + # pw + nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False), + nn.BatchNorm2d(hidden_dim), + nn.ReLU(inplace=True), + # dw + nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=True), + nn.BatchNorm2d(hidden_dim), + nn.ReLU(inplace=True), + # pw-linear + nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + ) + + self.conv2 = nn.Sequential( + # pw + nn.Conv2d(oup, hidden_dim, 1, 1, 0, bias=False), + nn.BatchNorm2d(hidden_dim), + nn.ReLU(inplace=False), + # dw + nn.Conv2d(hidden_dim, hidden_dim, 3, 1, 1, groups=hidden_dim, bias=False), + nn.BatchNorm2d(hidden_dim), + nn.ReLU(inplace=False), + # pw-linear + nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + ) + + def forward(self, x): + x = self.conv1(x) + + for _ in range(self.n): + x = x + self.conv2(x) + + return x + + + +# class MulReshapeArgMax(nn.Module): +# def __init__(self): +# super(MulReshapeArgMax, self).__init__() + + +# self.weight = torch.reshape(weight, (1,48,48)) + +# def forward(self, x): +# # x n,1,48,48 +# #print(self.weight.shape) +# x = x.mul_(self.weight) +# x = torch.reshape(x, (-1,2304)) +# x = torch.argmax(x, dim=1, keepdim=True) + +# return x + + + + +class Backbone(nn.Module): + def __init__(self): + super(Backbone, self).__init__() + #mobilenet v2 + + + input_channel = 32 + + self.features1 = nn.Sequential(*[ + conv_3x3_act(3, input_channel, 2), + dw_conv(input_channel, 16, 1), + InvertedResidual(16, 24, 2, 6, 1) + ]) + + self.features2 = InvertedResidual(24, 32, 2, 6, 2) + self.features3 = InvertedResidual(32, 64, 2, 6, 3) + + self.features4 = nn.Sequential(*[ + InvertedResidual(64, 96, 1, 6, 2), + InvertedResidual(96, 160, 2, 6, 2), + InvertedResidual(160, 320, 1, 6, 0), + conv_1x1_act(320,1280), + nn.Conv2d(1280, 64, 1, 1, 0, bias=False), + nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) + ]) + + + self.upsample2 = upsample(64, 32) + self.upsample1 = upsample(32, 24) + + self.conv3 = nn.Conv2d(64, 64, 1, 1, 0) + self.conv2 = nn.Conv2d(32, 32, 1, 1, 0) + self.conv1 = nn.Conv2d(24, 24, 1, 1, 0) + + self.conv4 = dw_conv3(24, 24, 1) + + + + def forward(self, x): + x = x/127.5-1 + + + f1 = self.features1(x) + #print(f1.shape)#1, 24, 48, 48] + + f2 = self.features2(f1) + #print(f2.shape)#1, 32, 24, 24] + + f3 = self.features3(f2) + #print(f3.shape)#[1, 64, 12, 12] + + f4 = self.features4(f3) + f3 = self.conv3(f3) + #print(f4.shape)#[1, 64, 12, 12] + f4 += f3 + + f4 = self.upsample2(f4) + f2 = self.conv2(f2) + f4 += f2 + + f4 = self.upsample1(f4) + f1 = self.conv1(f1) + f4 += f1 + + f4 = self.conv4(f4) + + return f4 + + + +class Header(nn.Module): + def __init__(self, num_classes, mode='train'): + super(Header, self).__init__() + + self.mode = mode + + #heatmaps, centers, regs, offsets + #Person keypoint heatmap + self.header_heatmaps = nn.Sequential(*[ + dw_conv3(24, 96), + nn.Conv2d(96, num_classes, 1, 1, 0, bias=True), + nn.Sigmoid() + ]) + + #Person center heatmap + self.header_centers = nn.Sequential(*[ + dw_conv3(24, 96), + nn.Conv2d(96, 1, 1, 1, 0, bias=True), + nn.Sigmoid(), + # MulReshapeArgMax() + ]) + + #Keypoint regression field: + self.header_regs = nn.Sequential(*[ + dw_conv3(24, 96), + nn.Conv2d(96, num_classes*2, 1, 1, 0, bias=True), + ]) + + #2D per-keypoint offset field + self.header_offsets = nn.Sequential(*[ + dw_conv3(24, 96), + nn.Conv2d(96, num_classes*2, 1, 1, 0, bias=True), + ]) + + def argmax2loc(self, x, h=48, w=48): + ## n,1 + y0 = torch.div(x,w).long() + x0 = torch.sub(x, y0*w).long() + return x0,y0 + + + def forward(self, x): + + res = [] + if self.mode=='train': + h1 = self.header_heatmaps(x) + h2 = self.header_centers(x) + h3 = self.header_regs(x) + h4 = self.header_offsets(x) + res = [h1,h2,h3,h4] + + + elif self.mode=='test': + pass + + elif self.mode=='all': + #print("Header: ", x.shape) + pass + # h1 = self.head1(x) + # # print(h1.shape) # n,24,48,48 + + # h2 = self.head2(x) # n,1 + # # print(h2.shape,h2) + # x0,y0 = self.argmax2loc(h2) #n,1 + # # print(x0,y0,x0.shape,y0.shape) + + # # h2_1 = torch.cat([y0,x0], -1).long().unsqueeze(1) + # #torch.Tensor([[0] for _ in range(x.shape[0])]), + # # print(h2_1.shape, h2_1) # n,2 + + # h3 = self.head3(x) # n,34,48,48 + # h2_1 = [] + # for i in range(x.shape[0]): + # #print(h3[i,:,y0[i],x0[i]].shape) + # h2_1.append(h3[i,:,y0[i],x0[i]]) + # h2_1 = torch.cat(h2_1, dim=1).transpose(1,0).reshape((-1,17,2)) + # #print(h2_1.shape) # n,17,2 + + # h2_2 = h2_1[:,:,0].add_(y0) #n,17 + # #print(h2_2.shape) + + # h2_2_tmp = [] + # # print(self.range_tensor.shape) #48,48,17 + # for i in range(x.shape[0]): + # h2_2_t = self.range_tensor.sub_(h2_2[i]) + # print(h2_2_t.shape) + # h2_2_tmp.append(h2_2_t) + # h2_2 = torch.cat(h2_2_tmp, dim=0) + # print(h2_2.shape) + + # h2_1 = h2_1[:,:,1].add_(x0) + + + # # h2_1 = h3[h2_1] + # # h2_1 = torch.index_select(h3, 0, h2_1) + # print(h2_1.shape) #n,17 + + # h4 = self.head4(x) # n,34,48,48 + # # print(h4.shape) + # b + else: + print("[ERROR] wrong mode.") + + + + return res + + +class MoveNet(nn.Module): + def __init__(self, num_classes=17, width_mult=1.,mode='train'): + super(MoveNet, self).__init__() + + self.backbone = Backbone() + + self.header = Header(num_classes, mode) + + + self._initialize_weights() + + + def forward(self, x): + x = self.backbone(x) # n,24,48,48 + # print(x.shape) + + x = self.header(x) + # print([x0.shape for x0 in x]) + + return x + + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + # m.weight.data.normal_(0, 0.01) + torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity='relu') + # torch.nn.init.xavier_normal_(m.weight.data) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + # elif isinstance(m, nn.Linear): + # m.weight.data.normal_(0, 0.01) + # m.bias.data.zero_() + + + + +if __name__ == "__main__": + from torchsummary import summary + + model = MoveNet().cuda() + print(summary(model, (3, 192, 192))) + + dummy_input1 = torch.randn(1, 3, 192, 192).cuda() + input_names = [ "input1"] #自己命名 + output_names = [ "output1" ] + + torch.onnx.export(model, dummy_input1, "pose.onnx", + verbose=True, input_names=input_names, output_names=output_names, do_constant_folding=True,opset_version=11) \ No newline at end of file diff --git a/pycore/moveenet/task/__init__.py b/pycore/moveenet/task/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example/movenet/lib/task/task.py b/pycore/moveenet/task/task.py similarity index 97% rename from example/movenet/lib/task/task.py rename to pycore/moveenet/task/task.py index b32675b9..8dee9947 100644 --- a/example/movenet/lib/task/task.py +++ b/pycore/moveenet/task/task.py @@ -15,11 +15,11 @@ # import torch.nn.functional as F from torch.utils.tensorboard import SummaryWriter -from lib.task.task_tools import getSchedu, getOptimizer, movenetDecode, clipGradient, restore_sizes, image_show -from lib.loss.movenet_loss import MovenetLoss -from lib.utils.utils import printDash, ensure_loc -# from lib.visualization.visualization import superimpose_pose -from lib.utils.metrics import myAcc, pck +from pycore.moveenet.task.task_tools import getSchedu, getOptimizer, movenetDecode, clipGradient, restore_sizes, image_show +from pycore.moveenet.loss.movenet_loss import MovenetLoss +from pycore.moveenet.utils.utils import printDash, ensure_loc +# from pycore.moveenet.visualization.visualization import superimpose_pose +from pycore.moveenet.utils.metrics import myAcc, pck from datasets.h36m.utils.parsing import movenet_to_hpecore @@ -126,7 +126,7 @@ def predict_online(self, img_in): if self.cfg['show_center']: centers = output[1].cpu().numpy()[0] - from lib.utils.utils import maxPoint + from pycore.moveenet.utils.utils import maxPoint cx, cy = maxPoint(centers) instant['center'] = np.array([cx[0][0],cy[0][0]])/centers.shape[1] diff --git a/example/movenet/lib/task/task_tools.py b/pycore/moveenet/task/task_tools.py similarity index 98% rename from example/movenet/lib/task/task_tools.py rename to pycore/moveenet/task/task_tools.py index 2a2ab0aa..f1ef1c67 100644 --- a/example/movenet/lib/task/task_tools.py +++ b/pycore/moveenet/task/task_tools.py @@ -7,7 +7,7 @@ import numpy as np import cv2 -from lib.utils.utils import maxPoint, extract_keypoints +from pycore.moveenet.utils.utils import maxPoint, extract_keypoints from datasets.h36m.utils.parsing import movenet_to_hpecore from csv import writer @@ -15,7 +15,7 @@ _range_weight_y = _range_weight_x.T -# _reg_weight = np.load("lib/data/my_weight_reg.npy") # 99x99 +# _reg_weight = np.load("../data/my_weight_reg.npy") # 99x99 def getSchedu(schedu, optimizer): diff --git a/pycore/moveenet/utils/__init__.py b/pycore/moveenet/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/example/movenet/lib/utils/metrics.py b/pycore/moveenet/utils/metrics.py similarity index 100% rename from example/movenet/lib/utils/metrics.py rename to pycore/moveenet/utils/metrics.py diff --git a/example/movenet/lib/utils/utils.py b/pycore/moveenet/utils/utils.py similarity index 96% rename from example/movenet/lib/utils/utils.py rename to pycore/moveenet/utils/utils.py index a3122ebb..d188e26d 100644 --- a/example/movenet/lib/utils/utils.py +++ b/pycore/moveenet/utils/utils.py @@ -1,269 +1,269 @@ -""" -@Fire -https://github.com/fire717 -""" -import os -import torch -import random -import numpy as np -import argparse - -from operator import itemgetter - - -# from config import cfg - -def setRandomSeed(seed=42): - """Reproducer for pytorch experiment. - - Parameters - ---------- - seed: int, optional (default = 2019) - Radnom seed. - - Example - ------- - setRandomSeed(seed=2019). - """ - random.seed(seed) - os.environ["PYTHONHASHSEED"] = str(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - torch.backends.cudnn.enabled = True - - -def printDash(num=50): - print(''.join(['-'] * num)) - - -############## task tool - -# _center_weight = np.load('lib/data/center_weight.npy') -# _center_weight = np.reshape(_center_weight,(48,48)) -# # _center_weight = _center_weight.repeat((1,7,1,1)) -# def heatmap2locate(heatmap, size = cfg['img_size']//4): -# """ -# input: -# dist: [batchsize, 7, 64, 64] - -# return: -# dist: [batchsize, 32] x0,y0,x1,y1,... -# """ - -# heatmap = heatmap*_center_weight -# # print(heatmap.shape) -# # b -# bs = heatmap.shape[0] -# point_count = heatmap.shape[1] -# h,w = heatmap.shape[2:] - -# heatmap = np.reshape(heatmap, (bs,point_count,-1)) -# max_id = np.argmax(heatmap, 2)#64, 16 -# # print(max_id[0]) -# y = max_id//size -# y = y[:, : ,np.newaxis] -# # print(x.shape) -# x = max_id%size -# x = x[:, : ,np.newaxis] -# # print(y.shape) -# # print(x[0]) -# # print(y[0]) - -# res = np.concatenate((x,y), -1) -# # print(res[0]) -# # print(res.shape) - -# res = np.reshape(res, (bs,-1))/size -# # print(res.shape) -# # print(res[0]) -# # b -# return res - - -_center_weight = np.load('/usr/local/src/hpe-core/example/movenet/lib/data/center_weight_origin.npy').reshape(48, 48) - - -# _center_weight = np.load("lib/data/my_weight_center.npy") - -def maxPoint(heatmap, center=True): - if len(heatmap.shape) == 3: - batch_size, h, w = heatmap.shape - c = 1 - - elif len(heatmap.shape) == 4: - # n,c,h,w - batch_size, c, h, w = heatmap.shape - # print(heatmap.shape) - - if center: - # print(heatmap.shape) - # print(heatmap[0][27:31,27:31]) - # print(_center_weight[27:31,27:31]) - heatmap = heatmap * _center_weight # 加权取最靠近中间的 - # print(heatmap[0][27:31,27:31]) - - heatmap = heatmap.reshape((batch_size, c, -1)) # 64,c, 48x48 - # print(heatmap.shape) - # print(heatmap[0,0,23*48+20:23*48+30]) - max_id = np.argmax(heatmap, 2) # 64,c, 1 - # print(max_id) - # print(max_id, heatmap[0][0][max_id]) - # print(np.max(heatmap)) - y = max_id // w - x = max_id % w - # bv - return x, y - - -def extract_keypoints(heatmap): - """ - 热力图解析为关键点 - """ - # print(heatmap.shape)#(1, 48, 48) - # from light openpose - - heatmap = heatmap[0] - - heatmap[heatmap < 0.1] = 0 - # print(heatmap.shape)#500, 375 - heatmap_with_borders = np.pad(heatmap, [(2, 2), (2, 2)], mode='constant') - # 1.上下左右分别扩张2像素，然后分别取上下左右中5个不同的heatmap - # print(heatmap_with_borders.shape)#504, 379 - heatmap_center = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 1:heatmap_with_borders.shape[1] - 1] - heatmap_left = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 2:heatmap_with_borders.shape[1]] - heatmap_right = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 0:heatmap_with_borders.shape[1] - 2] - heatmap_up = heatmap_with_borders[2:heatmap_with_borders.shape[0], 1:heatmap_with_borders.shape[1] - 1] - heatmap_down = heatmap_with_borders[0:heatmap_with_borders.shape[0] - 2, 1:heatmap_with_borders.shape[1] - 1] - # print(heatmap_center.shape, heatmap_left.shape,heatmap_right.shape,heatmap_up.shape,heatmap_down.shape,) - # (502, 377) (502, 377) (502, 377) (502, 377) (502, 377 - heatmap_peaks = (heatmap_center > heatmap_left) & \ - (heatmap_center > heatmap_right) & \ - (heatmap_center > heatmap_up) & \ - (heatmap_center > heatmap_down) - # 2.通过center和周围比找峰值 - heatmap_peaks = heatmap_peaks[1:heatmap_center.shape[0] - 1, 1:heatmap_center.shape[1] - 1] - # print(heatmap_peaks.shape)#500, 375 - # print(heatmap_peaks[22:26,22:26]) - - keypoints = list(zip(np.nonzero(heatmap_peaks)[1], np.nonzero(heatmap_peaks)[0])) # (w, h) - # np.nonzero返回数组a中非零元素的索引值数组 - # keypoints = sorted(keypoints, key=itemgetter(0)) - keypoints = sorted(keypoints, key=lambda x: (x[0] - 23.5) ** 2 + (x[1] - 23.5) ** 2) - # print(keypoints) - - suppressed = np.zeros(len(keypoints), np.uint8) - - keypoints_with_score_and_id = [] - for i in range(len(keypoints)): - if suppressed[i]: - continue - for j in range(i + 1, len(keypoints)): - if ((keypoints[i][0] - keypoints[j][0]) ** 2 + - (keypoints[i][1] - keypoints[j][1]) ** 2) ** 0.5 < 6: - # 如果两点距离小于6则忽略 - suppressed[j] = 1 - keypoint_with_score_and_id = [keypoints[i][0], keypoints[i][1], - heatmap[keypoints[i][1], keypoints[i][0]], - _center_weight[keypoints[i][1], keypoints[i][0]], - heatmap[keypoints[i][1], keypoints[i][0]] * _center_weight[ - keypoints[i][1], keypoints[i][0]]] - # print(keypoint_with_score_and_id) - # (43, 1, 0.46907818, 0) # [x,y,heatmap[x,y],total_keypoint_num] - keypoints_with_score_and_id.append(keypoint_with_score_and_id) - - # bb - - # print(keypoints_with_score_and_id) - keypoints_with_score_and_id = sorted(keypoints_with_score_and_id, - key=lambda x: x[-1], reverse=True) - x, y = keypoints_with_score_and_id[0][0], keypoints_with_score_and_id[0][1] - return x, y - -def arg_parser(cfg): - parser = argparse.ArgumentParser() - ##### Global Setting - parser.add_argument('--show_center', help='Show center of the detected human', action='store_true') - parser.add_argument('--write_output', help='Write output pose to a .csv file', action='store_true') - parser.add_argument('--GPU_ID', help='GPUs to use. Example: 0,1,2', default=cfg['GPU_ID']) - parser.add_argument('--dataset', help='Training dataset.', default=cfg['dataset'], - choices=["mpii", "coco", 'h36m', 'dhp19'], type=str) - parser.add_argument('--num_workers', help='Number of workers', default=cfg['num_workers'], type=int) - parser.add_argument('--random_seed', help='Random seed', default=cfg['random_seed'], type=int) - parser.add_argument('--cfg_verbose', '-v', help='Verbosity', default=cfg['cfg_verbose'], type=str) - parser.add_argument('--save_dir', help='Directory to save model checkpoints', default=cfg['save_dir'], type=str) - parser.add_argument('--num_classes', help='Number of joints in the dataset', default=cfg['num_classes'], type=int) - parser.add_argument('--img_size', help='Image size of the square image for training and output', - default=cfg['img_size'], type=int) - parser.add_argument('--label', help='Label for the run', default=cfg['label'], type=str) - - ##### Train Setting - parser.add_argument('--pre-separated_data', - help='Set true if the training and validation annotations are in separate files', - default=cfg['pre-separated_data'], type=str) - parser.add_argument('--training_data_split', - help='Percentage of data to use for training. Not used if pre-separated set to True', - default=cfg['training_data_split'], type=int) - parser.add_argument('--newest_ckpt', help='File containing the name of the latest model checkpoints', - default=cfg['newest_ckpt'], type=str) - parser.add_argument('--balance_data', help='Set true for data balancing', default=cfg['balance_data'], type=str) - parser.add_argument('--log_interval', help='How frequently to log output', default=cfg['log_interval'], type=int) - parser.add_argument('--save_best_only', help='Save only the best or improved model checkpoints', - default=cfg['save_best_only'], type=str) - parser.add_argument('--pin_memory', help='Pin memory', default=cfg['pin_memory'], type=bool) - parser.add_argument('--th', help='Threshold value, in percentage of head size, for accuracy calculation', - default=cfg['th'], type=int) - parser.add_argument('--from_scratch', help='Set true to begin training from scratch', default=cfg['from_scratch'], - type=str) - - ##### Train Hyperparameters - parser.add_argument('--learning_rate', help='Initial learning rate of the training paradigm', - default=cfg['learning_rate'], type=float) - parser.add_argument('--batch_size', help='Batch size', default=cfg['batch_size'], type=int) - parser.add_argument('--epochs', help='Epochs', default=cfg['epochs'], type=int) - parser.add_argument('--optimizer', help='Optimizer', default=cfg['optimizer'], type=str, choices=['Adam', 'SGD']) - parser.add_argument('--scheduler', help='scheduler', default=cfg['scheduler'], type=str) - parser.add_argument('--weight_decay', help='weight decay', default=cfg['weight_decay'], type=float) - parser.add_argument('--class_weight', help='class weight', default=cfg['class_weight']) - parser.add_argument('--clip_gradient', help='Clip gradient', default=cfg['clip_gradient']) - parser.add_argument('--w_heatmap', help='Weight for the heatmap loss', default=cfg['w_heatmap'], type=int) - parser.add_argument('--w_bone', help='Weight for the Bone length loss', default=cfg['w_bone'], type=int) - parser.add_argument('--w_center', help='Weight for the center loss', default=cfg['w_center'], type=int) - parser.add_argument('--w_reg', help='Weight for the keypoint regression loss', default=cfg['w_reg'], type=int) - parser.add_argument('--w_offset', help='Weight for the offset loss', default=cfg['w_offset'], type=int) - - #### Directory paths - parser.add_argument('--img_path', help='Directory to the images for training and validation', - default=cfg['img_path'], type=str) - parser.add_argument('--train_label_path', help='File (full or relative path) to training annotation file', - default=cfg['train_label_path'], type=str) - parser.add_argument('--val_label_path', help='File (full or relative path) to validation annotation file', - default=cfg['val_label_path'], type=str) - parser.add_argument('--test_img_path', help='Directory to the images for test', default=cfg['test_img_path'], - type=str) - parser.add_argument('--predict_output_path', help='Directory to save images after prediction (without GT)', - default=cfg['predict_output_path'], type=str) - parser.add_argument('--eval_img_path', help='Directory to the images for evaluation', default=cfg['eval_img_path'], - type=str) - parser.add_argument('--eval_label_path', help='File (full or relative path) to evaluation annotation file', - default=cfg['eval_label_path'], type=str) - parser.add_argument('--results_path', help='Path to the location where the resultant .csv files should be stored.', - default=cfg['eval_label_path'], type=str) - - args = parser.parse_args() - - for key, value in vars(args).items(): - cfg[key] = value - - return cfg - -def ensure_loc(path): - # TODO: add functionality for filenames as well. - if os.path.isdir(path): - return True - else: - os.makedirs(path) +""" +@Fire +https://github.com/fire717 +""" +import os +import torch +import random +import numpy as np +import argparse + +from operator import itemgetter + + +# from config import cfg + +def setRandomSeed(seed=42): + """Reproducer for pytorch experiment. + + Parameters + ---------- + seed: int, optional (default = 2019) + Radnom seed. + + Example + ------- + setRandomSeed(seed=2019). + """ + random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.enabled = True + + +def printDash(num=50): + print(''.join(['-'] * num)) + + +############## task tool + +# _center_weight = np.load('moveenet/data/center_weight.npy') +# _center_weight = np.reshape(_center_weight,(48,48)) +# # _center_weight = _center_weight.repeat((1,7,1,1)) +# def heatmap2locate(heatmap, size = cfg['img_size']//4): +# """ +# input: +# dist: [batchsize, 7, 64, 64] + +# return: +# dist: [batchsize, 32] x0,y0,x1,y1,... +# """ + +# heatmap = heatmap*_center_weight +# # print(heatmap.shape) +# # b +# bs = heatmap.shape[0] +# point_count = heatmap.shape[1] +# h,w = heatmap.shape[2:] + +# heatmap = np.reshape(heatmap, (bs,point_count,-1)) +# max_id = np.argmax(heatmap, 2)#64, 16 +# # print(max_id[0]) +# y = max_id//size +# y = y[:, : ,np.newaxis] +# # print(x.shape) +# x = max_id%size +# x = x[:, : ,np.newaxis] +# # print(y.shape) +# # print(x[0]) +# # print(y[0]) + +# res = np.concatenate((x,y), -1) +# # print(res[0]) +# # print(res.shape) + +# res = np.reshape(res, (bs,-1))/size +# # print(res.shape) +# # print(res[0]) +# # b +# return res + + +_center_weight = np.load('/usr/local/src/hpe-core/example/movenet/moveenet/data/center_weight_origin.npy').reshape(48, 48) + + +# _center_weight = np.load("moveenet/data/my_weight_center.npy") + +def maxPoint(heatmap, center=True): + if len(heatmap.shape) == 3: + batch_size, h, w = heatmap.shape + c = 1 + + elif len(heatmap.shape) == 4: + # n,c,h,w + batch_size, c, h, w = heatmap.shape + # print(heatmap.shape) + + if center: + # print(heatmap.shape) + # print(heatmap[0][27:31,27:31]) + # print(_center_weight[27:31,27:31]) + heatmap = heatmap * _center_weight # 加权取最靠近中间的 + # print(heatmap[0][27:31,27:31]) + + heatmap = heatmap.reshape((batch_size, c, -1)) # 64,c, 48x48 + # print(heatmap.shape) + # print(heatmap[0,0,23*48+20:23*48+30]) + max_id = np.argmax(heatmap, 2) # 64,c, 1 + # print(max_id) + # print(max_id, heatmap[0][0][max_id]) + # print(np.max(heatmap)) + y = max_id // w + x = max_id % w + # bv + return x, y + + +def extract_keypoints(heatmap): + """ + 热力图解析为关键点 + """ + # print(heatmap.shape)#(1, 48, 48) + # from light openpose + + heatmap = heatmap[0] + + heatmap[heatmap < 0.1] = 0 + # print(heatmap.shape)#500, 375 + heatmap_with_borders = np.pad(heatmap, [(2, 2), (2, 2)], mode='constant') + # 1.上下左右分别扩张2像素，然后分别取上下左右中5个不同的heatmap + # print(heatmap_with_borders.shape)#504, 379 + heatmap_center = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 1:heatmap_with_borders.shape[1] - 1] + heatmap_left = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 2:heatmap_with_borders.shape[1]] + heatmap_right = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 0:heatmap_with_borders.shape[1] - 2] + heatmap_up = heatmap_with_borders[2:heatmap_with_borders.shape[0], 1:heatmap_with_borders.shape[1] - 1] + heatmap_down = heatmap_with_borders[0:heatmap_with_borders.shape[0] - 2, 1:heatmap_with_borders.shape[1] - 1] + # print(heatmap_center.shape, heatmap_left.shape,heatmap_right.shape,heatmap_up.shape,heatmap_down.shape,) + # (502, 377) (502, 377) (502, 377) (502, 377) (502, 377 + heatmap_peaks = (heatmap_center > heatmap_left) & \ + (heatmap_center > heatmap_right) & \ + (heatmap_center > heatmap_up) & \ + (heatmap_center > heatmap_down) + # 2.通过center和周围比找峰值 + heatmap_peaks = heatmap_peaks[1:heatmap_center.shape[0] - 1, 1:heatmap_center.shape[1] - 1] + # print(heatmap_peaks.shape)#500, 375 + # print(heatmap_peaks[22:26,22:26]) + + keypoints = list(zip(np.nonzero(heatmap_peaks)[1], np.nonzero(heatmap_peaks)[0])) # (w, h) + # np.nonzero返回数组a中非零元素的索引值数组 + # keypoints = sorted(keypoints, key=itemgetter(0)) + keypoints = sorted(keypoints, key=lambda x: (x[0] - 23.5) ** 2 + (x[1] - 23.5) ** 2) + # print(keypoints) + + suppressed = np.zeros(len(keypoints), np.uint8) + + keypoints_with_score_and_id = [] + for i in range(len(keypoints)): + if suppressed[i]: + continue + for j in range(i + 1, len(keypoints)): + if ((keypoints[i][0] - keypoints[j][0]) ** 2 + + (keypoints[i][1] - keypoints[j][1]) ** 2) ** 0.5 < 6: + # 如果两点距离小于6则忽略 + suppressed[j] = 1 + keypoint_with_score_and_id = [keypoints[i][0], keypoints[i][1], + heatmap[keypoints[i][1], keypoints[i][0]], + _center_weight[keypoints[i][1], keypoints[i][0]], + heatmap[keypoints[i][1], keypoints[i][0]] * _center_weight[ + keypoints[i][1], keypoints[i][0]]] + # print(keypoint_with_score_and_id) + # (43, 1, 0.46907818, 0) # [x,y,heatmap[x,y],total_keypoint_num] + keypoints_with_score_and_id.append(keypoint_with_score_and_id) + + # bb + + # print(keypoints_with_score_and_id) + keypoints_with_score_and_id = sorted(keypoints_with_score_and_id, + key=lambda x: x[-1], reverse=True) + x, y = keypoints_with_score_and_id[0][0], keypoints_with_score_and_id[0][1] + return x, y + +def arg_parser(cfg): + parser = argparse.ArgumentParser() + ##### Global Setting + parser.add_argument('--show_center', help='Show center of the detected human', action='store_true') + parser.add_argument('--write_output', help='Write output pose to a .csv file', action='store_true') + parser.add_argument('--GPU_ID', help='GPUs to use. Example: 0,1,2', default=cfg['GPU_ID']) + parser.add_argument('--dataset', help='Training dataset.', default=cfg['dataset'], + choices=["mpii", "coco", 'h36m', 'dhp19'], type=str) + parser.add_argument('--num_workers', help='Number of workers', default=cfg['num_workers'], type=int) + parser.add_argument('--random_seed', help='Random seed', default=cfg['random_seed'], type=int) + parser.add_argument('--cfg_verbose', '-v', help='Verbosity', default=cfg['cfg_verbose'], type=str) + parser.add_argument('--save_dir', help='Directory to save model checkpoints', default=cfg['save_dir'], type=str) + parser.add_argument('--num_classes', help='Number of joints in the dataset', default=cfg['num_classes'], type=int) + parser.add_argument('--img_size', help='Image size of the square image for training and output', + default=cfg['img_size'], type=int) + parser.add_argument('--label', help='Label for the run', default=cfg['label'], type=str) + + ##### Train Setting + parser.add_argument('--pre-separated_data', + help='Set true if the training and validation annotations are in separate files', + default=cfg['pre-separated_data'], type=str) + parser.add_argument('--training_data_split', + help='Percentage of data to use for training. Not used if pre-separated set to True', + default=cfg['training_data_split'], type=int) + parser.add_argument('--newest_ckpt', help='File containing the name of the latest model checkpoints', + default=cfg['newest_ckpt'], type=str) + parser.add_argument('--balance_data', help='Set true for data balancing', default=cfg['balance_data'], type=str) + parser.add_argument('--log_interval', help='How frequently to log output', default=cfg['log_interval'], type=int) + parser.add_argument('--save_best_only', help='Save only the best or improved model checkpoints', + default=cfg['save_best_only'], type=str) + parser.add_argument('--pin_memory', help='Pin memory', default=cfg['pin_memory'], type=bool) + parser.add_argument('--th', help='Threshold value, in percentage of head size, for accuracy calculation', + default=cfg['th'], type=int) + parser.add_argument('--from_scratch', help='Set true to begin training from scratch', default=cfg['from_scratch'], + type=str) + + ##### Train Hyperparameters + parser.add_argument('--learning_rate', help='Initial learning rate of the training paradigm', + default=cfg['learning_rate'], type=float) + parser.add_argument('--batch_size', help='Batch size', default=cfg['batch_size'], type=int) + parser.add_argument('--epochs', help='Epochs', default=cfg['epochs'], type=int) + parser.add_argument('--optimizer', help='Optimizer', default=cfg['optimizer'], type=str, choices=['Adam', 'SGD']) + parser.add_argument('--scheduler', help='scheduler', default=cfg['scheduler'], type=str) + parser.add_argument('--weight_decay', help='weight decay', default=cfg['weight_decay'], type=float) + parser.add_argument('--class_weight', help='class weight', default=cfg['class_weight']) + parser.add_argument('--clip_gradient', help='Clip gradient', default=cfg['clip_gradient']) + parser.add_argument('--w_heatmap', help='Weight for the heatmap loss', default=cfg['w_heatmap'], type=int) + parser.add_argument('--w_bone', help='Weight for the Bone length loss', default=cfg['w_bone'], type=int) + parser.add_argument('--w_center', help='Weight for the center loss', default=cfg['w_center'], type=int) + parser.add_argument('--w_reg', help='Weight for the keypoint regression loss', default=cfg['w_reg'], type=int) + parser.add_argument('--w_offset', help='Weight for the offset loss', default=cfg['w_offset'], type=int) + + #### Directory paths + parser.add_argument('--img_path', help='Directory to the images for training and validation', + default=cfg['img_path'], type=str) + parser.add_argument('--train_label_path', help='File (full or relative path) to training annotation file', + default=cfg['train_label_path'], type=str) + parser.add_argument('--val_label_path', help='File (full or relative path) to validation annotation file', + default=cfg['val_label_path'], type=str) + parser.add_argument('--test_img_path', help='Directory to the images for test', default=cfg['test_img_path'], + type=str) + parser.add_argument('--predict_output_path', help='Directory to save images after prediction (without GT)', + default=cfg['predict_output_path'], type=str) + parser.add_argument('--eval_img_path', help='Directory to the images for evaluation', default=cfg['eval_img_path'], + type=str) + parser.add_argument('--eval_label_path', help='File (full or relative path) to evaluation annotation file', + default=cfg['eval_label_path'], type=str) + parser.add_argument('--results_path', help='Path to the location where the resultant .csv files should be stored.', + default=cfg['eval_label_path'], type=str) + + args = parser.parse_args() + + for key, value in vars(args).items(): + cfg[key] = value + + return cfg + +def ensure_loc(path): + # TODO: add functionality for filenames as well. + if os.path.isdir(path): + return True + else: + os.makedirs(path) From 8fc660ae005b5d92b7bd7b9c9455fc6414d290d9 Mon Sep 17 00:00:00 2001 From: Gaurvi Goyal Date: Wed, 25 Jan 2023 15:00:19 +0100 Subject: [PATCH 2/4] Minor change to accommodate the refactor. --- example/movenet/movenet_online.py | 2 +- example/op_detector_example_module/Dockerfile | 4 ++-- pycore/moveenet/loss/movenet_loss.py | 5 +++-- pycore/moveenet/utils/utils.py | 2 +- 4 files changed, 7 insertions(+), 6 deletions(-) diff --git a/example/movenet/movenet_online.py b/example/movenet/movenet_online.py index 14ad0541..5fb7c985 100644 --- a/example/movenet/movenet_online.py +++ b/example/movenet/movenet_online.py @@ -1,4 +1,4 @@ -from lib import init, MoveNet, Task +from pycore.moveenet import init, MoveNet, Task from config import cfg from pycore.moveenet.utils.utils import arg_parser diff --git a/example/op_detector_example_module/Dockerfile b/example/op_detector_example_module/Dockerfile index c000f480..0f5c2942 100644 --- a/example/op_detector_example_module/Dockerfile +++ b/example/op_detector_example_module/Dockerfile @@ -235,7 +235,7 @@ RUN python3 -m pip install --upgrade pip && \ python3 -m pip install torch==1.9.1 torchvision==0.10.1 pandas scipy opencv-python #ENV PYTHONHOME=/root/anaconda/bin -#ENV PYTHONPATH=/root/anaconda/moveenet:/root/anaconda/moveenet/python3.8:/root/anaconda/moveenet/python3.8/site-packages +#ENV PYTHONPATH=/root/anaconda/lib:/root/anaconda/lib/python3.8:/root/anaconda/lib/python3.8/site-packages #ENV PATH=$PYTHONHOME:$PATH # download hpe-core @@ -244,7 +244,7 @@ RUN cd $SOURCE_FOLDER && \ cd hpe-core && \ git checkout $HPE_BRANCH -# add /usr/local/moveenet to the library path, so that libcaffe.so compiled with openpose will be used +# add /usr/local/lib to the library path, so that libcaffe.so compiled with openpose will be used # instead of the one provided by the nvidia/cuda docker image ENV LD_LIBRARY_PATH /usr/local/lib:$LD_LIBRARY_PATH diff --git a/pycore/moveenet/loss/movenet_loss.py b/pycore/moveenet/loss/movenet_loss.py index 6677d246..829f6939 100644 --- a/pycore/moveenet/loss/movenet_loss.py +++ b/pycore/moveenet/loss/movenet_loss.py @@ -2,7 +2,7 @@ @Fire https://github.com/fire717 """ -import sys +import os.path import torch import math @@ -14,7 +14,8 @@ _img_size = 192 _feature_map_size = _img_size // 4 -_center_weight_path = '../data/center_weight_origin.npy' + +_center_weight_path = os.path.abspath(os.path.join(os.path.dirname(__file__),'../data/center_weight_origin.npy')) class JointBoneLoss(torch.nn.Module): diff --git a/pycore/moveenet/utils/utils.py b/pycore/moveenet/utils/utils.py index d188e26d..ad65b66e 100644 --- a/pycore/moveenet/utils/utils.py +++ b/pycore/moveenet/utils/utils.py @@ -85,7 +85,7 @@ def printDash(num=50): # return res -_center_weight = np.load('/usr/local/src/hpe-core/example/movenet/moveenet/data/center_weight_origin.npy').reshape(48, 48) +_center_weight = os.path.abspath(np.load(os.path.join(os.path.dirname(__file__),'../data/center_weight_origin.npy'))).reshape(48, 48) # _center_weight = np.load("moveenet/data/my_weight_center.npy") From 79f986889208c36243e5fd9ad2f24daf3601deb1 Mon Sep 17 00:00:00 2001 From: Gaurvi Goyal Date: Thu, 2 Feb 2023 14:07:12 +0100 Subject: [PATCH 3/4] Minor change to accommodate the refactor. --- example/movenet/evaluate.py | 1 + pycore/moveenet/config.py | 59 ++++++++++++++++++++++++++++++++++ pycore/moveenet/task/task.py | 39 +++++++++++++--------- pycore/moveenet/utils/utils.py | 3 +- 4 files changed, 86 insertions(+), 16 deletions(-) create mode 100644 pycore/moveenet/config.py diff --git a/example/movenet/evaluate.py b/example/movenet/evaluate.py index 923d4699..5842d219 100644 --- a/example/movenet/evaluate.py +++ b/example/movenet/evaluate.py @@ -7,6 +7,7 @@ import sys sys.path.append('.') +sys.path.append('../../../hpe-core') from pycore.moveenet import init, Data, MoveNet, Task from config import cfg diff --git a/pycore/moveenet/config.py b/pycore/moveenet/config.py new file mode 100644 index 00000000..2e28e060 --- /dev/null +++ b/pycore/moveenet/config.py @@ -0,0 +1,59 @@ +""" +@Fire +https://github.com/fire717 +""" + +cfg = { + ##### Global Setting + 'GPU_ID': '0,1', + "num_workers": 4, + "random_seed": 42, + "cfg_verbose": True, + "save_dir": '', + "num_classes": 13, + "width_mult": 1.0, + "img_size": 192, + 'label': '', + + ##### Train Setting + 'pre-separated_data': True, + 'training_data_split': 80, + "dataset": '', + 'balance_data': False, + 'log_interval': 10, + 'save_best_only': False, + + 'pin_memory': True, + 'newest_ckpt': '', + 'th': 50, # percentage of headsize + 'from_scratch': True, + + ##### Train Hyperparameters + 'learning_rate': 0.001, # 1.25e-4 + 'batch_size': 64, + 'epochs': 300, + 'optimizer': 'Adam', # Adam SGD + 'scheduler': 'MultiStepLR-70,100-0.1', # default SGDR-5-2 CVPR step-4-0.8 MultiStepLR + # multistepLR-<<>milestones>-<> + 'weight_decay': 0.001, # 5.e-4, # 0.0001, + 'class_weight': None, # [1., 1., 1., 1., 1., 1., 1., ] + 'clip_gradient': 5, # 1, + 'w_heatmap': 1, + 'w_bone': 20, + 'w_center': 1, + 'w_reg': 3, + 'w_offset': 1, + + ##### File paths + 'predict_output_path': '', + 'results_path': '', + "img_path": '', + "test_img_path": '', + "eval_img_path": '', + "eval_outputs": '', + "eval_label_path": '', + 'train_label_path': '', + 'val_label_path': '', + 'ckpt': '' +} + diff --git a/pycore/moveenet/task/task.py b/pycore/moveenet/task/task.py index 8dee9947..81848054 100644 --- a/pycore/moveenet/task/task.py +++ b/pycore/moveenet/task/task.py @@ -99,7 +99,7 @@ def predict(self, data_loader, save_dir): cv2.imwrite(os.path.join(save_dir, basename[:-4] + "_regs0.jpg"), cv2.resize(regs[0] * 255, (size, size))) - def predict_online(self, img_in): + def predict_online(self, img_in, write_csv=None, ts = None): self.model.eval() correct = 0 @@ -121,32 +121,38 @@ def predict_online(self, img_in): # img_size_original = img.shape img = img.to(self.device) + start_sample = time.time() output = self.model(img) instant = {} - if self.cfg['show_center']: - centers = output[1].cpu().numpy()[0] - from pycore.moveenet.utils.utils import maxPoint - cx, cy = maxPoint(centers) - instant['center'] = np.array([cx[0][0],cy[0][0]])/centers.shape[1] + try: + if self.cfg['show_center']: + centers = output[1].cpu().numpy()[0] + from pycore.moveenet.utils.utils import maxPoint + cx, cy = maxPoint(centers) + instant['center'] = np.array([cx[0][0],cy[0][0]])/centers.shape[1] + except KeyError: + pass pre = movenetDecode(output, None, mode='output', num_joints=self.cfg["num_classes"]) img = np.transpose(img[0].cpu().numpy(), axes=[1, 2, 0]) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) _, pose_pre = restore_sizes(img, pre, (int(img_size_original[0]), int(img_size_original[1]))) - if self.cfg['show_center']: - instant['center_heatmap'] = centers[0] - + try: + if self.cfg['show_center']: + instant['center_heatmap'] = centers[0] + except KeyError: + pass kps_2d = np.reshape(pose_pre, [-1, 2]) kps_hpecore = movenet_to_hpecore(kps_2d) kps_pre_hpecore = np.reshape(kps_hpecore, [-1]) # print(kps_pre_hpecore) - # row = self.create_row(ts,kps_pre_hpecore, delay=time.time()-start_sample) - # sample = '_'.join(os.path.basename(img_names[0]).split('_')[:-1]) - # write_path = os.path.join(self.cfg['results_path'],self.cfg['dataset'],sample,'movenet_cam2.csv') - # ensure_loc(os.path.dirname(write_path)) - # self.write_results(write_path, row) + if write_csv is not None: + # print('writing') + row = self.create_row(ts,kps_pre_hpecore, delay=time.time()-start_sample) + ensure_loc(os.path.dirname(write_csv)) + self.write_results(write_csv, row) instant['joints'] = kps_pre_hpecore return instant @@ -315,7 +321,10 @@ def modelLoad(self, model_path, data_parallel=False): init_epoch = int(str1.join(os.path.basename(model_path).split('_')[0][1:])) self.init_epoch = init_epoch print(model_path) - self.model.load_state_dict(torch.load(model_path, map_location=self.device)) + try: + self.model.load_state_dict(torch.load(model_path, map_location=self.device)) + except FileNotFoundError: + print('The checkpoint expected does not exist. Please check the path and filename.') if data_parallel: self.model = torch.nn.DataParallel(self.model) diff --git a/pycore/moveenet/utils/utils.py b/pycore/moveenet/utils/utils.py index ad65b66e..481634e2 100644 --- a/pycore/moveenet/utils/utils.py +++ b/pycore/moveenet/utils/utils.py @@ -85,7 +85,7 @@ def printDash(num=50): # return res -_center_weight = os.path.abspath(np.load(os.path.join(os.path.dirname(__file__),'../data/center_weight_origin.npy'))).reshape(48, 48) +_center_weight = np.load(os.path.abspath(os.path.join(os.path.dirname(__file__),'../data/center_weight_origin.npy'))).reshape(48, 48) # _center_weight = np.load("moveenet/data/my_weight_center.npy") @@ -210,6 +210,7 @@ def arg_parser(cfg): default=cfg['training_data_split'], type=int) parser.add_argument('--newest_ckpt', help='File containing the name of the latest model checkpoints', default=cfg['newest_ckpt'], type=str) + parser.add_argument('--ckpt', help='Filename with location to the checkpoints', default=cfg['ckpt'], type=str) parser.add_argument('--balance_data', help='Set true for data balancing', default=cfg['balance_data'], type=str) parser.add_argument('--log_interval', help='How frequently to log output', default=cfg['log_interval'], type=int) parser.add_argument('--save_best_only', help='Save only the best or improved model checkpoints', From ebe49d57cd76887d8f8aefe56c220c95bfc7f025 Mon Sep 17 00:00:00 2001 From: Gaurvi Goyal Date: Wed, 8 Feb 2023 17:15:29 +0100 Subject: [PATCH 4/4] bugfix --- example/movenet/config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/example/movenet/config.py b/example/movenet/config.py index 19ba4648..51bd915d 100644 --- a/example/movenet/config.py +++ b/example/movenet/config.py @@ -55,6 +55,7 @@ "eval_outputs": home + '/dhp-outputs/', "eval_label_path": home + "/dhp19_eros/poses.json", 'train_label_path': '', - 'val_label_path': '' + 'val_label_path': '', + 'ckpt': '' }