Algorithm-Repo/SourceCode/MachineLearningUtils/SemanticSeg/PytorchFramework/val.py

# D:/workplace/python
# -*- coding: utf-8 -*-
# @ File       : val.py
# @ Author     : Guido    LuXiaohao
# @ Date       : 2021/8/19
# @ Software   : PyCharm
# @ Description: 代码文件描述。

import os
import sys
import zipfile
from argparse import ArgumentParser
from copy import deepcopy
from pathlib import Path
from typing import Sequence

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from dataset import custom_collate_fn
from dataset.data_process import DataProcess
from dataset.utils import parse_metadata_info
from utils.config import Config
from utils.utils import convert_nested_tensors_to_device, select_device

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # project root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative path


def zipDir(dir_path, save_name=""):
    zip_file = dir_path + '.zip' if not save_name.endswith(".zip") else save_name  # 压缩后文件夹的名字
    z = zipfile.ZipFile(zip_file, 'w', zipfile.ZIP_DEFLATED)  # 参数一：文件夹名
    for dirpath, dirnames, filenames in os.walk(dir_path):
        fpath = dirpath.replace(dir_path, '')  # 这一句很重要，不replace的话，就从根目录开始复制
        fpath = fpath and fpath + os.sep or ''  # 这句话理解我也点郁闷，实现当前文件夹以及包含的所有文件的压缩
        for filename in filenames:
            z.write(os.path.join(dirpath, filename), fpath + filename)
            # print ('压缩成功')
    z.close()


def evaluate(cfg, model, eval_loader, scales=[1.0], overlap=0.0, save=False):
    model.eval()
    device = next(model.parameters()).device

    loss = 0
    loss_cnt = 0

    pixel_confidence_threshold = cfg['postprocess_args']['pixel_confidence_threshold']
    limit_area = cfg['postprocess_args']['limit_area']
    if not isinstance(limit_area, list):
        limit_area = [limit_area]
    limit_area_rate = cfg['postprocess_args']['limit_area_rate']
    if not isinstance(limit_area_rate, list):
        limit_area_rate = [limit_area_rate]
    limit_number = cfg['postprocess_args']['limit_number']
    if not isinstance(limit_number, list):
        limit_number = [limit_number]
    needed_rois_dict = cfg['metric_args']['needed_rois_dict']
    if not isinstance(needed_rois_dict, list):
        needed_rois_dict = [needed_rois_dict]
    class_id_map = cfg['metric_args']['class_id_map']
    if not isinstance(class_id_map, list):
        class_id_map = [class_id_map]

    # Metric method
    _metrics = cfg.metrics
    for idx, _metric in enumerate(_metrics):
        dataset_meta = dict(
            pixel_confidence_threshold=pixel_confidence_threshold,
            limit_area=limit_area[idx],
            limit_area_rate=limit_area_rate[idx],
            limit_number=limit_number[idx],
            needed_rois_dict=needed_rois_dict[idx],
            class_id_map=class_id_map[idx],
            class_names=cfg['roi_cat'][idx],
            crop=cfg['val_dataset'].get('crop_class_index', None),
        )
        _metric.init_dataset_meta(dataset_meta)

    pbar = tqdm(iterable=enumerate(eval_loader), total=len(eval_loader), desc='Evaluating on validation dataset')
    for step, (image, data_samples) in pbar:  # ---------------- start evaluating

        with torch.no_grad():
            image = image.to(device)
            data_samples = convert_nested_tensors_to_device(data_samples, device)
            loss_val = model(image, data_samples, mode='loss')
            predictions = model(image, data_samples, mode='predict')

        loss += sum(loss_val.values())
        loss_cnt += 1

        # Add step to data_samples for `Metric`
        for data_sample in data_samples:
            if isinstance(data_sample, Sequence):
                # adapt to multi-output model
                for i_data_sample in data_sample:
                    i_data_sample.set_metainfo(dict(step=step))
            else:
                data_sample.set_metainfo(dict(step=step))

        for idx, _metric in enumerate(_metrics):
            if isinstance(predictions[0], Sequence):
                task_predictions = [i_predictions[idx] for i_predictions in predictions]
                _metric.process(image, task_predictions)
            else:
                _metric.process(image, predictions)

    metrics = dict(loss=loss/loss_cnt)
    for _metric in _metrics:
        ret_metric = _metric.evaluate()
        metrics.update(ret_metric)

    return metrics


