计算分类结果.py

# -*- coding: gbk -*-
# usd on windows
"""
Created on Thu Oct 16 09:14:42 2014

@author: zhangjie
"""

import os
import sys


CONFIG_TXT = "config.txt"
DEBUG = False


def load_config(fp):
    """加载配置文件
    """
    config = dict()
    with open(fp, 'r') as fd:
        for line in fd.readlines():
            # 滤去空行
            if line.strip(' ').strip('\n').strip('\t') == '':
                continue
            # 读取 key value
            line = line.strip('\n')
            key, value = line.split(' ')
            config[key] = value
    # logging
    if DEBUG:
        for key, value in config.items():
            print key, ' --> ', value
    return config


class PeopleList(object):
    """
        people {
            name:string #人名
            face_num:int
            data:[{
                klass:string
                number:int
            }]
        }
    """
    
    def __init__(self):
        self.peoples = []
    
    def add_people(self, name):
        """
            添加一张人脸
        """
        assert(name != '')
        for people in self.peoples:
            if people['name'] == name:
                people['face_num'] += 1
                return None
        people = dict(name=name, face_num=0, data=[])
        self.peoples.append(people)
        return None
    
    def add_klass(self, people_name, klass_name, number=1):
        """
            添加人脸所在类别，一次添加一个
        """
        for people in self.peoples:
            if people['name'] == people_name:
                # search classes
                for k in people['data']:
                    if k['klass'] == klass_name:
                        k['number'] += number
                        return None
                people['data'].append(dict(klass=klass_name, number=number))
                return None
        raise Exception(' '.jon([people_name, "can't be found in people_list"]))
    
    def get_poeple(self, name):
        for people in self.peoples:
            if people['name'] == name:
                return poeple
        return None
        
    def get_peoples(self):
        return self.peoples


class KlassList(object):
    """
        klass {
            name:string #分类名
            face_num:int
            data:[{
                people:string
                number:int
            }]
        }
    """
    
    def __init__(self):
        self.klasses = []
    
    def add_klass(self, klass_name):
        for klass in self.klasses:
            if klass['name'] == klass_name:
                klass['face_num'] += 1
                return klass
        klass = dict(name=klass_name, face_num=1, data=[])
        self.klasses.append(klass)
        return klass
    
    def add_people(self, klass_name, people_name, number=1):
        for klass in self.klasses:
            if klass['name'] == klass_name:
                # search people
                for rv in klass['data']:
                    if rv['people'] == people_name:
                        rv['number'] += number
                        return None
                klass['data'].append(dict(people=people_name, number=number))
                return None
        raise Exception(' '.jon([klass_name, "can't be found in klass_list"]))

    def get_klass(self, name):
        for klass in self.klasses:
            if klass['name'] == name:
                return klass
        return None


people_list = PeopleList()
klass_list = KlassList()


def read_people(people_dir, result_dir):
    """
        读取人脸目录来统计每个人有几张脸
        读取分类结果目录来统计每个人的分类情况
        
        
    """
    if not os.path.exists(people_dir):
        raise Exception(' '.join([people_dir, "doesn't exists"]))
    if not os.path.exists(result_dir):
        raise Exception(' '.join([result_dir, "doesn't exists"]))

    # 统计每个人的人脸数
    for fname in os.listdir(people_dir):
        fp = os.path.join(people_dir, fname)
        if os.path.isfile(fp) and os.path.splitext(fname)[1] == '.jpg':
            # 取jpg文件
            # A_001_001.jpg
            people_name = fname.split('-')[0]
            people_list.add_people(people_name)

    # 统计人脸分类情况
    for fname in os.listdir(result_dir):
        fp = os.path.join(result_dir, fname)
        if os.path.isdir(fp):
            # 进入分类文件夹
            # 分类文件夹下仅有人脸图片 A_001_001.jpg
            for name in os.listdir(fp):
                assert(os.path.splitext(name)[1] == '.jpg')
                people_name = name.split('-')[0]
                klass_list.add_klass(fname)
                people_list.add_klass(people_name, fname)
                klass_list.add_people(fname, people_name)


def get_result(out_dir):
    """
        统计识别率等数值
    """
    with open(os.path.join(out_dir, 'statistics.txt'), 'w') as fd:
        for people in people_list.get_peoples():
            words = [people['name'], '拥有', str(people['face_num']), '张脸']
            fd.write(' '.join(words))
            fd.write('\n')
            
            max_face_num = 0
            max_klass_name = ''
            for rv in people['data']:
                words = [rv['klass'], str(rv['number'])]
                fd.write(' '.join(words))
                fd.write('\n')
                
                if rv['number'] > max_face_num:
                    max_face_num = rv['number']
                    max_klass_name = rv['klass']
            # 计算一些数值
            correct_rate = (max_face_num+0.0) / people['face_num']
            failed_num = len(people['data']) - 1
            klass_face_num = klass_list.get_klass(max_klass_name)['face_num']
            similar_rate = 1 - (max_face_num+0.0) / klass_face_num
            words = ["所属类别", str(max_klass_name), '\n']
            fd.write(' '.join(words))
            words = ["识别率", str(correct_rate), '\n']
            fd.write(' '.join(words))
            words = ["误认类数", str(failed_num), '\n']
            fd.write(' '.join(words))
            words = ["与他人近似度", str(similar_rate), '\n', '\n']
            fd.write(' '.join(words))


if __name__ == '__main__':
    
    config = load_config(CONFIG_TXT)
    assert(config['src_dir'] != None)
    assert(config['face_dir'] != None)
    assert(config['result_dir'] != None)
    read_people(config['face_dir'], config['result_dir'])
    get_result(config['result_dir'])