# -*- coding: UTF-8 -*-
'''
调用格式：
root@xxx:/centerface/prj-python# python demo.py --image_dir ./testImages --output_dir ./res/

验证是否正确：
python draw.py

'''
import cv2
from centerface import CenterFace

import argparse
import os
from six.moves import xrange
from lxml import etree, objectify
# import numpy as np

def camera():
    cap = cv2.VideoCapture(0)
    ret, frame = cap.read()
    h, w = frame.shape[:2]
    centerface = CenterFace(h, w)
    while True:
        ret, frame = cap.read()
        dets, lms = centerface(frame, threshold=0.5)
        for det in dets:
            boxes, score = det[:4], det[4]
            cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
        for lm in lms:
            cv2.circle(frame, (int(lm[0]), int(lm[1])), 2, (0, 0, 255), -1)
            cv2.circle(frame, (int(lm[2]), int(lm[3])), 2, (0, 0, 255), -1)
            cv2.circle(frame, (int(lm[4]), int(lm[5])), 2, (0, 0, 255), -1)
            cv2.circle(frame, (int(lm[6]), int(lm[7])), 2, (0, 0, 255), -1)
            cv2.circle(frame, (int(lm[8]), int(lm[9])), 2, (0, 0, 255), -1)
        cv2.imshow('out', frame)
        # Press Q on keyboard to stop recording
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()


def test_image():
    frame = cv2.imread('000388.jpg')
    h, w = frame.shape[:2]
    centerface = CenterFace(h, w)
    dets, lms = centerface(frame, threshold=0.35)
    for det in dets:
        boxes, score = det[:4], det[4]
        cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
    for lm in lms:
        cv2.circle(frame, (int(lm[0]), int(lm[1])), 2, (0, 0, 255), -1)
        cv2.circle(frame, (int(lm[2]), int(lm[3])), 2, (0, 0, 255), -1)
        cv2.circle(frame, (int(lm[4]), int(lm[5])), 2, (0, 0, 255), -1)
        cv2.circle(frame, (int(lm[6]), int(lm[7])), 2, (0, 0, 255), -1)
        cv2.circle(frame, (int(lm[8]), int(lm[9])), 2, (0, 0, 255), -1)
    cv2.imshow('out', frame)
    cv2.waitKey(0)


# -------------------------------------------------------------------------------------------------------------------------
IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']

def is_image_file(filename):
    """Checks if a file is an image.
      Args:
          filename (string): path to a file
      Returns:
          bool: True if the filename ends with a known image extension
    """
    filename_lower = filename.lower()
    return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS)

def get_imagelist_from_dir(dirpath):
    images = []
    for f in os.listdir(dirpath):
        if is_image_file(f):
            images.append(os.path.join(dirpath, f))
    return images



def parse_args():
    """Parse in command line arguments"""
    parser = argparse.ArgumentParser(description='Demonstrate centerface results')
   
    # parser.add_argument(
    #     '--cfg', dest='cfg_file', required=True,
    #     help='optional config file')

    # parser.add_argument('--load_ckpt', help='path of checkpoint to load')
    
    parser.add_argument(
        '--image_dir',
        help='directory to load images for demo')
    parser.add_argument(
        '--images', nargs='+',
        help='images to infer. Must not use with --image_dir')
    parser.add_argument(
        '--output_dir',
        help='directory to save demo results',
        default="infer_outputs")
    # parser.add_argument(
    #     '--vis',
    #     help='save visualized result or not',
    #     type=bool,
    #     default=False)

    args = parser.parse_args()

    return args


