MDMLM/eval.py at master · naivete5656/MDMLM · 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
import torch
import torch.nn as nn
from pathlib import Path
import matplotlib.pyplot as plt
from utils import target_peaks_gen, local_maxima, optimum, remove_outside_plot, show_res
import numpy as np


def eval_net(
    net,
    dataset,
    vis=None,
    vis_img=None,
    vis_gt=None,
    gpu=True,
    dist_peak=2,
    peak_thresh=100,
    dist_threshold=10,
    only_loss=False,
):
    criterion = nn.MSELoss()
    net.eval()
    losses = 0

    for iteration, data in enumerate(dataset):
        img = data["image"]
        target = data["gt"]
        if gpu:
            img = img.cuda()
            target = target.cuda()

        pred_img = net(img)

        loss = criterion(pred_img, target)
        losses += loss.data

        # if vis is not None:
        #     vis.images(img.cpu(), 1, win=vis_ori)
        #     vis.images(pred_img.cpu(), 1, win=vis_img)
        #     vis.close()

    return losses / iteration


def eval_net_cnn(
    net,
    dataset,
    vis=None,
    vis_img=None,
    vis_gt=None,
    gpu=True,
    dist_peak=2,
    peak_thresh=100,
    dist_threshold=10,
    only_loss=False,
):
    criterion = nn.BCELoss()
    net.eval()
    losses = 0

    for iteration, data in enumerate(dataset):
        img = data["image"]
        target = data["gt"]
        if gpu:
            img = img.cuda()
            target = target.cuda()

        pred_img = net(img)

        loss = criterion(pred_img, target)
        losses += loss.data

        # if vis is not None:
        #     vis.images(img.cpu(), 1, win=vis_ori)
        #     vis.images(pred_img.cpu(), 1, win=vis_img)
        #     vis.close()

    return losses / iteration