IBM_AI/vi-opton/opton_cls_train/eval.py

import logging
logging.basicConfig(
    format='%(asctime)s [%(levelname)s] - %(message)s',
    level=logging.INFO)

import os

from vi_cls.cls_pred import cls_pred
from config import config


class eval(cls_pred):

    def __init__(self, cfg, args=None):
        cls_pred.__init__(self, cfg, args)

        self.num_total = 0
        self.num_ok = 0
        self.num_ng = 0
        self.label_class = None
        self.pred_class = None

    #------------------------------------------------------------------------------
    # Predict Image Files
    #------------------------------------------------------------------------------

    def pred_file_startup(self, img_dir, names):
        self.num_total = len(names)
        self.num_ok = 0
        self.num_ng = 0

        # the parent folder name is the class_name
        self.label_class = os.path.split(img_dir)[1]
        if self.label_class == '':
            self.label_class = os.path.split(self.label_class)[1]

    def before_pred_file(self, img_file):
        split_num = 2
        # if img_file.find('100X') > 0:
        #     split_num = 2
        # elif img_file.find('200X') > 0:
        #     split_num = 2
        # elif img_file.find('500X') > 0:
        #     split_num = 1
        # elif img_file.find('1000X') > 0:
        #     split_num = 1
        self.set_arg('splits_num', split_num)
        

    def after_pred_file(self, det_dict):
        self.pred_class = ''

        if len(det_dict) > 0:
            dets = sorted(det_dict.items(), lambda x, y : cmp(x[1], y[1]), reverse=True)
            self.pred_class = dets[0][0]

    def get_pred_result_text(self, det_dict, im_inp):
        if self.pred_class is not None:
            pred_items = self.pred_class.rpartition('-')
            label_items = self.label_class.rpartition('-')
            if pred_items[0] == label_items[0]:
                self.num_ok = self.num_ok + 1
                return 'OK'

        self.num_ng = self.num_ng + 1
        return 'NG'

    def get_summary_text_lines(self):
        ok_percent = 0
        ng_percent = 0

        if self.num_total > 0:
            ok_percent = float(self.num_ok) / self.num_total
            ng_percent = float(self.num_ng) / self.num_total

        lines = []
        lines.append('')
        lines.append('')
        lines.append('{:>10}: {:5} ({:.2%})'.format('OK', self.num_ok, ok_percent))
        lines.append('{:>10}: {:5} ({:.2%})'.format('NG', self.num_ng, ng_percent))
        lines.append('')

        return lines

if __name__ == '__main__':
    cfg = config()
    e = eval(cfg)

    split_num = 2
    crop_size = 224

    e.set_arg('gray_image', 'False')
    e.set_arg('resize_w', crop_size * split_num)
    e.set_arg('resize_h', crop_size * split_num)
    e.set_arg('crop_w', crop_size)
    e.set_arg('crop_h', crop_size)

    root_dir = e._args.img_dir
    lines = []
    lines.append('')
    lines.append('')

    names = os.listdir(root_dir)
    for name in names:
        img_dir = os.path.join(root_dir, name)
        e.set_arg('img_dir', img_dir)

        e.pred_file()

        ok_percent = 0
        if e.num_total > 0:
            ok_percent = float(e.num_ok) / e.num_total
        lines.append('{:>20}: {:.2%} ({}/{})'.format(name, ok_percent, e.num_ok, e.num_total))
    
    summary_file = open(os.path.join(e._args.pred_dir, 'summary.txt'), 'w')
    for line in lines:
        summary_file.write(line + '\n')
        logging.info(line)
    summary_file.close()