cog-ified

Author: ericguizzo

README.md

@@ -5,6 +5,7 @@
 [![Build Status](https://travis-ci.org/idealo/image-super-resolution.svg?branch=master)](https://travis-ci.org/idealo/image-super-resolution)
 [![Docs](https://img.shields.io/badge/docs-online-brightgreen)](https://idealo.github.io/image-super-resolution/)
 [![License](https://img.shields.io/badge/License-Apache%202.0-orange.svg)](https://github.com/idealo/image-super-resolution/blob/master/LICENSE)
+<a href="https://replicate.ai/idealo/image-super-resolution"><img src="https://img.shields.io/static/v1?label=Replicate&message=Demo and Docker Image&color=darkgreen" height=20></a>
 
 The goal of this project is to upscale and improve the quality of low resolution images.
 
@@ -65,18 +66,18 @@ pip install 'h5py==2.10.0' --force-reinstall
 
 ## Pre-trained networks
 
-The weights used to produced these images are available directly when creating the model object. 
+The weights used to produced these images are available directly when creating the model object.
 
 Currently 4 models are available:
   - RDN: psnr-large, psnr-small, noise-cancel
   - RRDN: gans
- 
+
 Example usage:
 
 ```
 model = RRDN(weights='gans')
 ```
-  
+
 The network parameters will be automatically chosen.
 (see [Additional Information](#additional-information)).
 

cog.yaml

@@ -0,0 +1,8 @@
+build:
+  python_version: "3.7"
+  gpu: false
+  python_packages:
+    - ISR==2.2.0
+    - h5py==2.10.0 --force-reinstall
+
+predict: "predict.py:ISRPredictor"

predict.py

@@ -0,0 +1,43 @@
+import os
+import tempfile
+from pathlib import Path
+
+import cog
+import numpy as np
+from ISR.models import RDN, RRDN
+from PIL import Image
+
+
+class ISRPredictor(cog.Predictor):
+    def setup(self):
+        """Load the super-resolution ans noise canceling models"""
+        self.model_gans = RRDN(weights="gans")
+        self.model_noise_cancel = RDN(weights="noise-cancel")
+
+    @cog.input("input", type=Path, help="Image path")
+    @cog.input(
+        "type",
+        type=str,
+        default="super-resolution",
+        options=["super-resolution", "noise-cancel"],
+        help="Precessing type: super-resolution or noise-cancel",
+    )
+    def predict(self, input, type):
+        """Apply super-resolution or noise-canceling to input image"""
+        # compute super resolution
+        img = Image.open(str(input))
+        lr_img = np.array(img)
+
+        if type == "super-resolution":
+            img = self.model_gans.predict(np.array(img))
+        elif type == "noise-cancel":
+            img = self.model_noise_cancel.predict(np.array(img))
+        else:
+            raise NotImplementedError("Invalid processing type selected")
+
+        img = Image.fromarray(img)
+
+        output_path = Path(tempfile.mkdtemp()) / "output.png"
+        img.save(str(output_path), "PNG")
+
+        return output_path