CSRnet/Analysis.py at master · renzodgc/CSRnet · 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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
# -*- coding: utf-8 -*-
"""
Created on Sat Nov  3 14:50:14 2018

@author: lenovo
"""

# Python Imports

import cv2
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import h5py
import glob
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from pathlib import Path
from PIL import Image
from sklearn.metrics import mean_absolute_error
from matplotlib import cm as c
from keras.models import model_from_json

import warnings
warnings.filterwarnings("ignore")

# model_A_weights (Trained with ShanghaiTech A: Shanghai crowd pictures from the Internet)
MODELS_PATHS = [("A", "weights/model_A_weights.h5"), ("B", "weights/model_B_weights.h5")]
DATASET_PATH = "data"
MODEL_JSON = "models/Model.json"

RESULTS_PATH = "results"
VISUALIZATION_PATH = f"{RESULTS_PATH}/visualization"

def load_model(model):
    # Function to load and return neural network model

    json_file = open(MODEL_JSON, 'r')
    loaded_model_json = json_file.read()
    json_file.close()
    loaded_model = model_from_json(loaded_model_json)
    loaded_model.load_weights(model)
    return loaded_model

def create_img(path):
    # Function to load,normalize and return image
    im = Image.open(path).convert('RGB')

    im = np.array(im)

    im = im/255.0

    im[:,:,0]=(im[:,:,0]-0.485)/0.229
    im[:,:,1]=(im[:,:,1]-0.456)/0.224
    im[:,:,2]=(im[:,:,2]-0.406)/0.225

    im = np.expand_dims(im, axis  = 0)
    return im

def predict(path, model):
    # Function to load image,predict heat map, generate count and return (image , heat map, count)
    image = create_img(path)
    pred_hmap = model.predict(image)
    pred_count = np.sum(pred_hmap)
    return image, pred_hmap, pred_count

def run_test(test_image_path, test_gt_path, model, model_index):
    image, pred_hmap, pred_count = predict(test_image_path, model)
    gt_file = h5py.File(test_gt_path , 'r')
    gt_hmap = np.asarray(gt_file['density'])
    gt_count = np.sum(gt_hmap)
    process_test_results(image, pred_hmap, pred_count, gt_hmap, gt_count, test_image_path, model_index)
    return image, pred_hmap, pred_count, gt_hmap, gt_count

def process_test_results(img, pred_hmap, pred_count, gt_hmap, gt_count, image_path, model_index):
    print(f"Predict Count: {pred_count} out of {gt_count} annotated")
    image_name, image_ext = os.path.splitext(os.path.basename(image_path))
    original_image = cv2.imread(image_path)
    pred_hmap = pred_hmap.reshape(pred_hmap.shape[1], pred_hmap.shape[2])
    pred_image = cv2.resize(pred_hmap, dsize=(img.shape[2], img.shape[1]), interpolation=cv2.INTER_CUBIC)

    cv2.imwrite(os.path.join(VISUALIZATION_PATH, model_index, f"{image_name}-img{image_ext}"), original_image)
    fig = plt.imshow(pred_image, cmap = c.jet)
    plt.axis('off')
    fig.axes.get_xaxis().set_visible(False)
    fig.axes.get_yaxis().set_visible(False)
    plt.savefig(os.path.join(VISUALIZATION_PATH, model_index, f"{image_name}-img-gt{image_ext}"), bbox_inches='tight', pad_inches = 0)
    plt.close()
    fig = plt.imshow(gt_hmap, cmap = c.jet)
    plt.axis('off')
    fig.axes.get_xaxis().set_visible(False)
    fig.axes.get_yaxis().set_visible(False)
    plt.savefig(os.path.join(VISUALIZATION_PATH, model_index, f"{image_name}-img-pred{image_ext}"), bbox_inches='tight', pad_inches = 0)
    plt.close()

def test_dataset(index, model):
    print()
    print(f"--- Starting model {index} ---")
    print()
    img_paths = glob.glob(os.path.join(DATASET_PATH, "images", '*.jpg'))

    name = []
    y_true = []
    y_pred = []

    Path(os.path.join(VISUALIZATION_PATH, index)).mkdir(parents=True, exist_ok=True)

    for image_path in img_paths:
        name.append(image_path)
        gt_path = image_path.replace('.jpg','.h5').replace('images','ground_truth')
        img, pred_hmap, pred_count, gt_hmap, gt_count = run_test(image_path, gt_path, model, index)
        y_pred.append(pred_count)
        y_true.append(gt_count)

    data = pd.DataFrame({'name': name, 'y_pred': y_pred, 'y_true': y_true})
    data.to_csv(f'{RESULTS_PATH}/test_{index}.csv', sep=',')
    ans = mean_absolute_error(np.array(y_true), np.array(y_pred))
    print(f"--- Model {index} finished. Total MAE : {ans} ---")
    print()
    print()

if __name__ == "__main__":
    for index, model_path in MODELS_PATHS:
        model = load_model(model_path)
        test_dataset(index, model)