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))