diff --git a/README.md b/README.md index 54da127..3a4f5eb 100644 --- a/README.md +++ b/README.md @@ -7,22 +7,35 @@ Convert an image into an animated image using [`AnimeganV2`](https://github.com/ ## Code Example -Load an image from path './image.png'. +Load an image from path './test.png'. *Write the pipeline in simplified style*: ```python import towhee -from PIL import Image -import numpy - -pipeline = towhee.glob('./image.png').image_decode.cv2().img2img_translation.animegan(model_name = 'hayao') -img = pipeline.to_list()[0] -img = numpy.transpose(img, (1,2,0)) -img = Image.fromarray((img * 255).astype(numpy.uint8)) -img.show() + +towhee.glob('/Users/chenshiyu/workspace/data/pic/test.png') \ + .image_decode.cv2() \ + .img2img_translation.animegan(model_name = 'hayao') \ + .show() ``` +results1 + +*Write a same pipeline with explicit inputs/outputs name specifications:* + +```python +import towhee + +towhee.glob['path']('/Users/chenshiyu/workspace/data/pic/test.png') \ + .image_decode.cv2['path', 'origin']() \ + .img2img_translation.animegan['origin', 'transformed'](model_name = 'hayao') \ + .select('origin', 'transformed') \ + .show() +``` + +results1 + ## Factory Constructor Create the operator via the following factory method diff --git a/animegan.py b/animegan.py index 47ed90f..0982a80 100644 --- a/animegan.py +++ b/animegan.py @@ -1,5 +1,7 @@ import os +import numpy from pathlib import Path +from PIL import Image as PImage from torchvision import transforms from towhee import register @@ -30,6 +32,11 @@ class Animegan(Operator): @arg(1, to_image_color('RGB')) def __call__(self, image): img = self.tfms(image).unsqueeze(0) - print(img.shape) styled_image = self.model(img) + + styled_image = numpy.transpose(styled_image, (1,2,0)) + styled_image = PImage.fromarray((styled_image * 255).astype(numpy.uint8)) + styled_image = numpy.array(styled_image) + styled_image = styled_image[:, :, ::-1].copy() + return Image(styled_image, 'RGB') diff --git a/requirements.txt b/requirements.txt index 8414abf..85f820d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,4 @@ pathlib -torchvision \ No newline at end of file +torchvision +pillow +opencv-python \ No newline at end of file