feat image2sketch 新增风格上传 自定义风格

fix

app/service/image2sketch/infer.py

@@ -54,7 +54,7 @@ def load_img(filepath):
 
 if __name__ == '__main__':
     img_A = "/workspace/Semi_ref2sketch_code/datasets/ref_unpair/testA/real_Dress_732caedc416a0cbfedd0e6528040eac7.jpg_Img.jpg"
-    img_B = "/workspace/Semi_ref2sketch_code/datasets/ref_unpair/testC/styleA.png"
+    img_B = "/workspace/Semi_ref2sketch_code/datasets/ref_unpair/testC/style_3.png"
     from opt import Config
 
     opt = Config()  # get  test options
@@ -73,7 +73,7 @@ if __name__ == '__main__':
         model.eval()
     data = {}
     print(os.getcwd())
-    B = reference, _, _ = load_img(r"E:\workspace\trinity_client_aida\app\service\image2sketch\datasets\ref_unpair\testC\styleA.png")
+    B = reference, _, _ = load_img(r"/app/service/image2sketch/datasets/ref_unpair/testC/style_3.png")
     style_img = transform(reference)
     data['B'] = style_img
     data['B'] = data['B'].unsqueeze(0).to(device)

app/service/image2sketch/opt.py

@@ -46,8 +46,12 @@ class Config:
         self.num_test = 1000
         self.morm = 'batch'
         if DEBUG:
-            self.style_image = "service/image2sketch/datasets/ref_unpair/testC/20180422151845_stEe4.jpeg"
+            self.style_image1 = "service/image2sketch/datasets/ref_unpair/testC/style_1.jpg"
+            self.style_image2 = "service/image2sketch/datasets/ref_unpair/testC/style_2.jpeg"
+            self.style_image3 = "service/image2sketch/datasets/ref_unpair/testC/style_3.png"
             self.checkpoints_dir = 'service/image2sketch/checkpoints/'
         else:
             self.checkpoints_dir = 'app/service/image2sketch/checkpoints/'
-            self.style_image = "app/service/image2sketch/datasets/ref_unpair/testC/20180422151845_stEe4.jpeg"
+            self.style_image1 = "app/service/image2sketch/datasets/ref_unpair/testC/style_1.jpg"
+            self.style_image2 = "app/service/image2sketch/datasets/ref_unpair/testC/style_2.jpeg"
+            self.style_image3 = "app/service/image2sketch/datasets/ref_unpair/testC/style_3.png"

app/service/image2sketch/server.py

@@ -33,6 +33,7 @@ def tensor2im(input_image, imtype=np.uint8):
 class Image2SketchServer:
     def __init__(self, request_data):
         self.image_url = request_data.image_url
+        self.style_image_url = request_data.style_image_url
         self.sketch_bucket = request_data.sketch_bucket
         self.sketch_name = request_data.sketch_name
         self.opt = Config()
@@ -47,7 +48,15 @@ class Image2SketchServer:
         self.model.setup(self.opt)
         transform_list = [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
         transform = transforms.Compose(transform_list)
-        style_img = Image.open(self.opt.style_image).convert('L')
+        if request_data.default_style == "1":
+            style_img = Image.open(self.opt.style_image1).convert('L')
+        elif request_data.default_style == "2":
+            style_img = Image.open(self.opt.style_image2).convert('L')
+        elif request_data.default_style == "3":
+            style_img = Image.open(self.opt.style_image3).convert('L')
+        else:
+            style_img = oss_get_image(bucket=self.style_image_url.split('/')[0], object_name=self.style_image_url[self.style_image_url.find('/') + 1:], data_type="PIL")
+        style_img = style_img.convert('L')
         style_img = transform(style_img)
         self.data['B'] = style_img
         self.data['B'] = self.data['B'].unsqueeze(0).to(device)