def test_images(imglist,output_dir):

    num_images = len(imglist)

    for i in xrange(num_images):
        # print('img{}:{}'.format(i,imglist[i]))
        # im = cv2.imread(imglist[i])
        frame = cv2.imread(imglist[i])
        myboxes=[]
        assert frame is not None
        h, w = frame.shape[:2]
        centerface = CenterFace(h, w)
        dets, lms = centerface(frame, threshold=0.35)
        for det in dets:
            boxes, score = det[:4], det[4]
            cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
            # print( (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])))
            box=['Faces', 0, (int(boxes[0]),int(boxes[1]),int(boxes[2]),int(boxes[3]))]
            myboxes.append(box)
            
        # for lm in lms:
        #     cv2.circle(frame, (int(lm[0]), int(lm[1])), 2, (0, 0, 255), -1)
        #     cv2.circle(frame, (int(lm[2]), int(lm[3])), 2, (0, 0, 255), -1)
        #     cv2.circle(frame, (int(lm[4]), int(lm[5])), 2, (0, 0, 255), -1)
        #     cv2.circle(frame, (int(lm[6]), int(lm[7])), 2, (0, 0, 255), -1)
        #     cv2.circle(frame, (int(lm[8]), int(lm[9])), 2, (0, 0, 255), -1)
        cv2.imshow('out', frame)
        cv2.imwrite(output_dir+os.path.basename(imglist[i]).split('.')[0] + '.jpg' ,frame)
        write_xml(imglist[i],w,h,myboxes,output_dir)
        cv2.waitKey(500)
    
def write_xml(image_name,width,height,boxes,output_dir):
    """
    // 根据李韬提供的seemo程序，人车分离 type 定义
    static const int UNKNOWN = 0;               // 未知
    static const int PEDESTRAIN = 1;            // 行人
    static const int BIKE = 2;                  // 自行车
    static const int MOTOBIKE = 3;              // 摩托车，电动车
    static const int CAR = 4;                   // 小汽车
    static const int TRICYCLE = 5;              // 三轮车
    static const int BUS = 6;                   // 大巴
    static const int VAN = 7;                   // 面包车
    static const int TRUCK = 8;                 // 卡车
    """
    face_label_2_number_dict={0: 'face'}
    number_2_label_dict = {'bicycle': 0, 'car': 1, 'bus': 2, 'motorbike': 3, 'person': 4, 'truck': 5,
                           'tricycle': 6, 'plate': 7, 'face': 8}
    label_2_number_dict = {0: 'bicycle', 1: 'car', 2: 'bus', 3: 'motorbike', 4: 'person', 5: 'truck', 6: 'tricycle',
                           7: 'plate', 8:'face'}
    E = objectify.ElementMaker(annotate=False)
    anno_tree = E.annotation(
        E.folder('VOC2007'),
        E.filename(image_name),
        E.source(
            E.database('DETRAC'),
            E.annotation('DETRAC'),
            E.image('DETRAC'),
            E.url('None')
        ),
        E.size(
            E.width(width),
            E.height(height),
            E.depth(3)
        ),
        E.segmented(0),
    )
    E = objectify.ElementMaker(annotate=False)
    for bbox in boxes:
        if bbox[0]=='Faces':
            label=face_label_2_number_dict[bbox[1]]
        if label != 'unkown':
            number = number_2_label_dict[label]
            xmin = max(int(bbox[2][0]), 0)
            ymin = max(int(bbox[2][1]), 0)
            # xmax = min(int(bbox[2][0]) + int(bbox[2][2]), width)
            # ymax = min(int(bbox[2][1]) + int(bbox[2][3]), height)
            xmax = min(int(bbox[2][2]), width)
            ymax = min(int(bbox[2][3]), height)
            label = label_2_number_dict[number]
            anno_tree.append(E.object(
                E.name(label),
                E.bndbox(
                    E.xmin(xmin),
                    E.ymin(ymin),
                    E.xmax(xmax),
                    E.ymax(ymax)
                ),
                E.difficult(0),
                E.occlusion(0)
            )
            )
    target_xml_file = output_dir+os.path.basename(image_name).split('.')[0] + '.xml'
    print(target_xml_file)
    etree.ElementTree(anno_tree).write(target_xml_file, pretty_print=True)



if __name__ == '__main__':
    camera()
    # test_image()


    # args = parse_args()

    # if args.image_dir:
    #     imglist = get_imagelist_from_dir(args.image_dir)
    # else:
    #     imglist = args.images
    
    # if not os.path.exists(args.output_dir):
    #     os.makedirs(args.output_dir)
    
    # test_images(imglist,args.output_dir)
    
