Fix inference + scaling, update Gradio

LICENSE

@@ -0,0 +1,236 @@
+MIT License
+
+Copyright (c) 2021 Junlin Han
+Copyright (c) 2024 Yarflam on Hugging Face (same license)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+--------------------------- LICENSE FOR CUT -------------------------------
+Copyright (c) 2020, Taesung Park and Jun-Yan Zhu
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+--------------------------- LICENSE FOR CycleGAN -------------------------------
+-------------------https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix------
+Copyright (c) 2017, Jun-Yan Zhu and Taesung Park
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+--------------------------- LICENSE FOR stylegan2-pytorch ----------------------
+----------------https://github.com/rosinality/stylegan2-pytorch/----------------
+MIT License
+
+Copyright (c) 2019 Kim Seonghyeon
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+--------------------------- LICENSE FOR pix2pix --------------------------------
+BSD License
+
+For pix2pix software
+Copyright (c) 2016, Phillip Isola and Jun-Yan Zhu
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+----------------------------- LICENSE FOR DCGAN --------------------------------
+BSD License
+
+For dcgan.torch software
+
+Copyright (c) 2015, Facebook, Inc. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+Neither the name Facebook nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+--------------------------- LICENSE FOR StyleGAN2 ------------------------------
+--------------------------- Inherited from stylegan2-pytorch -------------------
+Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+
+
+Nvidia Source Code License-NC
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+"Nvidia Processors" means any central processing unit (CPU), graphics
+processing unit (GPU), field-programmable gate array (FPGA),
+application-specific integrated circuit (ASIC) or any combination
+thereof designed, made, sold, or provided by Nvidia or its affiliates.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+        License, each Licensor grants to you a perpetual, worldwide,
+	    non-exclusive, royalty-free, copyright license to reproduce,
+	        prepare derivative works of, publicly display, publicly perform,
+		    sublicense and distribute its Work and any resulting derivative
+		        works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+        if (a) you do so under this License, (b) you include a complete
+	    copy of this License with your distribution, and (c) you retain
+	        without modification any copyright, patent, trademark, or
+		    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+        terms apply to the use, reproduction, and distribution of your
+	    derivative works of the Work ("Your Terms") only if (a) Your Terms
+	        provide that the use limitation in Section 3.3 applies to your
+		    derivative works, and (b) you identify the specific derivative
+		        works that are subject to Your Terms. Notwithstanding Your Terms,
+			    this License (including the redistribution requirements in Section
+			        3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+        may be used or intended for use non-commercially. The Work or
+	    derivative works thereof may be used or intended for use by Nvidia
+	        or its affiliates commercially or non-commercially. As used herein,
+		    "non-commercially" means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+        against any Licensor (including any claim, cross-claim or
+	    counterclaim in a lawsuit) to enforce any patents that you allege
+	        are infringed by any Work, then your rights under this License from
+		    such Licensor (including the grants in Sections 2.1 and 2.2) will
+		        terminate immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+        Licensor's or its affiliates' names, logos, or trademarks, except
+	    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+        rights under this License (including the grants in Sections 2.1 and
+	    2.2) will terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

ModelLoader.py

@@ -1,5 +1,6 @@
 from models import create_model
 from util.get_transform import get_transform
+from util.util import tensor2im
 from PIL import Image
 import os
 
@@ -14,41 +15,43 @@ class Options(object):
             setattr(self, key, kwargs[key])
 
 class ModelLoader:
-    def __init__(self) -> None:
+    def __init__(self, gpu_ids='', max_img_wh=512) -> None:
         self.opt = Options({
             'isGradio': True, # Custom
-            'name': 'original',
-            'checkpoints_dir': ckp_path,
-            'gpu_ids': [],
-            'init_gain': 0.02,
-            'init_type': 'xavier',
-            'input_nc': 3,
+            'name': 'original', # Checkpoints name
+            'checkpoints_dir': ckp_path, # Checkpoint folder
+            'gpu_ids': gpu_ids.split(',') if gpu_ids else [],
+            'init_gain': 0.02, # Scaling Factor
+            'init_type': 'xavier', # list: 'normal', 'xavier', 'kaiming', 'orthogonal'
+            'input_nc': 3, # 3 -> RGB, 1 -> Grayscale
             'output_nc': 3,
             'isTrain': False,
             'model': 'cwr',
             'nce_idt': False,
             'nce_layers': '0',
-            'ndf': 64,
+            'ndf': 64, # Nb of discrim filters in the first conv layer
             'netD': 'basic',
             'netG': 'resnet_9blocks',
             'netF': 'mlp_sample',
             'netF_nc': 256,
-            'ngf': 64,
-            'no_antialias_up': None,
-            'no_antialias': None,
+            'ngf': 64, # Nb of gen filters in the last conv layer
+            'no_antialias_up': False,
+            'no_antialias': False,
             'no_dropout': True,
             'normD': 'instance',
             'normG': 'instance',
-            'preprocess': 'scale_width',
-            'num_threads': 0, # test code only supports num_threads = 1
+            'preprocess': 'yarflam_auto', # see more: util.get_transform
+            'dataroot': 'placeholder',
+            'num_threads': 1, # test code only supports num_threads = 1
             'batch_size': 1, # test code only supports batch_size = 1
-            'serial_batches': True, # disable data shuffling; comment this line if results on randomly chosen images are needed.
+            'serial_batches': False, # disable data shuffling; comment this line if results on randomly chosen images are needed.
             'no_flip': True, # no flip; comment this line if results on flipped images are needed.
             'display_id': -1, # no visdom display; the test code saves the results to a HTML file.
             'direction': 'AtoB', # inference
             'flip_equivariance': False,
-            'load_size': 1680,
-            'crop_size': 512,
+            'load_size': 1680, # not used
+            'crop_size': 512, # not used
+            'yarflam_img_wh': max_img_wh, # max width|height + auto scale down
         })
         self.transform = get_transform(self.opt, grayscale=False)
         self.model = None
@@ -62,9 +65,14 @@ class ModelLoader:
         # Loading
         print('Loading the image %s' % src)
         source = Image.open(src).convert('RGB')
-        img = self.transform(source)
+        img = self.transform(source).unsqueeze(0)
         print(img.shape)
         # Inference
-        self.model.set_input({ 'A': img, 'B': img, 'A_paths': src })
+        self.model.set_input({
+            'A': img, 'A_paths': src,
+            'B': img, 'B_paths': src
+        })
         self.model.forward()
-        print(self.model)
+        out_data = list(self.model.get_current_visuals().items())[1][1]
+        out_img = Image.fromarray(tensor2im(out_data))
+        return out_img

README.md

@@ -10,4 +10,18 @@ pinned: false
 short_description: Contrastive UnderWater Restoration
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# UnderWater
+
+Improve underwater photos.
+
+Contrastive UnderWater Restoration (CWR) - Gradio integration.
+
+Source: https://github.com/JunlinHan/CWR
+
+## Author
+
+- Yarflam - Gradio integration
+
+## License
+
+MIT License - [see more](LICENSE)

app.py

@@ -1,11 +1,24 @@
 import gradio as gr
+from ModelLoader import ModelLoader
+import os
 
-def greet(name):
-    return "Hello " + name + "!!"
+max_img_wh = 1024  # Set the maximum image size in pixels
+
+def inference(image, use_gpu):
+    gpu_ids = '0' if use_gpu else ''
+    model = ModelLoader(gpu_ids=gpu_ids, max_img_wh=max_img_wh)
+    model.load()
+    output_img = model.inference(src=image)
+    return output_img
 
 demo = gr.Interface(
-    fn=greet,
-    inputs="text",
-    outputs="text"
+    fn=inference,
+    inputs=[
+        gr.inputs.Image(type="pil", label="Input Image"),
+        gr.inputs.Checkbox(label="Use GPU", value=True)  # Precheck the GPU checkbox
+    ],
+    outputs=gr.outputs.Image(type="pil", label="Output Image"),
+    title="Model Demo",
+    description="Upload an image to see the model output."
 )
 demo.launch()

test.py

@@ -2,10 +2,12 @@ from ModelLoader import ModelLoader
 import os
 
 def main():
-    model = ModelLoader()
+    model = ModelLoader(gpu_ids='', max_img_wh=512)
     model.load()
     # Test
     sample = os.path.join(os.path.dirname(__file__), 'examples', 'rawimg.png')
-    model.inference(src=sample)
+    # sample = os.path.join(os.path.dirname(__file__), 'examples', 'test.jpg')
+    img = model.inference(src=sample)
+    img.show() # PIL Image
 
 main()

util/get_transform.py

@@ -18,6 +18,9 @@ def get_transform(opt, params=None, grayscale=False, method=Image.BICUBIC, conve
     elif 'scale_shortside' in opt.preprocess:
         transform_list.append(transforms.Lambda(lambda img: __scale_shortside(img, opt.load_size, opt.crop_size, method)))
 
+    if opt.preprocess == 'yarflam_auto':
+        transform_list.append(transforms.Lambda(lambda img: __scale_yarflam(img, opt.yarflam_img_wh, method)))
+
     if 'zoom' in opt.preprocess:
         if params is None:
             transform_list.append(transforms.Lambda(lambda img: __random_zoom(img, opt.load_size, opt.crop_size, method)))
@@ -110,6 +113,19 @@ def __scale_width(img, target_width, crop_width, method=Image.BICUBIC):
     h = int(max(target_width * oh / ow, crop_width))
     return img.resize((w, h), method)
 
+def __scale_yarflam(img, target_wh, method=Image.BICUBIC):
+    ow, oh = img.size
+    if max(ow, oh) <= target_wh:
+        return img
+    if ow > target_wh and oh > target_wh:
+        ratio = target_wh / max(ow, oh)
+        w, h = int(ow * ratio), int(oh * ratio)
+    elif ow > target_wh:
+        w, h = target_wh, int((oh / ow) * target_wh)
+    else:
+        w, h = int((ow / oh) * target_wh), target_wh
+    return img.resize((w, h), method)
+
 
 def __crop(img, pos, size):
     ow, oh = img.size