aidlab/AiDA_Python
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
fix
This commit is contained in:
zhouchengrong
2024-09-20 17:03:48 +08:00
parent 1385fde9ce
commit 2e07bc2de9
9 changed files with 28 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
app/service/image2sketch/server.py
View File

@@ -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)