修复

cyclegan.py

@@ -1,4 +1,3 @@
-import cv2
 import numpy as np
 import torch
 from PIL import Image
@@ -15,19 +14,11 @@ class CYCLEGAN(object):
         #   model_path指向logs文件夹下的权值文件
         #-----------------------------------------------#
         "model_path"        : 'model_data/G_model_B2A_last_epoch_weights.pth',
-        #-----------------------------------------------#
-        #   输入图像大小的设置
-        #-----------------------------------------------#
-        "input_shape"       : [112, 112],
-        #-------------------------------#
-        #   是否进行不失真的resize
-        #-------------------------------#
-        "letterbox_image"   : True,
         #-------------------------------#
         #   是否使用Cuda
         #   没有GPU可以设置成False
         #-------------------------------#
-        "cuda"              : False,
+        "cuda"              : True,
     }
 
     #---------------------------------------------------#
@@ -68,16 +59,6 @@ class CYCLEGAN(object):
         #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
         #---------------------------------------------------------#
         image       = cvtColor(image)
-        #---------------------------------------------------#
-        #   获得高宽
-        #---------------------------------------------------#
-        orininal_h  = np.array(image).shape[0]
-        orininal_w  = np.array(image).shape[1]
-        #---------------------------------------------------------#
-        #   给图像增加灰条，实现不失真的resize
-        #   也可以直接resize进行识别
-        #---------------------------------------------------------#
-        image_data, nw, nh = resize_image(image, (self.input_shape[1],self.input_shape[0]), self.letterbox_image)
         #---------------------------------------------------------#
         #   添加上batch_size维度
         #---------------------------------------------------------#
@@ -95,19 +76,7 @@ class CYCLEGAN(object):
             #---------------------------------------------------#
             #   转为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)]
-                
-            #---------------------------------------------------#
-            #   进行图片的resize
-            #---------------------------------------------------#
-            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR)
+            pr = pr.permute(1, 2, 0).cpu().numpy()      
             
         image = postprocess_output(pr)
         image = np.clip(image, 0, 255)