更强的模型

cyclegan.py

@@ -19,6 +19,10 @@ class CYCLEGAN(object):
         #-----------------------------------------------#
         "input_shape"       : [112, 112],
         #-------------------------------#
+        #   是否进行不失真的resize
+        #-------------------------------#
+        "letterbox_image"   : True,
+        #-------------------------------#
         #   是否使用Cuda
         #   没有GPU可以设置成False
         #-------------------------------#
@@ -64,9 +68,14 @@ class CYCLEGAN(object):
         #---------------------------------------------------------#
         image       = cvtColor(image)
         #---------------------------------------------------------#
+        #   给图像增加灰条，实现不失真的resize
+        #   也可以直接resize进行识别
+        #---------------------------------------------------------#
+        image_data, nw, nh = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
+        #---------------------------------------------------------#
         #   添加上batch_size维度
         #---------------------------------------------------------#
-        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image, dtype='float32')), (2, 0, 1)), 0)
+        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
         
         with torch.no_grad():
             images = torch.from_numpy(image_data)
@@ -80,10 +89,17 @@ class CYCLEGAN(object):
             #---------------------------------------------------#
             #   转为numpy
             #---------------------------------------------------#
-            pr = pr.permute(1, 2, 0).cpu().numpy()      
+            pr = pr.permute(1, 2, 0).cpu().numpy()
+            
+            #--------------------------------------#
+            #   将灰条部分截取掉
+            #--------------------------------------#
+            if nw is not None:
+                pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \
+                        int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]
+                
             
         image = postprocess_output(pr)
-        image = np.clip(image, 0, 255)
         image = Image.fromarray(np.uint8(image))
 
         return image

model_data/G_model_B2A_last_epoch_weights.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:03aa46c6ebd8da9196749e02e214687ea3dd7143976d030f92b91b0cb2496583
+oid sha256:1815cd8f77471a8712b9a80b20da4cd7afe7aad2b32ad48cd205d1c370a65dc2
 size 11888773