diff --git a/README.md b/README.md index 049dab42..31a24556 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ A minimal PyTorch implementation of YOLOv4. ![image](https://user-gold-cdn.xitu.io/2020/4/26/171b5a6c8b3bd513?w=768&h=576&f=jpeg&s=78882) -# 0. Weight +# 0. Weights Download ## 0.1 darkent - baidu(https://pan.baidu.com/s/1dAGEW8cm-dqK14TbhhVetA Extraction code:dm5b) @@ -77,11 +77,20 @@ you can use darknet2pytorch to convert it yourself, or download my converted mod ``` # 2. Inference -- download model weight https://drive.google.com/open?id=1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT + +- Load the pretrained darknet model and darknet weights to do the inference + +```sh +python demo.py -cfgfile -weightfile -imgfile ``` -python demo.py + +- Load pytorch weights (pth file) to do the inference + +```sh +python models.py ``` + # 3. Darknet2ONNX (Evolving) - **Pytorch version Recommended: 1.4.0** diff --git a/models.py b/models.py index efa82481..0bb8588e 100644 --- a/models.py +++ b/models.py @@ -405,10 +405,10 @@ def __init__(self, yolov4conv137weight=None, n_classes=80, yolo_layer_included=F self.down4 = DownSample4() self.down5 = DownSample5() # neck - self.neck = Neck() + self.neek = Neck() # yolov4conv137 if yolov4conv137weight: - _model = nn.Sequential(self.down1, self.down2, self.down3, self.down4, self.down5, self.neck) + _model = nn.Sequential(self.down1, self.down2, self.down3, self.down4, self.down5, self.neek) pretrained_dict = torch.load(yolov4conv137weight) model_dict = _model.state_dict() @@ -429,7 +429,7 @@ def forward(self, input): d4 = self.down4(d3) d5 = self.down5(d4) - x20, x13, x6 = self.neck(d5, d4, d3) + x20, x13, x6 = self.neek(d5, d4, d3) output = self.head(x20, x13, x6) return output @@ -455,7 +455,7 @@ def forward(self, input): model = Yolov4(n_classes=n_classes, yolo_layer_included=True) - pretrained_dict = torch.load(weightfile, map_location=torch.device('cpu')) + pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) model.load_state_dict(pretrained_dict) if namesfile == None: @@ -466,7 +466,7 @@ def forward(self, input): else: print("please give namefile") - use_cuda = 0 + use_cuda = True if use_cuda: model.cuda() diff --git a/tool/darknet2onnx.py b/tool/darknet2onnx.py index c73adefb..c44cb466 100644 --- a/tool/darknet2onnx.py +++ b/tool/darknet2onnx.py @@ -35,12 +35,12 @@ def transform_to_onnx(cfgfile, weightfile, batch_size=1): if len(sys.argv) == 3: cfgfile = sys.argv[1] weightfile = sys.argv[2] - fransform_to_onnx(cfgfile, weightfile) + transform_to_onnx(cfgfile, weightfile) elif len(sys.argv) == 4: cfgfile = sys.argv[1] weightfile = sys.argv[2] batch_size = int(sys.argv[3]) - fransform_to_onnx(cfgfile, weightfile, batch_size) + transform_to_onnx(cfgfile, weightfile, batch_size) else: print('Please execute this script this way:\n') print(' python darknet2onnx.py ')