From 09f3bf0a89e93013b70e976f0507116d5e8d5ed1 Mon Sep 17 00:00:00 2001
From: shiyu22 <shiyu.chen@zilliz.com>
Date: Wed, 6 Apr 2022 17:54:42 +0800
Subject: [PATCH] Covert data type with cv2

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
---
 README.md        | 31 ++++++++++++++++++++++---------
 animegan.py      |  9 ++++++++-
 requirements.txt |  4 +++-
 3 files changed, 33 insertions(+), 11 deletions(-)

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()
 ```
 
+<img src="./results1.png" alt="results1" style="height:120px;"/>
+
+*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()
+```
+
+<img src="./results2.png" alt="results1" style="height:120px;"/>
+
 ## 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