semantic_segmentation_model/evaluate.py at main · 1900690105/semantic_segmentation_model · 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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import torch
from torch.utils.data import DataLoader

from config import *
from utils.dataset import SegmentationDataset
from models.unet import UNet
from utils.metrics import calculate_iou, pixel_accuracy, approximate_map
from utils.metrics import map50_score

# -------------------------------
# Device
# -------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# -------------------------------
# Load Dataset
# -------------------------------
val_dataset = SegmentationDataset(VAL_IMG_DIR, VAL_MASK_DIR)
val_loader = DataLoader(val_dataset, batch_size=1)

print("Total validation samples:", len(val_dataset))

# -------------------------------
# Load Model
# -------------------------------
model = UNet(NUM_CLASSES).to(device)
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()

# -------------------------------
# Evaluation
# -------------------------------
ious = []
accs = []
maps = []
map50_scores = []

with torch.no_grad():
    for images, masks in val_loader:
        images, masks = images.to(device), masks.to(device)

        outputs = model(images)

        score = map50_score(outputs, masks)
        map50_scores.append(score)

        # IoU
        iou = calculate_iou(outputs, masks, NUM_CLASSES)
        ious.append(iou)

        # Pixel Accuracy
        acc = pixel_accuracy(outputs, masks)
        accs.append(acc)

        # Approx mAP
        m = approximate_map(outputs, masks)
        maps.append(m)

# -------------------------------
# Results
# -------------------------------
if len(ious) > 0:
    print("Mean IoU:", sum(ious) / len(ious))
else:
    print("IoU not calculated")

if len(accs) > 0:
    print("Pixel Accuracy:", sum(accs) / len(accs))
else:
    print("Accuracy not calculated")

if len(maps) > 0:
    print("Approx mAP:", sum(maps) / len(maps))
else:
    print("mAP not calculated")

if len(map50_scores) > 0:
    print("mAP@50:", sum(map50_scores) / len(map50_scores))

# import torch
# from torch.utils.data import DataLoader
# from config import *
# from utils.dataset import SegmentationDataset
# from models.unet import UNet
# from utils.metrics import calculate_iou

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# val_dataset = SegmentationDataset(VAL_IMG_DIR, VAL_MASK_DIR)
# val_loader = DataLoader(val_dataset, batch_size=1)

# model = UNet(NUM_CLASSES).to(device)
# model.load_state_dict(torch.load("model.pth"))
# model.eval()

# ious = []

# with torch.no_grad():
#     for images, masks in val_loader:
#         images, masks = images.to(device), masks.to(device)
#         outputs = model(images)

#         iou = calculate_iou(outputs, masks, NUM_CLASSES)
#         ious.append(iou)

# print("Mean IoU:", sum(ious)/len(ious))