def create_model(num_classes, device):
# https://download.pytorch.org/models/vgg16-397923af.pth
# 如果使用mobilenetv2的话就下载对应预训练权重并注释下面三行，接着把mobilenetv2模型对应的两行代码注释取消掉
vgg_feature = vgg(model_name="vgg16",
weights_path="./backbone/vgg16.pth").features
backbone = torch.nn.Sequential(*list(
vgg_feature._modules.values())[:-1]) # 删除feature中最后的maxpool层
backbone.out_channels = 512
# https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
# backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features
# backbone.out_channels = 1280 # 设置对应backbone输出特征矩阵的channels
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def create_model(num_classes):
vgg_feature = vgg(model_name="vgg16").features
backbone = torch.nn.Sequential(*list(
vgg_feature._modules.values())[:-1]) # 删除feature中最后的maxpool层
backbone.out_channels = 512
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
return model
def build_ssd(phase, size=300, num_classes=5):
if phase != "test" and phase != "train":
print("ERROR: Phase: " + phase + " not recognized")
return
if size != 300:
print("ERROR: You specified size " + repr(size) + ". However, " +
"currently only SSD300 (size=300) is supported!")
return
base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
add_extras(extras[str(size)], 1024),
mbox[str(size)], num_classes)
# print("-----------")
# for i,v in enumerate(base_):
# print(i,v)
# print("-----------")
# for i,e in enumerate(extras_):
# print(i,e)
# print("-----------")
# for i,h in enumerate(head_):
# print(i,h)
return SSD(phase, size, base_, extras_, head_, num_classes)
def main(parser_data):
device = torch.device(
parser_data.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
data_transform = {"val": transforms.Compose([transforms.ToTensor()])}
# read class_indict
label_json_path = './coco80_indices.json'
assert os.path.exists(
label_json_path), "json file {} dose not exist.".format(
label_json_path)
json_file = open(label_json_path, 'r')
category_index = json.load(json_file)
coco_root = parser_data.data_path
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = parser_data.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
# load validation data set
val_dataset = CocoDetection(coco_root, "val", data_transform["val"])
val_dataset_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=val_dataset.collate_fn)
# create model
vgg_feature = vgg(model_name="vgg16",
weights_path="./backbone/vgg16.pth").features
backbone = torch.nn.Sequential(*list(
vgg_feature._modules.values())[:-1]) # 删除feature中最后的maxpool层
backbone.out_channels = 512
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
# num_classes equal 80 + background classes
model = FasterRCNN(backbone=backbone,
num_classes=parser_data.num_classes + 1,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# 载入你自己训练好的模型权重
weights_path = parser_data.weights
assert os.path.exists(weights_path), "not found {} file.".format(
weights_path)
weights_dict = torch.load(weights_path, map_location=device)
model.load_state_dict(weights_dict['model'])
# print(model)
model.to(device)
# evaluate on the val dataset
cpu_device = torch.device("cpu")
coco91to80 = val_dataset.coco91to80
coco80to91 = dict([(str(v), k) for k, v in coco91to80.items()])
results = []
model.eval()
with torch.no_grad():
for image, targets in tqdm(val_dataset_loader, desc="validation..."):
# 将图片传入指定设备device
image = list(img.to(device) for img in image)
# inference
outputs = model(image)
outputs = [{k: v.to(cpu_device)
for k, v in t.items()} for t in outputs]
# 遍历每张图像的预测结果
for target, output in zip(targets, outputs):
if len(output) == 0:
continue
img_id = int(target["image_id"])
per_image_boxes = output["boxes"]
# 对于coco_eval, 需要的每个box的数据格式为[x_min, y_min, w, h]
# 而我们预测的box格式是[x_min, y_min, x_max, y_max]，所以需要转下格式
per_image_boxes[:, 2:] -= per_image_boxes[:, :2]
per_image_classes = output["labels"]
per_image_scores = output["scores"]
# 遍历每个目标的信息
for object_score, object_class, object_box in zip(
per_image_scores, per_image_classes, per_image_boxes):
object_score = float(object_score)
# 要将类别信息还原回coco91中
coco80_class = int(object_class)
coco91_class = int(coco80to91[str(coco80_class)])
# We recommend rounding coordinates to the nearest tenth of a pixel
# to reduce resulting JSON file size.
object_box = [round(b, 2) for b in object_box.tolist()]
res = {
"image_id": img_id,
"category_id": coco91_class,
"bbox": object_box,
"score": round(object_score, 3)
}
results.append(res)
# accumulate predictions from all images
# write predict results into json file
json_str = json.dumps(results, indent=4)
with open('predict_tmp.json', 'w') as json_file:
json_file.write(json_str)
# accumulate predictions from all images
coco_true = val_dataset.coco
coco_pre = coco_true.loadRes('predict_tmp.json')
coco_evaluator = COCOeval(cocoGt=coco_true,
cocoDt=coco_pre,
iouType="bbox")
coco_evaluator.evaluate()
coco_evaluator.accumulate()
coco_evaluator.summarize()
# calculate COCO info for all classes
coco_stats, print_coco = summarize(coco_evaluator)
# calculate voc info for every classes(IoU=0.5)
voc_map_info_list = []
for i in range(len(category_index)):
stats, _ = summarize(coco_evaluator, catId=i)
voc_map_info_list.append(" {:15}: {}".format(
category_index[str(i + 1)], stats[1]))
print_voc = "\n".join(voc_map_info_list)
print(print_voc)
# 将验证结果保存至txt文件中
with open("record_mAP.txt", "w") as f:
record_lines = [
"COCO results:", print_coco, "", "mAP(IoU=0.5) for each category:",
print_voc
]
f.write("\n".join(record_lines))
