HierCore/infer.py at main · DSBA-Lab/HierCore · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import os
import json
import numpy as np
from tqdm.auto import tqdm
from omegaconf import OmegaConf

import utils
import backbones
import common
from models import create_model
from dataset import create_dataset, DatasetSplit
from main import get_sampler


def infer_time(savedir):
    cfg = OmegaConf.load(os.path.join(savedir, 'config.yaml'))

    savedir = os.path.join(cfg.DEFAULT.savedir, cfg.DATASET.classname)
    device = utils.set_torch_device(gpu_ids=cfg.DEFAULT.device_ids)

    testset = create_dataset(
        dataname  = cfg.DATASET.name,
        source    = cfg.DATASET.datadir,
        classname = cfg.DATASET.classname,
        resize    = cfg.DATASET.resize,
        imagesize = cfg.DATASET.imagesize,
        split     = DatasetSplit.TEST,
    )

    # create model
    params = {}
    if cfg.MODEL.name == "HierarchicalPatchCore":
        params.update({
            "semantic_layer_to_extract_from" : cfg.MODEL.get("semantic_layer"),
            "semantic_batch_size"            : cfg.MODEL.get("semantic_batch_size"),
        })

    model = create_model(
        modelname                = cfg.MODEL.name,
        device                   = device,
        backbone                 = backbones.load(cfg.MODEL.backbone),
        layers_to_extract_from   = cfg.MODEL.layers,
        input_shape              = testset.imagesize,
        pretrain_embed_dimension = cfg.MODEL.pretrained_embed_dim,
        target_embed_dimension   = cfg.MODEL.target_embed_dim,
        patchsize                = cfg.MODEL.patchsize,
        anomaly_scorer_num_nn    = cfg.MODEL.anomaly_nn,
        featuresampler           = get_sampler(name=cfg.MODEL.sampler, percentage=cfg.MODEL.sampler_ratio, device=device, seed=cfg.DEFAULT.seed, **cfg.MODEL.get('params', {})),
        nn_method                = common.FaissNN(on_gpu=cfg.MODEL.faiss.use_gpu, num_workers=cfg.MODEL.faiss.num_workers),
        **params
    )

    model.load_from_path(
        load_path     = savedir,
        backbone_name = cfg.MODEL.backbone,
        nn_method     = common.FaissNN(on_gpu=True, num_workers=cfg.MODEL.faiss.num_workers)
    )


    class_names = np.array([c[0] for c in testset.data_to_iterate])

    class_idx = {}
    for c in np.unique(class_names):
        class_idx[c] = np.where(class_names==c)[0]


    warmup_iter = 5
    repeat_iter = 100
    infer = {}
    for c, idx in tqdm(class_idx.items(), total=len(class_idx), desc='Class'):
        img = testset[idx[0]]['image'].unsqueeze(0)
        for i in range(warmup_iter):
            _ = model.predict(img)

        repeat = []
        for i in tqdm(range(repeat_iter), desc='repeat', leave=False):
            r = model.predict(img)
            repeat.append(r['infer_time']['FPS'])

        infer[c] = np.mean(repeat)

    infer['all'] = np.mean(list(infer.values()))

    json.dump(infer, open(os.path.join(savedir, 'FPS.json'), 'w'), indent='\t')


if __name__ == "__main__":
    cfg = OmegaConf.from_cli()
    infer_time(cfg.savedir)