| from pathlib import Path | |
| from inference import Predictor as MyPredictor | |
| from utils import read_image | |
| import cv2 | |
| import tempfile | |
| from utils.image_processing import resize_image, normalize_input, denormalize_input | |
| import numpy as np | |
| from cog import BasePredictor, Path, Input | |
| class Predictor(BasePredictor): | |
| def setup(self): | |
| pass | |
| def predict( | |
| self, | |
| image: Path = Input(description="Image"), | |
| model: str = Input( | |
| description="Style", | |
| default='Hayao:v2', | |
| choices=[ | |
| 'Hayao', | |
| 'Shinkai', | |
| 'Hayao:v2' | |
| ] | |
| ) | |
| ) -> Path: | |
| version = model.split(":")[-1] | |
| predictor = MyPredictor(model, version) | |
| print(f"image path: {str(image)}") | |
| img = read_image(str(image)) | |
| anime_img = predictor.transform(resize_image(img))[0] | |
| out_path = Path(tempfile.mkdtemp()) / "out.png" | |
| cv2.imwrite(str(out_path), anime_img[..., ::-1]) | |
| return out_path | |