exercise_2/myTelloProject-master/myScripts/AlphaPose_control.py

import av
import numpy
import tellopy
import cv2
import os
import json
import math
import time
import sys
import traceback

import torch
from torch.autograd import Variable
import torch.nn.functional as F
import torchvision.transforms as transforms

import torch.nn as nn
import torch.utils.data
import numpy as np
from AlphaPose.opt import opt

from AlphaPose.dataloader import ImageLoader, DetectionLoader, DetectionProcessor, DataWriter, Mscoco
from AlphaPose.yolo.util import write_results, dynamic_write_results
from AlphaPose.SPPE.src.main_fast_inference import *

import os
import sys
from tqdm import tqdm
import time
from AlphaPose.fn import getTime

class  AlphaPose_control:
    # 设置参数
    args = opt
    args.inputpath = '../AlphaPose/duan_alphapose/photo/'
    args.outputpath = '../AlphaPose/duan_alphapose/'
    args.sp = True
    args.dataset = 'coco'

    img_path = args.inputpath

    if not args.sp:
        torch.multiprocessing.set_start_method('forkserver', force=True)
        torch.multiprocessing.set_sharing_strategy('file_system')
    def __init__(self):
        pass

    # 获取图像
    def GetImage(self,args):
        inputpath = args.inputpath
        inputlist = args.inputlist
        mode = args.mode

        if not os.path.exists(args.outputpath):
            os.mkdir(args.outputpath)

        for root, dirs, files in os.walk(inputpath):
            im_names = files
        return im_names

    # 下载模型
    def downModel(self):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # Load pose model
        print('Loading YOLO model...')
        pose_dataset = Mscoco()
        if self.args.fast_inference:
            pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
        else:
            pose_model = InferenNet(4 * 1 + 1, pose_dataset)
        pose_model.to(device)
        pose_model.eval()

        return pose_model

    # 处理图像，提取关键点
    def Alphapose(self,im_names, pose_model, ):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # Load input images
        data_loader = ImageLoader(im_names, batchSize=self.args.detbatch, format='yolo').start()

        # Load detection loader
        sys.stdout.flush()
        det_loader = DetectionLoader(data_loader, batchSize=self.args.detbatch).start()
        det_processor = DetectionProcessor(det_loader).start()
        runtime_profile = {
            'dt': [],
            'pt': [],
            'pn': []
        }

        # Init data writer
        writer = DataWriter(self.args.save_video).start()

        data_len = data_loader.length()
        im_names_desc = tqdm(range(data_len))

        batchSize = self.args.posebatch
        for i in im_names_desc:
            start_time = getTime()
            with torch.no_grad():
                (inps, orig_img, im_name, boxes, scores, pt1, pt2) = det_processor.read()
                if boxes is None or boxes.nelement() == 0:
                    writer.save(None, None, None, None, None, orig_img, im_name.split('/')[-1])
                    continue
                ckpt_time, det_time = getTime(start_time)
                runtime_profile['dt'].append(det_time)
                # Pose Estimation

                datalen = inps.size(0)
                leftover = 0
                if (datalen) % batchSize:
                    leftover = 1
                num_batches = datalen // batchSize + leftover
                hm = []
                for j in range(num_batches):
                    inps_j = inps[j * batchSize:min((j + 1) * batchSize, datalen)].to(device)
                    hm_j = pose_model(inps_j)
                    hm.append(hm_j)
                hm = torch.cat(hm)
                ckpt_time, pose_time = getTime(ckpt_time)
                runtime_profile['pt'].append(pose_time)
                hm = hm.cpu()
                writer.save(boxes, scores, hm, pt1, pt2, orig_img, im_name.split('/')[-1])

                ckpt_time, post_time = getTime(ckpt_time)
                runtime_profile['pn'].append(post_time)

            if self.args.profile:
                # TQDM
                im_names_desc.set_description(
                    'det time: {dt:.3f} | pose time: {pt:.2f} | post processing: {pn:.4f}'.format(
                        dt=np.mean(runtime_profile['dt']), pt=np.mean(runtime_profile['pt']),
                        pn=np.mean(runtime_profile['pn']))
                )

        print('Finish Model Running.')
        if (self.args.save_img or self.args.save_video) and not self.args.vis_fast:
            print('===========================> Rendering remaining images in the queue...')
            print(
                '===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).')
        while (writer.running()):
            pass
        writer.stop()
        final_result = writer.results()
        # write_json(final_result, args.outputpath)
        try:
            if final_result[0]['result']:
                return final_result[0]['result'][0]['keypoints']
            else:
                return None
        except:
            return None

    # 根据关键点分析动作，绘制图像
    def PoseFind(self,point_results):
        '''point_result是关键点信息，一个tensor数组
        {0,  "Nose"},
        {1,  "LEye"},
        {2,  "REye"},
        {3,  "LEar"},
        {4,  "REar"},
        {5,  "LShoulder"},
        {6,  "RShoulder"},
        {7,  "LElbow"},
        {8,  "RElbow"},
        {9,  "LWrist"},
        {10, "RWrist"},
        {11, "LHip"},
        {12, "RHip"},
        {13, "LKnee"},
        {14, "Rknee"},
        {15, "LAnkle"},
        {16, "RAnkle"},
        '''

        # 登记LH左手、LE左肘、LS左肩、RH右手、RE右肘、RS右肩数据
        # LS_X = int(numbers[0]['keypoints'][15])
        LS_Y = int(point_results[5][1].item())
        # RS_X = int(numbers[0]['keypoints'][18])
        RS_Y = int(point_results[6][1].item())

        # LE_X = int(numbers[0]['keypoints'][21])
        LE_Y = int(point_results[7][1].item())
        # RE_X = int(numbers[0]['keypoints'][24])
        RE_Y = int(point_results[8][1].item())

        # LH_X = int(numbers[0]['keypoints'][27])
        LH_Y = int(point_results[9][1].item())
        # RH_X = int(numbers[0]['keypoints'][30])
        RH_Y = int(point_results[10][1].item())

        # 以双眼间距的两倍作为参照
        Leye_x = int(point_results[1][0].item())
        Leye_y = int(point_results[1][1].item())
        Reye_x = int(point_results[2][0].item())
        Reye_y = int(point_results[2][1].item())
        len = int(math.sqrt(math.pow(Leye_x - Reye_x, 2) + math.pow(Leye_y - Reye_y, 2))) * 2

        # 判断姿势
        pose = 0
        if LH_Y - LE_Y >= len and LE_Y - LS_Y >= len and RH_Y - RE_Y >= len and RE_Y - RS_Y >= len:
            pose = 1  # 双垂
            print('---------------------------双垂--------------------------')
        if LE_Y - LH_Y >= len and abs(LE_Y - LS_Y) <= len and RE_Y - RH_Y >= len and abs(RE_Y - RS_Y) <= len:
            pose = 2  # 双平举
            print('---------------------------双平举--------------------------')
        if LS_Y - LE_Y >= len and LE_Y - LH_Y >= len and RS_Y - RE_Y >= len and RE_Y - RH_Y >= len:
            pose = 3  # 双高举
            print('---------------------------双高举--------------------------')
        if abs(LH_Y - LE_Y) <= len and abs(LE_Y - LS_Y) <= len and abs(RH_Y - RE_Y) <= len and abs(RE_Y - RS_Y) <= len:
            pose = 4  # 双伸
            print('---------------------------双伸--------------------------')
        if abs(LH_Y - LE_Y) <= len and abs(LE_Y - LS_Y) <= len and RE_Y - RH_Y >= len and abs(RE_Y - RS_Y) <= len:
            pose = 5  # 左伸右平举
            print('---------------------------左伸右平举--------------------------')
        if abs(LH_Y - LE_Y) <= len and abs(LE_Y - LS_Y) <= len and RH_Y - RE_Y >= len and RE_Y - RS_Y >= len:
            pose = 6  # 左伸右垂
            print('---------------------------左伸右垂--------------------------')
        if LE_Y - LH_Y >= len and abs(LE_Y - LS_Y) <= len and abs(RH_Y - RE_Y) <= len and abs(RE_Y - RS_Y) <= len:
            pose = 7  # 左平举右伸
            print('---------------------------左平举右伸--------------------------')
        if LH_Y - LE_Y >= len and LE_Y - LS_Y >= len and abs(RH_Y - RE_Y) <= len and abs(RE_Y - RS_Y) <= len:
            pose = 8  # 左垂右伸
            print('---------------------------左垂右伸--------------------------')
        if LS_Y - LE_Y >= len and LE_Y - LH_Y >= len and RH_Y - RE_Y >= len and RE_Y - RS_Y >= len:
            pose = 9  # 左高举右垂
            print('---------------------------左高举右垂--------------------------')
        if LH_Y - LE_Y >= len and LE_Y - LS_Y >= len and RS_Y - RE_Y >= len and RE_Y - RH_Y >= len:
            pose = 10  # 右高举左垂
            print('---------------------------右高举左垂--------------------------')
        else:
            pass

            # 绘制关键点可视化
            # 0鼻子，1左眼，2右眼，3左耳，4右耳，5左肩，6右肩，7左肘，8右肘，9左腕，10右腕，11左臀，12右臀，13左膝，14右膝，15左踝，16右踝
            point_list = [(int(point_results[0][0].item()), int(point_results[0][1].item())),
                          (int(point_results[1][0].item()), int(point_results[1][1].item())),
                          (int(point_results[2][0].item()), int(point_results[2][1].item())),
                          (int(point_results[3][0].item()), int(point_results[3][1].item())),
                          (int(point_results[4][0].item()), int(point_results[4][1].item())),
                          (int(point_results[5][0].item()), int(point_results[5][1].item())),
                          (int(point_results[6][0].item()), int(point_results[6][1].item())),
                          (int(point_results[7][0].item()), int(point_results[7][1].item())),
                          (int(point_results[8][0].item()), int(point_results[8][1].item())),
                          (int(point_results[9][0].item()), int(point_results[9][1].item())),
                          (int(point_results[10][0].item()), int(point_results[10][1].item())),
                          (int(point_results[11][0].item()), int(point_results[11][1].item())),
                          (int(point_results[12][0].item()), int(point_results[12][1].item())),
                          (int(point_results[13][0].item()), int(point_results[13][1].item())),
                          (int(point_results[14][0].item()), int(point_results[14][1].item())),
                          (int(point_results[15][0].item()), int(point_results[15][1].item())),
                          (int(point_results[16][0].item()), int(point_results[16][1].item()))]
            img = cv2.imread(self.img_path + 'frame.jpg')
            for point in point_list:
                cv2.circle(img, point, 1, (0, 0, 255), 4)

            # 绘制关键点连线
            cv2.line(img, point_list[0], point_list[1], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[0], point_list[2], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[1], point_list[3], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[2], point_list[4], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[9], point_list[7], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[7], point_list[5], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[5], point_list[6], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[6], point_list[8], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[8], point_list[10], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[5], point_list[11], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[6], point_list[12], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[11], point_list[12], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[11], point_list[13], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[13], point_list[15], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[12], point_list[14], (0, 255, 0), 1, 4)
            cv2.line(img, point_list[14], point_list[16], (0, 255, 0), 1, 4)

            # cv2.imshow('AlphaPose', img)
            # cv2.waitKey(1)
            return pose,img

    # 删除文件函数，用于清空文件夹内图像
    def del_files(self,path_file):
        ls = os.listdir(path_file)
        for i in ls:
            f_path = os.path.join(path_file, i)
            # 判断是否是一个目录,若是,则递归删除
            if os.path.isdir(f_path):
                self.del_files(f_path)
            else:
                os.remove(f_path)