update readme and other small faults

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 <cfgFile> -weightfile <weightFile> -imgfile <imgFile>
 ```
-python demo.py <cfgFile> <weightFile> <imgFile>
+
+- Load pytorch weights (pth file) to do the inference
+
+```sh
+python models.py <num_classes> <weightfile> <imgfile> <namefile(optional)>
 ```
 
+
 # 3. Darknet2ONNX (Evolving)
 
 - **Pytorch version Recommended: 1.4.0**

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

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 <cfgFile> <weightFile>')