# pytorh相关的工具包
import torch
import torchvision
import torchvision.transforms as transforms
# from torchvision.io import read_image
# 读取图像的工具包
from PIL import Image
# 绘图
import matplotlib.pyplot as plt
数据加载
img = Image.open('img2.jpg')
(994, 734)
<matplotlib.image.AxesImage at 0x180689d6688><Figure size 432x288 with 1 Axes>
transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((800,800))])
torch.Size([3, 800, 800])
plt.imshow(img.permute(1,2,0))
<matplotlib.image.AxesImage at 0x7f97b5ce4278><Figure size 432x288 with 1 Axes>
模型加载
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" to C:\Users\14767/.cache\torch\hub\checkpoints\fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
HBox(children=(FloatProgress(value=0.0, max=167502836.0), HTML(value='')))
FasterRCNN(
(transform): GeneralizedRCNNTransform(
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
Resize(min_size=(800,), max_size=1333, mode='bilinear')
)
(backbone): BackboneWithFPN(
(body): IntermediateLayerGetter(
(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=0.0)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): Bottleneck(
(conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=0.0)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64, eps=0.0)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256, eps=0.0)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): FrozenBatchNorm2d(256, eps=0.0)
)
)
(1): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=0.0)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64, eps=0.0)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256, eps=0.0)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=0.0)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64, eps=0.0)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256, eps=0.0)
(relu): ReLU(inplace=True)
)
)
(layer2): Sequential(
(0): Bottleneck(
(conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=0.0)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=0.0)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=0.0)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(512, eps=0.0)
)
)
(1): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=0.0)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=0.0)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=0.0)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=0.0)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=0.0)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=0.0)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=0.0)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=0.0)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=0.0)
(relu): ReLU(inplace=True)
)
)
(layer3): Sequential(
(0): Bottleneck(
(conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=0.0)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=0.0)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=0.0)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(1024, eps=0.0)
)
)
(1): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=0.0)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=0.0)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=0.0)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=0.0)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=0.0)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=0.0)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=0.0)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=0.0)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=0.0)
(relu): ReLU(inplace=True)
)
(4): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=0.0)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=0.0)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=0.0)
(relu): ReLU(inplace=True)
)
(5): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=0.0)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=0.0)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=0.0)
(relu): ReLU(inplace=True)
)
)
(layer4): Sequential(
(0): Bottleneck(
(conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512, eps=0.0)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512, eps=0.0)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048, eps=0.0)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(2048, eps=0.0)
)
)
(1): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512, eps=0.0)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512, eps=0.0)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048, eps=0.0)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512, eps=0.0)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512, eps=0.0)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048, eps=0.0)
(relu): ReLU(inplace=True)
)
)
)
(fpn): FeaturePyramidNetwork(
(inner_blocks): ModuleList(
(0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
(1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
(2): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
(3): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
)
(layer_blocks): ModuleList(
(0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)
(extra_blocks): LastLevelMaxPool()
)
)
(rpn): RegionProposalNetwork(
(anchor_generator): AnchorGenerator()
(head): RPNHead(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
(bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))
)
)
(roi_heads): RoIHeads(
(box_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(7, 7), sampling_ratio=2)
(box_head): TwoMLPHead(
(fc6): Linear(in_features=12544, out_features=1024, bias=True)
(fc7): Linear(in_features=1024, out_features=1024, bias=True)
)
(box_predictor): FastRCNNPredictor(
(cls_score): Linear(in_features=1024, out_features=91, bias=True)
(bbox_pred): Linear(in_features=1024, out_features=364, bias=True)
)
)
)
模型预测
# coco数据集的class，共90个类别：人，自行车，火车，。。。
coco_names = [
'__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
/Users/mac/opt/anaconda3/envs/dltorch/lib/python3.7/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:2157.)
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
[{'boxes': tensor([[ 31.7399, 328.4528, 460.8738, 793.4400],
[ 5.6514, 88.6435, 216.9703, 635.1252],
[367.4497, 225.4498, 762.0249, 789.4284],
[607.8991, 152.3864, 795.5750, 372.0394],
[183.5409, 34.5758, 396.0948, 661.7242],
[364.9885, 43.9209, 525.7601, 586.4974],
[ 13.4759, 211.1637, 290.4755, 748.9048],
[307.3405, 4.9721, 439.2374, 111.8531],
[383.3587, 180.4190, 578.4739, 772.2529],
[407.8741, 139.6278, 681.9370, 533.0029],
[711.5732, 5.5169, 758.7169, 120.2260],
[479.1790, 144.8082, 689.6121, 400.8835],
[244.1306, 36.7394, 487.5280, 744.1151],
[546.8978, 474.4509, 568.8634, 552.9698],
[742.9993, 314.1062, 798.9455, 607.5062],
[742.6909, 317.4082, 796.8760, 619.5782],
[524.8398, 475.1719, 576.5832, 769.7741],
[509.4732, 142.5110, 623.7419, 301.2690],
[469.8165, 142.4818, 800.0000, 396.9042],
[178.5640, 253.6984, 659.5800, 785.6712],
[365.3784, 474.1786, 470.2338, 769.1154],
[160.4500, 265.7687, 387.0724, 610.2078],
[481.0141, 153.0264, 720.3257, 699.9622],
[724.8939, 343.9777, 798.2111, 726.9193],
[439.7558, 568.3492, 517.4477, 787.6546],
[730.4789, 277.5797, 798.6329, 788.0072],
[409.8385, 474.7333, 540.6846, 793.6932],
[720.9894, 114.6369, 753.7390, 180.1143],
[690.4725, 317.8017, 798.0597, 600.9917],
[548.9393, 76.3451, 613.9292, 172.6772],
[720.4436, 118.2152, 754.2703, 179.3526],
[556.6137, 479.1154, 576.5483, 542.2100],
[514.3457, 148.3877, 619.4744, 501.8499],
[197.1474, 368.7211, 431.1885, 764.1635],
[333.9739, 126.6342, 411.8326, 199.1039],
[659.6864, 131.1151, 716.2268, 168.4116],
[318.3166, 531.7626, 407.8087, 717.9441],
[348.7296, 1.2086, 435.9537, 75.7988]], grad_fn=<StackBackward0>),
'labels': tensor([ 1, 1, 1, 72, 1, 1, 1, 72, 1, 1, 86, 1, 1, 32, 63, 62, 32, 1,
72, 1, 32, 1, 1, 63, 32, 62, 32, 44, 63, 72, 86, 32, 1, 1, 79, 85,
32, 72]),
'scores': tensor([0.9966, 0.9899, 0.9894, 0.9877, 0.9701, 0.9598, 0.8676, 0.7946, 0.6888,
0.6239, 0.5967, 0.5822, 0.5759, 0.5339, 0.5014, 0.4489, 0.3447, 0.3405,
0.2513, 0.2067, 0.1989, 0.1979, 0.1844, 0.1780, 0.1268, 0.0884, 0.0864,
0.0813, 0.0766, 0.0748, 0.0731, 0.0704, 0.0678, 0.0637, 0.0592, 0.0554,
0.0535, 0.0532], grad_fn=<IndexBackward0>)}]
list(pred[0]['labels'].numpy())
[1,
1,
1,
72,
1,
1,
1,
72,
1,
1,
86,
1,
1,
32,
63,
62,
32,
1,
72,
1,
32,
1,
1,
63,
32,
62,
32,
44,
63,
72,
86,
32,
1,
1,
79,
85,
32,
72]
pred_class = [coco_names[i] for i in list(pred[0]['labels'].numpy())]
['person',
'person',
'person',
'tv',
'person',
'person',
'person',
'tv',
'person',
'person',
'vase',
'person',
'person',
'tie',
'couch',
'chair',
'tie',
'person',
'tv',
'person',
'tie',
'person',
'person',
'couch',
'tie',
'chair',
'tie',
'bottle',
'couch',
'tv',
'vase',
'tie',
'person',
'person',
'oven',
'clock',
'tie',
'tv']
pred[0]['boxes'].detach().numpy()
array([[ 31.739914 , 328.45282 , 460.87378 , 793.44 ],
[ 5.65139 , 88.64346 , 216.97028 , 635.12524 ],
[367.4497 , 225.44983 , 762.0249 , 789.4284 ],
[607.8991 , 152.38641 , 795.575 , 372.03937 ],
[183.54091 , 34.575836 , 396.09485 , 661.72424 ],
[364.9885 , 43.92093 , 525.76013 , 586.49744 ],
[ 13.475891 , 211.1637 , 290.47552 , 748.90485 ],
[307.3405 , 4.9721336, 439.23743 , 111.85315 ],
[383.3587 , 180.41904 , 578.47394 , 772.2529 ],
[407.8741 , 139.62776 , 681.937 , 533.0029 ],
[711.57324 , 5.516918 , 758.7169 , 120.22604 ],
[479.17905 , 144.80818 , 689.61206 , 400.88348 ],
[244.13058 , 36.73944 , 487.52798 , 744.1151 ],
[546.8978 , 474.45093 , 568.8634 , 552.96985 ],
[742.9993 , 314.10623 , 798.9455 , 607.5062 ],
[742.6909 , 317.40823 , 796.876 , 619.57825 ],
[524.8398 , 475.17194 , 576.5832 , 769.7741 ],
[509.47318 , 142.51099 , 623.74194 , 301.26895 ],
[469.81653 , 142.48177 , 800. , 396.90424 ],
[178.56403 , 253.69843 , 659.58 , 785.6712 ],
[365.37836 , 474.17865 , 470.23376 , 769.1154 ],
[160.45001 , 265.76874 , 387.07245 , 610.2078 ],
[481.0141 , 153.0264 , 720.32574 , 699.96216 ],
[724.8939 , 343.97772 , 798.21106 , 726.91925 ],
[439.75577 , 568.3492 , 517.4477 , 787.6546 ],
[730.47894 , 277.57974 , 798.6329 , 788.0072 ],
[409.8385 , 474.73334 , 540.68463 , 793.6932 ],
[720.98944 , 114.63695 , 753.73895 , 180.11427 ],
[690.47253 , 317.80173 , 798.0597 , 600.9917 ],
[548.93933 , 76.3451 , 613.9292 , 172.67722 ],
[720.4436 , 118.215164 , 754.27026 , 179.35258 ],
[556.6137 , 479.11536 , 576.5483 , 542.21 ],
[514.3457 , 148.38771 , 619.47437 , 501.8499 ],
[197.14737 , 368.72107 , 431.1885 , 764.16345 ],
[333.97394 , 126.63425 , 411.83258 , 199.10388 ],
[659.6864 , 131.11514 , 716.2268 , 168.41156 ],
[318.3166 , 531.7626 , 407.80872 , 717.9441 ],
[348.72955 , 1.2085876, 435.95374 , 75.7988 ]],
dtype=float32)
pred_box = [[(box[0],box[1]),(box[2],box[3])] for box in list(pred[0]['boxes'].detach().numpy())]
[[(31.739914, 328.45282), (460.87378, 793.44)],
[(5.65139, 88.64346), (216.97028, 635.12524)],
[(367.4497, 225.44983), (762.0249, 789.4284)],
[(607.8991, 152.38641), (795.575, 372.03937)],
[(183.54091, 34.575836), (396.09485, 661.72424)],
[(364.9885, 43.92093), (525.76013, 586.49744)],
[(13.475891, 211.1637), (290.47552, 748.90485)],
[(307.3405, 4.9721336), (439.23743, 111.85315)],
[(383.3587, 180.41904), (578.47394, 772.2529)],
[(407.8741, 139.62776), (681.937, 533.0029)],
[(711.57324, 5.516918), (758.7169, 120.22604)],
[(479.17905, 144.80818), (689.61206, 400.88348)],
[(244.13058, 36.73944), (487.52798, 744.1151)],
[(546.8978, 474.45093), (568.8634, 552.96985)],
[(742.9993, 314.10623), (798.9455, 607.5062)],
[(742.6909, 317.40823), (796.876, 619.57825)],
[(524.8398, 475.17194), (576.5832, 769.7741)],
[(509.47318, 142.51099), (623.74194, 301.26895)],
[(469.81653, 142.48177), (800.0, 396.90424)],
[(178.56403, 253.69843), (659.58, 785.6712)],
[(365.37836, 474.17865), (470.23376, 769.1154)],
[(160.45001, 265.76874), (387.07245, 610.2078)],
[(481.0141, 153.0264), (720.32574, 699.96216)],
[(724.8939, 343.97772), (798.21106, 726.91925)],
[(439.75577, 568.3492), (517.4477, 787.6546)],
[(730.47894, 277.57974), (798.6329, 788.0072)],
[(409.8385, 474.73334), (540.68463, 793.6932)],
[(720.98944, 114.63695), (753.73895, 180.11427)],
[(690.47253, 317.80173), (798.0597, 600.9917)],
[(548.93933, 76.3451), (613.9292, 172.67722)],
[(720.4436, 118.215164), (754.27026, 179.35258)],
[(556.6137, 479.11536), (576.5483, 542.21)],
[(514.3457, 148.38771), (619.47437, 501.8499)],
[(197.14737, 368.72107), (431.1885, 764.16345)],
[(333.97394, 126.63425), (411.83258, 199.10388)],
[(659.6864, 131.11514), (716.2268, 168.41156)],
[(318.3166, 531.7626), (407.80872, 717.9441)],
[(348.72955, 1.2085876), (435.95374, 75.7988)]]
# 图像读取
img = Image.open('img2.jpg').resize((800,800))
# 图像展示
fig = plt.imshow(img)
# 将框绘制在图像上
for index, boxes in enumerate(pred_box):
fig.axes.add_patch(plt.Rectangle(
xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
fill=False, edgecolor="blue", linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
pred_score = list(pred[0]['scores'].detach().numpy())
[0.99662447,
0.98991066,
0.9894002,
0.9877265,
0.9700843,
0.95981795,
0.86762965,
0.79457456,
0.6888022,
0.6238965,
0.5967382,
0.5821962,
0.5758556,
0.53390706,
0.50142986,
0.4489189,
0.3446656,
0.34053934,
0.25132728,
0.20674044,
0.19888856,
0.19792202,
0.18443534,
0.17802832,
0.12675783,
0.08840103,
0.08642756,
0.08125935,
0.07662633,
0.0747837,
0.07308833,
0.07040265,
0.06783704,
0.063683495,
0.0592302,
0.055396155,
0.053475436,
0.053160064]
pred_t = [pred_score.index(score) for score in pred_score if score > the]
pred_box = pred_box[:index+1]
[[(31.739914, 328.45282), (460.87378, 793.44)],
[(5.65139, 88.64346), (216.97028, 635.12524)],
[(367.4497, 225.44983), (762.0249, 789.4284)],
[(607.8991, 152.38641), (795.575, 372.03937)],
[(183.54091, 34.575836), (396.09485, 661.72424)],
[(364.9885, 43.92093), (525.76013, 586.49744)],
[(13.475891, 211.1637), (290.47552, 748.90485)],
[(307.3405, 4.9721336), (439.23743, 111.85315)]]
# 图像读取
img = Image.open('img2.jpg').resize((800,800))
# 图像展示
fig = plt.imshow(img)
# 将框绘制在图像上
for index, boxes in enumerate(pred_box):
fig.axes.add_patch(plt.Rectangle(
xy=(boxes[0]), width=boxes[1][0]-boxes[0][0], height=boxes[1][1]-boxes[0][1],
fill=False, edgecolor="blue", linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
['person', 'person', 'person', 'tv', 'person', 'person', 'person', 'tv']
模型结构详解
backbone
transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((800,800))])
img = transform(img)
img.shape
torch.Size([3, 800, 800])
imgs,targets= model.transform([img],targets=None)
<torchvision.models.detection.image_list.ImageList at 0x7f97769e3908>
feature_body= model.backbone.body(imgs.tensors)
C2,C3,C4,C5 = feature_body.values()
torch.Size([1, 256, 200, 200])
torch.Size([1, 512, 100, 100])
torch.Size([1, 1024, 50, 50])
torch.Size([1, 2048, 25, 25])
feature_fpn=model.backbone.fpn(feature_body)
P2,P3,P4,P5,P6=feature_fpn.values()
torch.Size([1, 256, 200, 200])
torch.Size([1, 256, 100, 100])
torch.Size([1, 256, 50, 50])
torch.Size([1, 256, 25, 25])
torch.Size([1, 256, 13, 13])
RPN
anchors = model.rpn.anchor_generator(imgs,[P2,P3,P4,P5,P6])
[tensor([[ -23., -11., 23., 11.],
[ -16., -16., 16., 16.],
[ -11., -23., 11., 23.],
...,
[ 370., 551., 1094., 913.],
[ 476., 476., 988., 988.],
[ 551., 370., 913., 1094.]])]
torch.Size([159882, 4])
# 将生成的anchor绘制在图像上
img_anchor = Image.open('img2.jpg').resize((800,800))
# 绘制图像
fig = plt.imshow(img_anchor)
# 将框绘制在图像上
for index, boxes in enumerate(anchors[0]):
if index <300:
fig.axes.add_patch(plt.Rectangle(
xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
fill=False, edgecolor="red", linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
cls_logits = model.rpn.head.cls_logits(P3)
torch.Size([1, 3, 100, 100])
bbox_pred = model.rpn.head.bbox_pred(P4)
torch.Size([1, 12, 50, 50])
proposal = model.rpn(imgs,feature_fpn)
([tensor([[162.3166, 82.6454, 390.0235, 638.2616],
[ 0.0000, 45.9522, 237.6339, 632.3745],
[170.6081, 36.5363, 371.3253, 548.1898],
...,
[373.4832, 426.8497, 396.0713, 497.3481],
[317.5831, 7.1501, 417.4209, 70.2656],
[395.1561, 506.0393, 427.8982, 546.8923]])],
{})
torch.Size([1000, 4])
# 候选区域的绘制
img = Image.open('img2.jpg').resize((800,800))
# 显示图像
fig = plt.imshow(img)
# 将候选区域显示在图像上
for index, boxes in enumerate(proposal[0][0]):
if index <3:
fig.axes.add_patch(plt.Rectangle(
xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
fill=False, edgecolor="blue", linewidth=0.5))
plt.show()
<Figure size 432x288 with 1 Axes>
pool_region_list= model.roi_heads.box_roi_pool(feature_fpn,proposal[0],imgs.image_sizes)
torch.Size([1000, 256, 7, 7])
detection = model.roi_heads(feature_fpn,proposal[0],imgs.image_sizes)
detection_boxes = detection[0][0]['boxes'].detach().numpy()
img = Image.open('img2.jpg').resize((800,800))
# 显示图像
fig = plt.imshow(img)
# 将检测结果绘制在图像上
for index, boxes in enumerate(detection_boxes):
fig.axes.add_patch(plt.Rectangle(
xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
fill=False, edgecolor="purple", linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
模型训练
数据加载
from voc_data_util import VocDataset,VOC_BBOX_LABEL_NAMES
dataset = VocDataset(
root='datasets/VOCdevkit/VOC2007/',
transforms=VocDataset.get_transform())
torch.Size([3, 800, 800])
import numpy as np
img = img.permute(1,2,0).numpy().astype(np.uint8)
<matplotlib.image.AxesImage at 0x7f9753d7b518><Figure size 432x288 with 1 Axes>
{'boxes': tensor([[238.4000, 233.0330, 550.4000, 648.6486]], dtype=torch.float64),
'labels': tensor([7]),
'image_id': tensor([3]),
'area': tensor([33735.]),
'iscrowd': tensor([0])}
fig = plt.imshow(img)
# 将标注信息绘制在图像上
for index, boxes in enumerate(targets['boxes'].numpy()):
fig.axes.add_patch(
plt.Rectangle(xy=(boxes[0], boxes[1]),
width=boxes[2] - boxes[0],
height=boxes[3] - boxes[1],
fill=False,
edgecolor="green",
linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
模型加载
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=False, num_classes=21, pretrained_backbone=True)
FasterRCNN(
(transform): GeneralizedRCNNTransform(
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
Resize(min_size=(800,), max_size=1333, mode='bilinear')
)
(backbone): BackboneWithFPN(
(body): IntermediateLayerGetter(
(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=1e-05)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): Bottleneck(
(conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=1e-05)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64, eps=1e-05)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256, eps=1e-05)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(1): FrozenBatchNorm2d(256, eps=1e-05)
)
)
(1): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=1e-05)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64, eps=1e-05)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256, eps=1e-05)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(64, eps=1e-05)
(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(64, eps=1e-05)
(conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(256, eps=1e-05)
(relu): ReLU(inplace=True)
)
)
(layer2): Sequential(
(0): Bottleneck(
(conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=1e-05)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=1e-05)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=1e-05)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(512, eps=1e-05)
)
)
(1): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=1e-05)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=1e-05)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=1e-05)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=1e-05)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=1e-05)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=1e-05)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(128, eps=1e-05)
(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(128, eps=1e-05)
(conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(512, eps=1e-05)
(relu): ReLU(inplace=True)
)
)
(layer3): Sequential(
(0): Bottleneck(
(conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=1e-05)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=1e-05)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=1e-05)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(1024, eps=1e-05)
)
)
(1): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=1e-05)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=1e-05)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=1e-05)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=1e-05)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=1e-05)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=1e-05)
(relu): ReLU(inplace=True)
)
(3): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=1e-05)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=1e-05)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=1e-05)
(relu): ReLU(inplace=True)
)
(4): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=1e-05)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=1e-05)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=1e-05)
(relu): ReLU(inplace=True)
)
(5): Bottleneck(
(conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(256, eps=1e-05)
(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(256, eps=1e-05)
(conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(1024, eps=1e-05)
(relu): ReLU(inplace=True)
)
)
(layer4): Sequential(
(0): Bottleneck(
(conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512, eps=1e-05)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512, eps=1e-05)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048, eps=1e-05)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
(1): FrozenBatchNorm2d(2048, eps=1e-05)
)
)
(1): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512, eps=1e-05)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512, eps=1e-05)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048, eps=1e-05)
(relu): ReLU(inplace=True)
)
(2): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d(512, eps=1e-05)
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d(512, eps=1e-05)
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d(2048, eps=1e-05)
(relu): ReLU(inplace=True)
)
)
)
(fpn): FeaturePyramidNetwork(
(inner_blocks): ModuleList(
(0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
(1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
(2): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
(3): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
)
(layer_blocks): ModuleList(
(0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)
(extra_blocks): LastLevelMaxPool()
)
)
(rpn): RegionProposalNetwork(
(anchor_generator): AnchorGenerator()
(head): RPNHead(
(conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))
(bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))
)
)
(roi_heads): RoIHeads(
(box_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(7, 7), sampling_ratio=2)
(box_head): TwoMLPHead(
(fc6): Linear(in_features=12544, out_features=1024, bias=True)
(fc7): Linear(in_features=1024, out_features=1024, bias=True)
)
(box_predictor): FastRCNNPredictor(
(cls_score): Linear(in_features=1024, out_features=21, bias=True)
(bbox_pred): Linear(in_features=1024, out_features=84, bias=True)
)
)
)
RPN正负样本
rpn_label_matchs,rpn_target_deltas= model.rpn.assign_targets_to_anchors(anchors,[targets])
[tensor([0., 0., 0., ..., 0., 0., 0.])]
[tensor([[238.4000, 233.0330, 550.4000, 648.6486],
[238.4000, 233.0330, 550.4000, 648.6486],
[238.4000, 233.0330, 550.4000, 648.6486],
...,
[238.4000, 233.0330, 550.4000, 648.6486],
[238.4000, 233.0330, 550.4000, 648.6486],
[238.4000, 233.0330, 550.4000, 648.6486]], dtype=torch.float64)]
positive_anchor_id= torch.where(torch.eq(rpn_label_matchs[0],1))[0].numpy()
(8,)
torch.Size([159882, 4])
fig = plt.imshow(img)
# 遍历所有的正样本的anchor，绘制在图像上
for idx in positive_anchor_id:
print(idx)
if idx == 158513:
boxes = anchors[0][idx]
fig.axes.add_patch(plt.Rectangle(
xy=(boxes[0],boxes[1]), width=boxes[2]-boxes[0], height=boxes[3]-boxes[1],
fill=False, edgecolor="red", linewidth=1))
plt.show()
158510
158513
158516
158519
158585
158588
158591
158594
<Figure size 432x288 with 1 Axes>
negative_anchor_id= torch.where(torch.eq(rpn_label_matchs[0],-1))[0].numpy()
(249,)
fig = plt.imshow(img)
# 遍历所有的正样本的anchor，绘制在图像上
for idx in negative_anchor_id:
# print(idx)
# 负样本的anchor
boxes = anchors[0][idx]
# 取余20为0的负样本进行展示
fig.axes.add_patch(
plt.Rectangle(xy=(boxes[0], boxes[1]),
width=boxes[2] - boxes[0],
height=boxes[3] - boxes[1],
fill=False,
edgecolor="white",
linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
regression_targets = model.rpn.box_coder.encode(rpn_target_deltas,anchors)
regression_targets[0].size()
torch.Size([159882, 4])
proposals,matched_idx,labels,regression_targets= model.roi_heads.select_training_samples(proposal[0],[targets])
torch.Size([512, 4])
regression_targets[0].size()
torch.Size([512, 4])
[tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 7, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 7])]
positive_proposal_id=torch.where(torch.gt(labels[0],0))[0].numpy()
# 绘制部分候选区域的负样本
fig = plt.imshow(img)
# 遍历所有的正样本的候选区域，绘制在图像上
for idx in positive_proposal_id:
# 正样本的proposals
boxes = proposals[0][idx]
# 取余10为0的负样本进行展示
fig.axes.add_patch(
plt.Rectangle(xy=(boxes[0], boxes[1]),
width=boxes[2] - boxes[0],
height=boxes[3] - boxes[1],
fill=False,
edgecolor="white",
linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
negitivate_proposal_id = torch.where(torch.eq(labels[0],0))[0].numpy()
# 绘制部分候选区域的负样本
fig = plt.imshow(img)
# 遍历所有的负样本的候选区域，绘制在图像上
for idx in negitivate_proposal_id:
# 负样本的proposals
boxes = proposals[0][idx]
# 取余10为0的负样本进行展示
if idx % 10 == 0:
fig.axes.add_patch(
plt.Rectangle(xy=(boxes[0], boxes[1]),
width=boxes[2] - boxes[0],
height=boxes[3] - boxes[1],
fill=False,
edgecolor="white",
linewidth=1))
plt.show()
<Figure size 432x288 with 1 Axes>
模型训练
from torch.utils.data import DataLoader
def collate_fn(batch):
return tuple(zip(*batch))
train_data_loader = DataLoader(dataset,batch_size=2,shuffle = True,collate_fn = collate_fn)
for i,data in enumerate(train_data_loader):
image,targets = data
print(len(image))
print(targets)
break
2
({'boxes': tensor([[224.0000, 140.8000, 425.6000, 522.6666]], dtype=torch.float64), 'labels': tensor([8]), 'image_id': tensor([3183]), 'area': tensor([22554.]), 'iscrowd': tensor([0])}, {'boxes': tensor([[230.4000, 256.0000, 696.0000, 605.8667]], dtype=torch.float64), 'labels': tensor([3]), 'image_id': tensor([2603]), 'area': tensor([47724.]), 'iscrowd': tensor([0])})
# 优化器
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,lr= 0.0001,momentum=0.9,weight_decay=0.0001)
# 遍历每个轮次进行训练
# 1.参数的初始化
iters = 0
loss_value = 0
total_train_loss = []
# 2.遍历每个轮次进行训练
for epoch in range(2):
train_loss = []
# 3.每个轮次中遍历每个batch
for images, targets in train_data_loader:
# 4.把batch数据送入网络中进行预测：损失
loss_dict = model(images, targets)
print(loss_dict)
loss = sum(loss for loss in loss_dict.values())
loss_value = loss.item()
train_loss.append(loss_value)
# 5.进行反向传播：
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch:{epoch},Batch:{iters},loss:{loss_value}")
iters+=1
break
# 6.保存一此列的参数
epoch_loss = np.mean(train_loss)
total_train_loss.append(epoch_loss)
break
{'loss_classifier': tensor(110.5280, grad_fn=<NllLossBackward0>), 'loss_box_reg': tensor(22.8108, grad_fn=<DivBackward0>), 'loss_objectness': tensor(7.9993, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rpn_box_reg': tensor(7.4245, dtype=torch.float64, grad_fn=<DivBackward0>)}
Epoch:0,Batch:0,loss:148.76265928275987