init

app.py

@@ -19,7 +19,7 @@ import cv2
 from io import BytesIO
 sys.path.insert(0, 'p2c')
 
-from test_onnx import Photo2Cartoon
+from test import Photo2Cartoon
 
 
 ORIGINAL_REPO_URL = 'https://github.com/minivision-ai/photo2cartoon'
@@ -47,9 +47,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument('--allow-screenshot', action='store_true')
     return parser.parse_args()
 
-
-
-
 def run(
     image,
     p2c : Photo2Cartoon,

p2c/test.py

@@ -7,25 +7,22 @@ import argparse
 from utils import Preprocess
 
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--photo_path', type=str, help='input photo path')
-parser.add_argument('--save_path', type=str, help='cartoon save path')
-args = parser.parse_args()
-
-os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
-
 class Photo2Cartoon:
     def __init__(self):
         self.pre = Preprocess()
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         self.net = ResnetGenerator(ngf=32, img_size=256, light=True).to(self.device)
-        
-        assert os.path.exists('./models/photo2cartoon_weights.pt'), "[Step1: load weights] Can not find 'photo2cartoon_weights.pt' in folder 'models!!!'"
-        params = torch.load('./models/photo2cartoon_weights.pt', map_location=self.device)
+
+        curPath = os.path.abspath(os.path.dirname(__file__))
+
+        #assert os.path.exists('./models/photo2cartoon_weights.pt'), "[Step1: load weights] Can not find 'photo2cartoon_weights.pt' in folder 'models!!!'"
+        params = torch.load(os.path.join(curPath, 'models/photo2cartoon_weights.pt'), map_location=self.device)
         self.net.load_state_dict(params['genA2B'])
         print('[Step1: load weights] success!')
 
-    def inference(self, img):
+    def inference(self, in_path):
+        img = cv2.cvtColor(cv2.imread(in_path), cv2.COLOR_BGR2RGB)
+
         # face alignment and segmentation
         face_rgba = self.pre.process(img)
         if face_rgba is None:
@@ -49,15 +46,9 @@ class Photo2Cartoon:
         cartoon = np.transpose(cartoon.cpu().numpy(), (1, 2, 0))
         cartoon = (cartoon + 1) * 127.5
         cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
-        cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
+        #cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
         print('[Step3: photo to cartoon] success!')
         return cartoon
 
 
-if __name__ == '__main__':
-    img = cv2.cvtColor(cv2.imread(args.photo_path), cv2.COLOR_BGR2RGB)
-    c2p = Photo2Cartoon()
-    cartoon = c2p.inference(img)
-    if cartoon is not None:
-        cv2.imwrite(args.save_path, cartoon)
-        print('Cartoon portrait has been saved successfully!')
+