with open(name_path, "r") as f: classes = [i.strip() for i in f.readlines()] try: im_list = [ osp.join(osp.realpath('.'), images, img) for img in os.listdir(images) ] except NotADirectoryError: im_list = [] im_list.append(osp.join(osp.realpath('.'), images)) except FileNotFoundError: print("No file or directory with the name {}".format(images)) exit() net = Yolov3("test", input_wh, anchors, cfg["anchors_mask"], num_classes) state_dict = torch.load(weightsfile) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7] if head == 'module.': name = k[7:] # remove `module.` else: name = k new_state_dict[name] = v if cuda: net.cuda() cudnn.benchmark = True net.load_state_dict(new_state_dict) print("load weights successfully.....")
cfg = voc_config test_dataset = VOCDetection(cfg["root"], datasets_dict["VOC2007"], input_wh) elif dataset[0] == "C": cfg = coco_config test_dataset = COCODetection(cfg["root"], datasets_dict["COCOval"], input_wh) else: print("only support VOC and COCO datasets !!!") print("load test_dataset successfully.....") with open(cfg["name_path"], "r") as f: classes = [i.strip() for i in f.readlines()] net = Yolov3("test", input_wh, cfg["anchors"], cfg["anchors_mask"], cfg["num_classes"]) state_dict = torch.load(weightsfile) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7] if head == 'module.': name = k[7:] # remove `module.` else: name = k new_state_dict[name] = v if cuda: net.cuda() cudnn.benchmark = True net.load_state_dict(new_state_dict)
pass else: bn = layer[i].bn conv = layer[i].conv ptr = copy_weights(bn, conv, ptr, weights) fp.close() if __name__ == '__main__': args = arg_parse() weightfile = args.weights input_wh = args.input_wh version = args.version save_name = args.save_name if version == "voc": cfg = voc_config yolov3 = Yolov3("train", input_wh, cfg["anchors"], cfg["anchors_mask"], cfg["num_classes"]) elif version == "coco": cfg = coco_config yolov3 = Yolov3("train", input_wh, cfg["anchors"], cfg["anchors_mask"], cfg["num_classes"]) elif version == "darknet53": cfg = voc_config num_blocks = [1,2,8,8,4] yolov3 = Darknet53(num_blocks) else: print("Unkown version !!!") import sys sys.exit() load_weights(weightfile, yolov3, version) # name = "convert_yolo_" + version + ".pth" # save_path = os.path.join("./weights", name)