Add files via upload

Author: bubbliiiing

frcnn.py

@@ -19,10 +19,11 @@
 #--------------------------------------------#
 class FRCNN(object):
     _defaults = {
-        "model_path": 'model_data/voc_weights_resnet.pth',
-        "classes_path": 'model_data/voc_classes.txt',
-        "confidence": 0.5,
-        "backbone": "resnet50"
+        "model_path"    : 'model_data/voc_weights_resnet.pth',
+        "classes_path"  : 'model_data/voc_classes.txt',
+        "confidence"    : 0.5,
+        "iou"           : 0.45,
+        "backbone"      : "resnet50"
     }
 
     @classmethod
@@ -86,9 +87,11 @@ def detect_image(self, image):
         old_height = image_shape[0]
         old_image = copy.deepcopy(image)
         width,height = get_new_img_size(old_width,old_height)
-        image = image.resize([width,height])
+
+        image = image.resize([width,height], Image.BICUBIC)
         photo = np.array(image,dtype = np.float32)/255
         photo = np.transpose(photo, (2, 0, 1))
+        
         with torch.no_grad():
             images = []
             images.append(photo)
@@ -97,7 +100,7 @@ def detect_image(self, image):
 
             roi_cls_locs, roi_scores, rois, roi_indices = self.model(images)
             decodebox = DecodeBox(self.std, self.mean, self.num_classes)
-            outputs = decodebox.forward(roi_cls_locs, roi_scores, rois, height=height, width=width, score_thresh = self.confidence)
+            outputs = decodebox.forward(roi_cls_locs, roi_scores, rois, height = height, width = width, nms_iou = self.iou, score_thresh = self.confidence)
             if len(outputs)==0:
                 return old_image
             bbox = outputs[:,:4]
@@ -107,6 +110,7 @@ def detect_image(self, image):
             bbox[:, 0::2] = (bbox[:, 0::2])/width*old_width
             bbox[:, 1::2] = (bbox[:, 1::2])/height*old_height
             bbox = np.array(bbox,np.int32)
+
         image = old_image
         thickness = (np.shape(old_image)[0] + np.shape(old_image)[1]) // old_width*2
         font = ImageFont.truetype(font='model_data/simhei.ttf',size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))

get_dr_txt.py

@@ -3,32 +3,37 @@
 #   具体视频教程可查看
 #   https://www.bilibili.com/video/BV1zE411u7Vw
 #----------------------------------------------------#
-from frcnn import FRCNN
-from PIL import Image
-from torch.autograd import Variable
-import torch
-import numpy as np
+import copy
 import os
+
+import numpy as np
+import torch
 import torch.backends.cudnn as cudnn
+from PIL import Image, ImageDraw, ImageFont
+from torch.autograd import Variable
 from torch.nn import functional as F
-from utils.utils import loc2bbox, nms, DecodeBox
+from tqdm import tqdm
+
+from frcnn import FRCNN
 from nets.frcnn import FasterRCNN
 from nets.frcnn_training import get_new_img_size
-from PIL import Image, ImageFont, ImageDraw
-import copy
+from utils.utils import DecodeBox, loc2bbox, nms
+
 
 class mAP_FRCNN(FRCNN):
     #---------------------------------------------------#
     #   检测图片
     #---------------------------------------------------#
     def detect_image(self,image_id,image):
-        self.confidence = 0.05
+        self.confidence = 0.01
+        self.iou        = 0.45
         f = open("./input/detection-results/"+image_id+".txt","w") 
         image_shape = np.array(np.shape(image)[0:2])
         old_width = image_shape[1]
         old_height = image_shape[0]
         width,height = get_new_img_size(old_width,old_height)
-        image = image.resize([width,height])
+        
+        image = image.resize([width,height], Image.BICUBIC)
         photo = np.array(image,dtype = np.float32)/255
         photo = np.transpose(photo, (2, 0, 1))
         with torch.no_grad():
