major update

Author: 2damin

README.md

@@ -94,11 +94,12 @@ python main.py --mode onnx --cfg ./cfg/yolov3.cfg --gpus 0 --checkpoint ${saved_
 
 target tensorrt version is 5.x
 
-1. change upsample layer
+1. **ONNX_EXPORT = True** in 'model/yolov3.py'
+   
    tensorrt(v5.x) is not support upsample scale factor, so you have to change upsample layer not using scale factor.
 
 ```{r, engine='bash', count_lines}
-python main.py --mode onnx --cfg ./cfg/yolov3.cfg --gpus 0 --pretrained ./darknet53.conv.74 --checkpoint ${saved_checkpoint_path}
+python main.py --mode onnx --cfg ./cfg/yolov3.cfg --gpus 0 --checkpoint ${saved_checkpoint_path}
 ```
 
 ### option
@@ -111,6 +112,8 @@ python main.py --mode onnx --cfg ./cfg/yolov3.cfg --gpus 0 --pretrained ./darkne
 
 --checkpoint (optional) : the path of saved model checkpoint. Use it when you want to load the previous train, or you want to test(evaluate) the model.
 
+--pretrained (optional) : the path of darknet pretrained weights. Use it when you want to fine-tuning the model.
+
 
 
 ## Visualize training graph

dataloader/data_transforms.py

@@ -14,13 +14,11 @@ def get_transformations(cfg_param = None, is_train = None):
     if is_train:
         data_transform = tf.Compose([AbsoluteLabels(),
                                      DefaultAug(),
-                                     PadSquare(),
                                      RelativeLabels(),
                                      ResizeImage(new_size = (cfg_param['in_width'], cfg_param['in_height'])),
                                      ToTensor(),])
     elif not is_train:
         data_transform = tf.Compose([AbsoluteLabels(),
-                                     PadSquare(),
                                      RelativeLabels(),
                                      ResizeImage(new_size = (cfg_param['in_width'], cfg_param['in_height'])),
                                      ToTensor(),]) 
@@ -51,12 +49,10 @@ def __init__(self, max_objects=50, is_debug=False):
     def __call__(self, data):
         image, labels = data
         if self.is_debug == False:
-            image = torch.tensor(np.transpose(np.array(image, dtype=float) / 255,(2,0,1)),dtype=torch.float32)
+            image = torch.tensor(np.ascontiguousarray(np.transpose(np.array(image, dtype=float) / 255,(2,0,1))),dtype=torch.float32)
         elif self.is_debug == True:
             image = torch.tensor(np.array(image, dtype=float),dtype=torch.float32)
         labels = torch.FloatTensor(np.array(labels))
-        # filled_labels = np.zeros((self.max_objects, 5), np.float32)
-        # filled_labels[range(len(labels))[:self.max_objects]] = labels[:self.max_objects]
         return image, labels
 
 class KeepAspect(object):

dataloader/yolodata.py

@@ -5,13 +5,9 @@
 import matplotlib
 import matplotlib.pyplot as plt
 from PIL import Image
-import matplotlib.patches as patches
 import numpy as np
 from util.tools import *
-from . import data_transforms
 import cv2
-import torchvision
-
 class Yolodata(Dataset):
     file_dir = ""
     anno_dir = ""
@@ -59,7 +55,8 @@ def __getitem__(self, index):
 
         with open(img_path, 'rb') as f:
             img = np.array(Image.open(img_path).convert('RGB'), dtype=np.uint8)
-            #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            #img = cv2.imread(img_path)
+            #img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
             img_origin_h, img_origin_w = img.shape[:2]
 
         #if anno_dir is didnt exist, Test dataset
@@ -79,12 +76,7 @@ def __getitem__(self, index):
                     #skip when no data
                     if len(gt_data) < 5:
                         continue
-                    #ignore very tiny GTs
                     cx, cy, w, h = float(gt_data[1]), float(gt_data[2]), float(gt_data[3]), float(gt_data[4])
-                    #trunc = float(gt_data[5]) if len(gt_data) > 5 else 0
-                    #occ = float(gt_data[6]) if len(gt_data) > 6 else 0
-                    # if w <= 20 or float(gt_data[7]) - float(gt_data[5]) <= 20:
-                    #     continue
                     bbox.append([float(gt_data[0]), cx, cy, w, h])
 
             #Change gt_box type
@@ -107,10 +99,9 @@ def __getitem__(self, index):
                 return
             return img, target_data, anno_path
         else:
-            bbox = np.array([[0,0,0,0,0]])
+            bbox = np.array([[0,0,0,0,0]], dtype=np.float64)
             img, _ = self.transform((img, bbox))
             return img, None, None
 
-
     def __len__(self):
         return len(self.img_data)

demo/demo.py

@@ -18,27 +18,28 @@ def run(self):
                 continue
             input_img, _, _ = batch
 
+            #drawBox(input_img.detach().numpy()[0])
+            #np.save("torch_input.npy",input_img.detach().numpy())
+            
             input_img = input_img.to(self.device, non_blocking=True)
 
             num_batch = input_img.shape[0]
             with torch.no_grad():
                 start_time = time.time()
 
                 output = self.model(input_img)
-                            
-                # _, output_list = self.yololoss.compute_loss(output, targets = None, nw = nw, nh = nh, yolo_layers = self.model.yolo_layers)
-                best_box_list = non_max_suppression(output, conf_thres=0.25, iou_thres=0.45)
+                best_box_list = non_max_suppression(output,
+                                                    conf_thres=0.4,
+                                                    iou_thres=0.45)
 
                 for b in range(num_batch):
-                                    
                     if best_box_list[b] is None:
                         continue
+                    print(best_box_list[b])
                     final_box_list = [bbox for bbox in best_box_list[b] if bbox[4] > 0.5]
                     print("final :", final_box_list)
 
                     if final_box_list is None:
                         continue
                     show_img = input_img[b].detach().cpu().numpy()
                     drawBoxlist(show_img, final_box_list, mode=1)
-
-        

eval/evaluator.py

@@ -3,7 +3,6 @@
 import PIL, time
 from train.loss import *
 import csv
-import pandas as pd
 from terminaltables import AsciiTable
 
 class Evaluator:

main.py

@@ -1,4 +1,5 @@
 import os,sys
+from pandas import cut
 # os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
 # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -58,9 +59,12 @@ def collate_fn(batch):
     imgs, targets, anno_path = list(zip(*batch))
 
     imgs = torch.stack([img for img in imgs])
+    
+    if targets[0] is None or anno_path[0] is None:
+        return imgs, None, None
+
     for i, boxes in enumerate(targets):
         boxes[:, 0] = i
-
     targets = torch.cat(targets,0)
 
     return imgs, targets, anno_path
@@ -202,6 +206,21 @@ def demo(cfg_param = None, using_gpus = None):
     demo_loader = DataLoader(data, batch_size = 1, num_workers = 0, pin_memory = True, drop_last = False, shuffle = False, collate_fn=collate_fn)
 
     model = DarkNet53(args.cfg, cfg_param)
+    model.eval()
+
+    #load pre-trained darknet weights
+    if args.pretrained is not None:
+        print("load pretrained model")
+        model.load_darknet_weights(args.pretrained)
+        chkpt = {'epoch': -1,
+                 'best_fitness': None,
+                 'training_results': None,
+                 'model': model.state_dict(),
+                 'optimizer': None}
+        target = args.pretrained.replace(".weights", ".pth")
+        torch.save(chkpt, target)
+    else:
+        model.initialize_weights()
 
     if args.checkpoint is not None:
         print("load pretrained model ", args.checkpoint)
@@ -220,6 +239,13 @@ def demo(cfg_param = None, using_gpus = None):
 
     model = model.to(device)
     model.eval()
+    
+    if args.pretrained is not None:
+        darknet_weights_name = args.pretrained.replace(".weights", "_new.weights")
+    elif args.checkpoint is not None:
+        darknet_weights_name = args.checkpoint.replace(".pth", ".weights")
+    model.save_darknet_weights(darknet_weights_name, cutoff=-1)
+
     torch.backends.cudnn.benchmark = True
 
     demo = Demo(model, data, demo_loader, device, cfg_param)
@@ -233,6 +259,8 @@ def torch2onnx(cfg_param = None, using_gpus = None):
     cfg_param['batch'] = 1
     model = DarkNet53(args.cfg, cfg_param)
 
+    if args.pretrained is not None:
+        model.load_darknet_weights(args.pretrained)
     #Set the device what you use, GPU or CPU
     for i in using_gpus:
         print("GPU total memory : {} free memory : {}".format(get_memory_total_MiB(i), get_memory_free_MiB(i)))
@@ -260,28 +288,15 @@ def torch2onnx(cfg_param = None, using_gpus = None):
     #If checkpoint is existed, load the previous checkpoint.
     checkpoint = None
     if args.checkpoint is not None:
-        print("load pretrained model ", args.checkpoint)
-        checkpoint = torch.load(args.checkpoint)
-        if len(using_gpus) == 0:
-            post_name = ""
-        else:
-            pass
-        #     post_name = "module."
-        # for key, value in checkpoint['model_state_dict'].copy().items():
-        #     new_key = post_name + key
-        #     checkpoint['model_state_dict'][new_key] = checkpoint['model_state_dict'].pop(key)
+        print("load checkpoint model ", args.checkpoint)
+        checkpoint = torch.load(args.checkpoint, map_location=torch.device('cpu'))
         model.load_state_dict(checkpoint['model_state_dict'])
-
-    # if len(using_gpus) > 0:
-    #     yolo_model = model.module
-    # else:
-    #     yolo_model = model
 
     model.eval()
 
     darknet_weights_name = args.checkpoint.replace(".pth", ".weights")
     onnx_weights_name = args.checkpoint.replace(".pth", ".onnx")
-    model.save_darknet_weights(darknet_weights_name)
+    model.save_darknet_weights(darknet_weights_name, cutoff=-1)
 
     #Pre-check the model structure and size of parameters
     #summary.summary(yolo_model, input_size=(3, cfg_param["in_width"], cfg_param["in_height"]), device='cuda' if device == torch.device('cuda') else 'cpu')