pose_model_train.py

# -*- coding: utf-8 -*-
"""pose_model_train.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/121KVgt1ju5orPimS-1TiuZK_c9faRYBR
"""

import sys
import numpy as np
import cv2
import math
import os, random

import torch.optim
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from model.unipose import unipose

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
from sklearn.externals import joblib


# Load pretrained unipose model
model = unipose(dataset='COCO', num_classes=16, backbone='resnet', output_stride=16, sync_bn=True, freeze_bn=False, stride=8)
model = model.cuda()
checkpoint = torch.load('pretrained/UniPose_COCO.pth')
p = checkpoint 

state_dict = model.state_dict()
model_dict = {}

for k,v in p.items():
    if k in state_dict:
        model_dict[k] = v

state_dict.update(model_dict)
model.load_state_dict(state_dict)


# Get keypoints from maps
def get_kpts(maps, img_h = 368.0, img_w = 368.0):
    maps = maps.clone().cpu().data.numpy()
    map_6 = maps[0]

    kpts = []
    for m in map_6[1:]:
        h, w = np.unravel_index(m.argmax(), m.shape)
        x = int(w * img_w / m.shape[1])
        y = int(h * img_h / m.shape[0])
        kpts.append([x,y])
    return kpts


# Get keypoint coordinates from img
def unipose_write(img_path):
    center   = [184, 184]
    img  = np.array(cv2.resize(cv2.imread(img_path),(368,368)), dtype=np.float32)
    img  = img.transpose(2, 0, 1)
    img  = torch.from_numpy(img)
    mean = [128.0, 128.0, 128.0]
    std  = [256.0, 256.0, 256.0]
    for t, m, s in zip(img, mean, std):
        t.sub_(m).div_(s)

    img = torch.unsqueeze(img, 0)

    model.eval()

    input_var = img.cuda()
    heat = model(input_var)
    heat = F.interpolate(heat, size=input_var.size()[2:], mode='bilinear', align_corners=True)

    kpts = get_kpts(heat, img_h=368.0, img_w=368.0)

    im = cv2.resize(cv2.imread(img_path),(368,368))
    for i, kpt in enumerate(kpts):
        x, y = kpt
        cv2.circle(im, (x, y), 5, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
        cv2.putText(im, str(i), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1, lineType=cv2.LINE_AA)
    
    return im, kpts


# Calculate degree by coordinates
def calculate_degree(point1, point2, point3):
    o1 = math.atan2((point1[1]-point2[1]),(point1[0]-point2[0]))
    o2 = math.atan2((point3[1]-point2[1]),(point3[0]-point2[0]))

    deg1=abs((o1-o2)*180/math.pi)
    deg2=360-deg1

    return min(deg1,deg2)


# Get a list of angles for each joint 
def angular_calculate(input_points):
    elbow_left_degree=calculate_degree(input_points[8],input_points[6],input_points[4]) # 6
    elbow_right_degree=calculate_degree(input_points[9],input_points[7],input_points[5]) # 7
    shoulder_left_degree=calculate_degree(input_points[6],input_points[4],input_points[10]) #4
    shoulder_right_degree=calculate_degree(input_points[7],input_points[5],input_points[11]) #5
    pelvis_left_degree=calculate_degree(input_points[4],input_points[10],input_points[12]) #10
    pelvis_right_degree=calculate_degree(input_points[5],input_points[11],input_points[13]) #11
    knee_left_degree=calculate_degree(input_points[10],input_points[12],input_points[14]) #12
    knee_right_degree=calculate_degree(input_points[11],input_points[13],input_points[15]) #13
    
    points_angular_list = [elbow_left_degree, elbow_right_degree, shoulder_left_degree, 
                          shoulder_right_degree, pelvis_left_degree, pelvis_right_degree,
                          knee_left_degree, knee_right_degree]
    
    return points_angular_list


# Make dataset for each class
def load_dataset(class_name, base_dir='dataset'):
    X_list, Y_list = [], []
    
    # load positive data
    positive_img_dir = os.path.join(base_dir, class_name)
    positive_img_list = os.listdir(positive_img_dir)
    
    for positive_img in positive_img_list:
      positive_img = os.path.join(positive_img_dir, positive_img)
      temp = []
      img, points=unipose_write(positive_img)
      for i,j in points:
        temp.append(i)
        temp.append(j)
      temp += angular_calculate(points)
      X_list.append(temp)
      Y_list.append(1)

    # load negative data
    for wrong_class in os.listdir(base_dir):
        if wrong_class == class_name:
            continue
        negative_img_dir = os.path.join(base_dir, wrong_class)
        negative_img_list = os.listdir(negative_img_dir)

        n = len(X_list) // 4
        while n >= 0:
            negative_img = random.choice(negative_img_list)
            negative_img = os.path.join(negative_img_dir, negative_img)
            temp = []
            img, points=unipose_write(negative_img)
            for i,j in points:
              temp.append(i)
              temp.append(j)
            temp += angular_calculate(points)
            X_list.append(temp)
            Y_list.append(0)
            n -= 1

    X_train, X_test, y_train, y_test = train_test_split(X_list, Y_list, random_state=42)
    
    return X_train, X_test, y_train, y_test


def train_model(X_train, y_train, class_name):
    # feature Scaling
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)

    # training and predictions
    classifier = SVC(probability=True)
    classifier.fit(X_train, y_train)

    # calculate train accuracy
    y_pred = classifier.predict(X_train)
    train_accu = accuracy_score(y_train, y_pred)
    print(f'{class_name} 클래스 학습 정확도: {train_accu}')

    return classifier, scaler


def test_model(X_test, y_test, classifier, scaler, class_name):
    # feature scaling
    X_test = scaler.transform(X_test)

    # predict
    y_pred = classifier.predict(X_test)
    prob = classifier.predict_proba(X_test)

    # calculate test accuracy
    test_accu = accuracy_score(y_test, y_pred)
    print(f'{class_name} 클래스 예측 정확도: {test_accu}')

    return prob, y_pred


def finalize_model(X_train, X_test, y_train, y_test, class_name, model_dir='classifier'):
    # Use All Data
    X_train += X_test
    y_train += y_test

    # Feature Scaling
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)

    # Training
    classifier = SVC(probability=True)
    classifier.fit(X_train, y_train)

    # Save Model
    model_name = class_name + '_model.pkl'
    model_name = os.path.join(model_dir, model_name)
    joblib.dump(classifier, model_name)
    print(f'{class_name} 클래스 분류 모델 저장 완료')

    # Save Scaler
    scaler_name = class_name + '_scaler.pkl'
    scaler_name = os.path.join(model_dir, scaler_name)
    joblib.dump(scaler, scaler_name)
    print(f'{class_name} 클래스 분류 스케일러 저장 완료')

    return

# if __name__ == "__main__":
#     for class_name in ['plank', 'pullup', 'pushup', 'squat']:
#     X_train, X_test, y_train, y_test = load_dataset(class_name=class_name)
#     finalize_model(X_train, X_test, y_train, y_test, class_name)