@@ -39,7 +44,7 @@ def detect_image(self,image_id,image):
 
             roi_cls_locs, roi_scores, rois, roi_indices = self.model(images)
             decodebox = DecodeBox(self.std, self.mean, self.num_classes)
-            outputs = decodebox.forward(roi_cls_locs, roi_scores, rois, height=height, width=width, score_thresh = self.confidence)
+            outputs = decodebox.forward(roi_cls_locs, roi_scores, rois, height = height, width = width, nms_iou = self.iou, score_thresh = self.confidence)
             if len(outputs)==0:
                 return 
             bbox = outputs[:,:4]
@@ -71,12 +76,11 @@ def detect_image(self,image_id,image):
     os.makedirs("./input/images-optional")
 
 
-for image_id in image_ids:
+for image_id in tqdm(image_ids):
     image_path = "./VOCdevkit/VOC2007/JPEGImages/"+image_id+".jpg"
     image = Image.open(image_path)
-    image.save("./input/images-optional/"+image_id+".jpg")
+    # image.save("./input/images-optional/"+image_id+".jpg")
     frcnn.detect_image(image_id,image)
-    print(image_id," done!")
 
 
 print("Conversion completed!")

nets/frcnn.py

@@ -63,4 +63,9 @@ def forward(self, x, scale=1.):
         # print(np.shape(rois))
         # print(roi_indices)
         roi_cls_locs, roi_scores = self.head.forward(h, rois, roi_indices)
-        return roi_cls_locs, roi_scores, rois, roi_indices
+        return roi_cls_locs, roi_scores, rois, roi_indices
+
+    def freeze_bn(self):
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()

nets/rpn.py

@@ -8,21 +8,12 @@
 import numpy as np
 
 
-'''
-一些建议的参数设置：
-VGG：SGD优化器，冻结时学习率1e-3，解冻时学习率1e-4
-    nets.rpn中ProposalCreator的n_train_post_nms=2000；
-    utils.utils中ProposalTargetCreator的pos_ratio=0.25；
-RESNET50：Adam优化器，冻结时学习率1e-4，解冻时学习率1e-5
-    nets.rpn中ProposalCreator的n_train_post_nms=300；
-    utils.utils中ProposalTargetCreator的pos_ratio=0.5;
-'''
 class ProposalCreator():
     def __init__(self,
                  mode,
                  nms_thresh=0.7,
                  n_train_pre_nms=12000,
-                 n_train_post_nms=300,
+                 n_train_post_nms=2000,
                  n_test_pre_nms=3000,
                  n_test_post_nms=300,
                  min_size=16

train.py

@@ -75,15 +75,6 @@ def fit_ont_epoch(net,epoch,epoch_size,epoch_size_val,gen,genval,Epoch):
     print('Saving state, iter:', str(epoch+1))
     torch.save(model.state_dict(), 'logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.pth'%((epoch+1),total_loss/(epoch_size+1),val_toal_loss/(epoch_size_val+1)))
 
-'''
-一些建议的参数设置：
-VGG：SGD优化器，冻结时学习率1e-3，解冻时学习率1e-4
-    nets.rpn中ProposalCreator的n_train_post_nms=2000；
-    utils.utils中ProposalTargetCreator的pos_ratio=0.25；
-RESNET50：Adam优化器，冻结时学习率1e-4，解冻时学习率1e-5
-    nets.rpn中ProposalCreator的n_train_post_nms=300；
-    utils.utils中ProposalTargetCreator的pos_ratio=0.5;
-'''
 if __name__ == "__main__":
     # 参数初始化
     annotation_path = '2007_train.txt'
@@ -118,22 +109,12 @@ def fit_ont_epoch(net,epoch,epoch_size,epoch_size_val,gen,genval,Epoch):
     num_val = int(len(lines)*val_split)
     num_train = len(lines) - num_val
 
-    '''
-    一些建议的参数设置：
-    VGG：SGD优化器，冻结时学习率1e-3，解冻时学习率1e-4
-        nets.rpn中ProposalCreator的n_train_post_nms=2000；
-        utils.utils中ProposalTargetCreator的pos_ratio=0.25；
-    RESNET50：Adam优化器，冻结时学习率1e-4，解冻时学习率1e-5
-        nets.rpn中ProposalCreator的n_train_post_nms=300；
-        utils.utils中ProposalTargetCreator的pos_ratio=0.5;
-    '''
     if True:
         lr = 1e-4
         Init_Epoch = 0
-        Freeze_Epoch = 25
+        Freeze_Epoch = 50
 
         optimizer = optim.Adam(model.parameters(),lr,weight_decay=5e-4)
-        # optimizer = optim.SGD(model.parameters(),lr,weight_decay=5e-4,momentum=0.9)
         lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)
 
         if Use_Data_Loader:
