用于目标检测的CNN

1import torch
2import torchvision
3
4# 加载预训练的 FasterRCNN 神经网络
5weights = torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT
6fasterrcnn = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(weights=weights,
7    progress=False)
8fasterrcnn.eval()
9transform = weights.transforms()
10
11# 类别信息位于权重对象的 meta 字典中
12categories = weights.meta["categories"]
13print(f"categories = {categories}")

1import torch
2import torchvision
3import urllib.request
4from PIL import Image
5import numpy as np
6import cv2
7
8# 加载预训练的 FasterRCNN 神经网络
9weights = torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_320_FPN_Weights.DEFAULT
10fasterrcnn = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(weights=weights,
11    progress=False)
12fasterrcnn.eval()
13transform = weights.transforms()
14categories = weights.meta["categories"]
15
16# 街景图像路径
17image_filepaths = ["./images/streets/tramway.jpg",
18                   "./images/streets/flinders_station.jpg",
19                   "./images/streets/cars.jpg"]
20annotated_imgs = []
21batch_list = []
22
23for image_filepath in image_filepaths:
24    image_pil = Image.open(image_filepath)
25    image_arr = np.array(image_pil)
26    image_arr = cv2.cvtColor(image_arr, cv2.COLOR_RGB2BGR)
27    annotated_imgs.append(image_arr)
28    batch_list.append(transform(image_pil))
29
30# Faster-RCNN 接受张量列表作为输入
31batch_results = fasterrcnn(batch_list)
32print(f"batch_results: \n{batch_results}")

1# ... 续上
2def random_color(index):
3    color = ((index * 371)%256, (index * 1169)%256, (index * 947)%256)
4    if color == (0, 0, 0):
5        color = (200, 0, 0)
6    return color
7
8for img_ndx in range(len(batch_results)):
9    labels = [categories[i] for i in batch_results[img_ndx]["labels"]]
10    number_of_boxes = len(labels)
11    annotated_img = annotated_imgs[img_ndx]
12    for box_ndx in range(number_of_boxes):
13        roi_x1y1_x2y2 = batch_results[img_ndx]["boxes"][box_ndx].int().numpy()
14        p1 = (roi_x1y1_x2y2[0], roi_x1y1_x2y2[1])
15        p2 = (roi_x1y1_x2y2[2], roi_x1y1_x2y2[3])
16        color = random_color(box_ndx)
17        cv2.rectangle(annotated_img, p1, p2, color=color, thickness=2)
18        cv2.putText(annotated_img, labels[box_ndx], p1, cv2.FONT_HERSHEY_SIMPLEX,
19                    0.5, color=(0, 0, 0), thickness=3)
20        cv2.putText(annotated_img, labels[box_ndx], p1, cv2.FONT_HERSHEY_SIMPLEX,
21                    0.5, color=color, thickness=1)
22        cv2.imwrite(f"./output/annotated_{img_ndx}.png", annotated_img)

1import torchvision
2import urllib.request
3from PIL import Image
4import os
5import numpy as np
6import cv2
7
8def random_color(index):
9    color = ((index * 371)%256, (index * 1169)%256, (index * 947)%256)
10    if color == (0, 0, 0):
11        color = (200, 0, 0)
12    return color
13
14# 加载带有 MobileNetV3 主干网络的预训练 SSDLite320 神经网络
15weights = torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights.DEFAULT
16ssdlite = torchvision.models.detection.ssdlite320_mobilenet_v3_large(
17    weights=weights, progress=False)
18ssdlite.eval()
19
20transform = weights.transforms()
21categories = weights.meta["categories"]
22
23image_filepaths = ["images/flickrImage.jpeg"]
24
25annotated_imgs = []
26batch_list = []
27
28for img_ndx in range(len(image_filepaths)):
29    with Image.open(image_filepaths[img_ndx]) as image_pil:
30        
31        image_arr = np.array(image_pil)
32        image_arr = cv2.cvtColor(image_arr, cv2.COLOR_RGB2BGR)
33        annotated_imgs.append(image_arr)
34        batch_list.append(transform(image_pil))
35
36batch_results = ssdlite(batch_list)
37print(f"batch_results: \n{batch_results}")
38
39for img_ndx in range(len(batch_results)):
40    labels = [categories[i] for i in batch_results[img_ndx]["labels"]]
41    number_of_boxes = len(labels)
42    annotated_img = annotated_imgs[img_ndx]
43    for box_ndx in range(number_of_boxes):
44      roi_x1y1_x2y2 = batch_results[img_ndx]["boxes"][box_ndx].int().numpy()
45      p1 = (roi_x1y1_x2y2[0], roi_x1y1_x2y2[1])
46      p2 = (roi_x1y1_x2y2[2], roi_x1y1_x2y2[3])
47      color = random_color(box_ndx)
48      cv2.rectangle(annotated_img, p1, p2, color=color, thickness=2)
49      cv2.putText(annotated_img, labels[box_ndx], p1, cv2.FONT_HERSHEY_SIMPLEX,
50                  0.5, color=(0, 0, 0), thickness=3)
51      cv2.putText(annotated_img, labels[box_ndx], p1, cv2.FONT_HERSHEY_SIMPLEX,
52                  0.5, color=color, thickness=1)
53      cv2.imwrite("./output/annotated_{img_ndx}.png", annotated_img)

1# ... 续上
2
3batch_results = ssdlite(batch_list)
4print(f"batch_results: \n{batch_results}")
5
6minimum_score = 0.9
7for img_ndx in range(len(batch_results)):
8    labels = [categories[i] for i in batch_results[img_ndx]["labels"]]
9    number_of_boxes = len(labels)
10    annotated_img = annotated_imgs[img_ndx]
11    for box_ndx in range(number_of_boxes):
12        # 读取置信度分数
13        score = batch_results[img_ndx]["scores"][box_ndx]
14        if score >= minimum_score:
15            roi_x1y1_x2y2 = batch_results[img_ndx]["boxes"][box_ndx].int().numpy()
16            p1 = (roi_x1y1_x2y2[0], roi_x1y1_x2y2[1])
17            p2 = (roi_x1y1_x2y2[2], roi_x1y1_x2y2[3])
18            color = random_color(box_ndx)
19            cv2.rectangle(annotated_img, p1, p2, color=color, thickness=2)
20            cv2.putText(annotated_img, labels[box_ndx], p1, cv2.FONT_HERSHEY_SIMPLEX,
21                        0.5, color=(0, 0, 0), thickness=3)
22            cv2.putText(annotated_img, labels[box_ndx], p1, cv2.FONT_HERSHEY_SIMPLEX,
23                        0.5, color=color, thickness=1)
24            cv2.imwrite("./output/annotated_{img_ndx}.png", annotated_img)