def parse_args():
    parser = ArgumentParser(description='Semantic segmentation with pytorch')
    parser.add_argument(
        '--cfg',
        help="training configuration file",
        type=str,
        default=ROOT / 'model.yaml')
    parser.add_argument(
        '--weights',
        help="initial weights path",
        type=str,
        default='')
    parser.add_argument(
        '--batch_size',
        help="total batch size for all GPUs, -1 for autobatch",
        type=int,
        default=8)
    parser.add_argument(
        '--num_workers',
        help=" the number of parallel threads",
        type=int,
        default=1)
    parser.add_argument(
        '--token',
        help="AI平台训练码",
        type=str,
        default="4925EC4929684AA0ABB0173B03CFC8FF")
    parser.add_argument(
        '--val_ratio',
        help="在线随机选取验证集（验证集占总数据集比例），剩下的作为训练集",
        type=float,
        default=0.)
    parser.add_argument(
        '--device',
        help="cuda device, i.e. 0 or 0,1,2,3 or cpu",
        type=str,
        default="0")

    return parser.parse_args()


def run(args):

    weights = args.weights
    device = select_device(args.device, args.batch_size)

    if not args.cfg:
        raise RuntimeError('No configuration file specified.')

    cfg = Config(args.cfg, batch_size=args.batch_size)

    # Dataset properties
    tasks, part_cat, roi_cat, cat_id, label_map = \
        parse_metadata_info(metadata=deepcopy(cfg["metadata"]))
    cfg.update(**{
        'tasks': tasks,
        'part_cat': part_cat,
        'roi_cat': roi_cat,
        'cat_id': cat_id,
        'label_map': label_map
    })

    # Data preprocessing
    data_processor = DataProcess(
        token=args.token,
        val_ratio=args.val_ratio,
        image_cat=cfg["metric_args"]["needed_imageresults_dict"],
        part_cat=part_cat,
        roi_cat=roi_cat,
        cat_id=cat_id,
        ignore=cfg.dic.get("ignore", None),
        video_mode=cfg.dic.get("split_by_snippet", False))

    trainval_indexes = data_processor.get_train_val_indexes
    train_data_index_list = trainval_indexes["train"]
    val_data_index_list = trainval_indexes["val"]

    # update dataset info
    cfg['train_dataset'].update({
        'token': args.token,
        'tasks': tasks,
        'data_index_list': train_data_index_list,
        'class_index_map': cat_id})
    cfg['val_dataset'].update({
        'token': args.token,
        'tasks': tasks,
        'data_index_list': val_data_index_list,
        'class_index_map': cat_id})

    val_dataset = cfg.val_dataset
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        collate_fn=custom_collate_fn,
        pin_memory=True,
        drop_last=True if args.batch_size > 1 else False)

    # Model
    model = cfg.model
    model = model.to(device)
    pretrained = weights.endswith('.pth')
    if pretrained:
        ckpt = torch.load(weights, map_location='cpu')
        csd = ckpt.get("ema") or ckpt['model']
        model.load_state_dict(csd, strict=True)  # load

    # Start evaluation
    metrics = evaluate(cfg, model, val_loader)
    # Calculate the mean value over tasks
    miou = np.mean([v for m, v in metrics.items() if m.endswith('.mIoU')])
    all_precision = np.mean([v for m, v in metrics.items() if m.endswith('.aPrecision')])
    all_recall = np.mean([v for m, v in metrics.items() if m.endswith('.aRecall')])
    map50 = np.mean([v for m, v in metrics.items() if m.endswith('.mAP50')])
    map = np.mean([v for m, v in metrics.items() if m.endswith('.mAP')])

    print('Validating {}\n'
          'miou: {:.4f}\n'
          'total precision: {:.4f}, total recall: {:.4f}, map50: {:.4f}, map50-95: {:.4f}\n'
          .format(args.weights, miou, all_precision, all_recall, map50, map))
 

if __name__ == '__main__':
    args = parse_args()
    run(args)