@@ -158,18 +139,18 @@ def fit_ont_epoch(net,epoch,epoch_size,epoch_size_val,gen,genval,Epoch):
         # ------------------------------------#
         #   由于batch==1所以冻结bn层
         # ------------------------------------#
-        model = model.eval()
+        model.freeze_bn()
 
         for epoch in range(Init_Epoch,Freeze_Epoch):
             fit_ont_epoch(model,epoch,epoch_size,epoch_size_val,gen,gen_val,Freeze_Epoch)
             lr_scheduler.step()
 
     if True:
         lr = 1e-5
-        Freeze_Epoch = 25
-        Unfreeze_Epoch = 50
+        Freeze_Epoch = 50
+        Unfreeze_Epoch = 100
+
         optimizer = optim.Adam(model.parameters(),lr,weight_decay=5e-4)
-        # optimizer = optim.SGD(model.parameters(),lr,weight_decay=5e-4,momentum=0.9)
         lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)
 
         if Use_Data_Loader:
@@ -194,7 +175,7 @@ def fit_ont_epoch(net,epoch,epoch_size,epoch_size_val,gen,genval,Epoch):
         # ------------------------------------#
         #   由于batch==1所以冻结bn层
         # ------------------------------------#
-        model = model.eval()
+        model.freeze_bn()
 
         for epoch in range(Freeze_Epoch,Unfreeze_Epoch):
             fit_ont_epoch(model,epoch,epoch_size,epoch_size_val,gen,gen_val,Unfreeze_Epoch)

utils/utils.py

@@ -87,7 +87,7 @@ def __init__(self, std, mean, num_classes):
         self.mean = mean
         self.num_classes = num_classes + 1    
 
-    def forward(self, roi_cls_locs, roi_scores, rois, height, width, score_thresh):
+    def forward(self, roi_cls_locs, roi_scores, rois, height, width, nms_iou, score_thresh):
 
         rois = torch.Tensor(rois)
 
@@ -130,24 +130,16 @@ def forward(self, roi_cls_locs, roi_scores, rois, height, width, score_thresh):
 
             prob_l_index = np.argsort(prob_l)[::-1]
             detections_class = detections_class[prob_l_index]
-            nms_out = nms(detections_class, 0.3)
+            nms_out = nms(detections_class, nms_iou)
             if outputs==[]:
                 outputs = nms_out
             else:
                 outputs = np.concatenate([outputs, nms_out],axis=0)
         return outputs
-'''
-一些建议的参数设置：
-VGG：SGD优化器，冻结时学习率1e-3，解冻时学习率1e-4
-    nets.rpn中ProposalCreator的n_train_post_nms=2000；
-    utils.utils中ProposalTargetCreator的pos_ratio=0.25；
-RESNET50：Adam优化器，冻结时学习率1e-4，解冻时学习率1e-5
-    nets.rpn中ProposalCreator的n_train_post_nms=300；
-    utils.utils中ProposalTargetCreator的pos_ratio=0.5;
-'''
+        
 class ProposalTargetCreator(object):
     def __init__(self,n_sample=128,
-                 pos_ratio=0.5, pos_iou_thresh=0.5,
+                 pos_ratio=0.25, pos_iou_thresh=0.5,
                  neg_iou_thresh_hi=0.5, neg_iou_thresh_lo=0
                  ):
         self.n_sample = n_sample