perf(file): 更新目录检查和创建功能

Author: zjZSTU

py/batch_detect.py

@@ -31,11 +31,11 @@
 dst_img_dir = os.path.join(dst_root_dir, 'imgs')
 tmp_json_dir = os.path.join(dst_root_dir, '.tmp_files')
 
-file.check_dir(dst_root_dir)
-file.check_dir(dst_target_dir)
-file.check_dir(dst_pred_dir)
-file.check_dir(dst_img_dir)
-file.check_dir(tmp_json_dir)
+file.make_dir(dst_root_dir, is_rm=True)
+file.make_dir(dst_target_dir, is_rm=True)
+file.make_dir(dst_pred_dir, is_rm=True)
+file.make_dir(dst_img_dir, is_rm=True)
+file.make_dir(tmp_json_dir, is_rm=True)
 
 
 def get_transform():

py/lib/train.py

@@ -92,7 +92,6 @@ def train_model(data_loader, model, criterion, optimizer, lr_scheduler, num_epoc
                 best_model_weights = copy.deepcopy(model.cpu().state_dict())
                 model = model.to(device)
 
-                file.check_dir('../models')
                 file.save_model(best_model_weights, '../models/checkpoint_yolo_v1_%d.pth' % (epoch))
                 print('save model')
 
@@ -118,6 +117,7 @@ def train_model(data_loader, model, criterion, optimizer, lr_scheduler, num_epoc
 
     criterion = MultiPartLoss(S=S, B=B, C=C)
     optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
-    lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.96)
+    lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.9)
 
+    file.make_dir('../models')
     train_model(data_loader, model, criterion, optimizer, lr_scheduler, num_epochs=50, device=device)

py/lib/utils/file.py

@@ -16,10 +16,12 @@
 import glob
 
 
-def check_dir(data_dir):
-    if os.path.exists(data_dir):
-        shutil.rmtree(data_dir)
-    os.mkdir(data_dir)
+def make_dir(data_dir, is_rm=False):
+    if is_rm:
+        if os.path.exists(data_dir):
+            shutil.rmtree(data_dir)
+    if not os.path.exists(data_dir):
+        os.mkdir(data_dir)
 
 
 def parse_location_xml(xml_path):