-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain_evaluate_people_code.py
161 lines (131 loc) · 6.54 KB
/
main_evaluate_people_code.py
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
"""
Project: Object detector and segmentation tutorial https://github.com/juancarlosmiranda/object_detector_tutorial
Author: Juan Carlos Miranda. https://github.com/juancarlosmiranda
Date: February 2021
Description:
Sample code from the TorchVision 0.3 Object Detection Finetuning Tutorial. http://pytorch.org/tutorials/intermediate/torchvision_tutorial.html
Explained by David Macedo.
https://github.com/dlmacedo/starter-academic/blob/master/content/courses/deeplearning/notebooks/pytorch/finetuning_torchvision_models_tutorial.ipynb
Adapted by Juan Carlos Miranda as a programming practice, February 2021.
Use:
"""
import os
import time
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
import torchvision.transforms.functional as F
from helpers.helper_examples import merge_masks
from PIL import Image
def get_model_instance_segmentation(num_classes):
# load an instance segmentation model pre-trained pre-trained on COCO
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# now get the number of input features for the mask classifier
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
print('get_model_instance_segmentation(num_classes):')
print('in_features -->', in_features)
print('in_features_mask -->', in_features_mask)
print('hidden_layer -->', hidden_layer)
print('num_classes -->', num_classes)
return model
def main_evaluate_people_loop():
print('------------------------------------')
print('MAIN STORY OBJECT DETECTION EVALUATION')
print('------------------------------------')
main_path_project = os.path.abspath('.')
# -------------------------------------------
# Datasets
# -------------------------------------------
dataset_folder = os.path.join('dataset', 'people')
path_dataset = os.path.join(main_path_project, dataset_folder)
path_images_folder = 'images'
path_dataset_images = os.path.join(path_dataset, path_images_folder)
# -------------------------------------------
# Open image with Pillow.Image.open() and torchvision.io.read_image()
# -------------------------------------------
img_to_eval_name = 'PATT_01_MIRANDA.png'
img_to_eval_name = 'kayaking_20211120.png'
img_to_eval_name = '20210523_red_cross.png'
device_selected = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
path_img_to_eval = os.path.join(path_dataset_images, img_to_eval_name)
p_img_to_eval = Image.open(path_img_to_eval) # {PngImageFile}
img_to_eval_float32 = F.to_tensor(p_img_to_eval)
img_to_eval_list = [img_to_eval_float32.to(device_selected)]
# -------------------------------------------
# Trained parameters for models
# -------------------------------------------
trained_model_folder = 'trained_model'
trained_model_path = os.path.join(main_path_project, trained_model_folder)
file_name_model = 'model_maskrcnn_20230329_173739.pth'
file_model_path = os.path.join(trained_model_path, file_name_model)
# -------------------------------------------
# ------------------------------------------
# Model initialization for object prediction
# -------------------------------------------
score_threshold = 0.5
num_classes = 2
start_time_model_load = time.time()
model = get_model_instance_segmentation(num_classes)
checkpoint = torch.load(file_model_path)
model.load_state_dict(checkpoint) # equivalent to maskrcnn_resnet50_fpn(), but with file loading
model.to(device_selected)
model.eval() # enabling evaluation mode
end_time_model_load = time.time()
# ----------------------------
# -------------------------------------
# Image evaluation with model
# -------------------------------------
start_time_eval = time.time() # this is the evaluation
# Data type int_input {Tensor:3}, tensor_input {Tensor:1}
with torch.no_grad():
predictions_model = model(img_to_eval_list)
end_time_eval = time.time()
# -------------------------------------
# Managing prediction, making something here (filtering, extracting)
# -------------------------------------
pred_scores = predictions_model[0]['scores'].detach().cpu().numpy()
pred_masks = predictions_model[0]['masks']
# -------------------------------------
# Filtering predictions according to rules
# -------------------------------------
masks_filtered = pred_masks[pred_scores >= score_threshold]
final_masks = masks_filtered > 0.5 # to clean bad pixels
final_masks = final_masks.squeeze(1) # ?
# -------------------------------------
# It displays the results on the screen according to the colours.
# -------------------------------------
merged_masks = merge_masks(final_masks)
merged_binary_img = Image.fromarray(merged_masks.mul(255).byte().cpu().numpy())
merged_binary_img.show('binary mask to show')
# -------------------------------------
# Display data on screen
# -------------------------------------
total_time_model_load = end_time_model_load - start_time_model_load
total_time_eval = end_time_eval - start_time_eval
process_time_eval = total_time_model_load + total_time_eval
w, h = p_img_to_eval.size
print('------------------------------------')
print(f'Main parameters')
print('------------------------------------')
print(f'path_dataset_images={path_dataset_images}')
print(f'path_img_to_evaluate_01={path_img_to_eval}')
print(f'Image size width={w} height={h}')
print(f'device_selected={device_selected}')
print(f'score_threshold={score_threshold}')
print(f'Trained model file_model_path={file_model_path}')
print(f'model={type(model).__name__}')
print(f'total_time_model_load={total_time_model_load}')
print(f'total_time_eval={total_time_eval}')
print(f'process_time_eval={process_time_eval}')
if __name__ == "__main__":
main_evaluate_people_loop()