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
| import matplotlib.pyplot as plt
from PIL import Image
import os
# Can put the results of different models into this function, and it will return the average value of IoU
def display_annotated_results(imgId, valid_boxes, model_name, color='g', ax=None):
# Load the image
imgInfo = coco.loadImgs(imgId)[0]
image_path = os.path.join(cocoRoot, dataType, imgInfo['file_name'])
image = Image.open(image_path)
# Get the correct bounding box results
annIds = coco.getAnnIds(imgIds=imgInfo['id'])
anns = coco.loadAnns(annIds)
bbox_tlist_anns = torch.tensor([ann["bbox"] for ann in anns]) # tensor.shape[2,4]
# because our bounding box is x,y,w,h which is the coordinate of the lower left corner of the box (x,y) + the length and width of the box
# But torchvision Calculate the box_iou must give the coordinates of the lower left corner (x,y) and the coordinates of the upper right corner (x2,y2), so we need to Calculate (x2,y2) through (x+w, y+h) to get the coordinates of the upper right corner
# x,y,w,h -> x1,y1,x2,y2 = x,y,x+w,y+h
bbox_tlist_anns[:, 2] = bbox_tlist_anns[:, 0] + bbox_tlist_anns[:, 2]
bbox_tlist_anns[:, 3] = bbox_tlist_anns[:, 1] + bbox_tlist_anns[:, 3]
# From resultsvalid_boxes, we only need the box part, so we use (box, _, _)
# Use stack because we want to stack all the boxes together to become a tensor
bbox_tlist_model = torch.stack([box for box, _, _ in valid_boxes]) # turn [4] to tensor.shape[2,4]
# use box_iou 來Calculate IoU
iou = torchvision.ops.box_iou(bbox_tlist_anns, bbox_tlist_model) # get IoU
# Get the maximum value of each predicted box in ann, and then Calculate the average value of IoU
avg_iou = np.mean([t.cpu().detach().numpy().max() for t in iou]) # calculate the mean of IoU
# display image label
all_labels = set()
# Start drawing the prediction box
for boxes in valid_boxes:
# Get the information of the prediction box, including box, label, and accuracy
box, label, score = boxes
# Get the text information of the label: load category name by category ID
label = coco.loadCats(label.item())[0]["name"]
# Save the label for later display
all_labels.add(label)
# Because the results returned by the model are two coordinates, the lower left corner and the upper right corner, so we need to convert them into x,y,w,h form and put them into Rectangle
x, y, x2, y2 = box.detach().numpy() # x,y,w,h -> x,y,x2-x,y2-y
rect = Rectangle((x, y), x2 - x, y2 - y, linewidth=2, edgecolor=color, facecolor='none')
# Draw the picture: if you need to sort the picture, you can specify where to draw the picture through ax
if ax is None:
# gca can get the current axes, if not, it will automatically create one, and then draw the prediction box through add_patch
plt.gca().add_patch(rect)
# Draw the label on the prediction box
plt.text(x, y, f'{label}', fontsize=10, color='w', backgroundcolor=color)
else:
# add_patch can add a patch to the current axes, and then draw the prediction box on the ax
ax.add_patch(rect)
# Draw the label on the prediction box
ax.text(x, y, f'{label}', fontsize=10, color='w', backgroundcolor=color)
# display image and give it a title, the title is the label that appears in this picture, and the average value of IoU
if ax is None:
plt.axis('off')
plt.title(f'{model_name}: {all_labels} \n IoU: {avg_iou:.4f}', color=color)
plt.imshow(image)
plt.show()
else:
ax.axis('off')
ax.set_title(f'{model_name}: {all_labels} \n I0U: {avg_iou:.4f}', color=color)
ax.imshow(image)
return avg_iou
res_iou = []
mobile_iou = []
# Recursively call each id, where id is one of the 10 random ids we selected above
for i in range(len(valid_ids)):
# Create a 1x3 grid of images, each image sized 15x5
fig, axs = plt.subplots(1, 3, figsize=(15, 5))
# Draw the truth image and display it in the center
plot_image_with_annotations(coco, cocoRoot, dataType, valid_ids[i], ax=axs[1])
# Draw the predicted results from two different models on the left and right sides respectively, and return the IoU
i_mobil_iou = display_annotated_results(valid_ids[i], valid_boxes_mobile[i], "mobile", color='g', ax=axs[0])
i_res_iou = display_annotated_results(valid_ids[i], valid_boxes_res[i], "ResNet", color='b', ax=axs[2])
# Store the IoU of each image to assess the overall performance of the model
mobile_iou.append(i_mobil_iou)
res_iou.append(i_res_iou)
# Organize the layout
plt.tight_layout()
# Print the mean of the IoU list
print("ResNet: Avg.", np.mean(res_iou), "; each IoU:", res_iou)
print("MobileNet: Avg.", np.mean(mobile_iou), "; each IoU:", mobile_iou)
|
