Update app.py

app.py

@@ -1,40 +1,37 @@
-import os
-
+import os, sys
+import argparse
 import cv2
 import gradio as gr
 import torch
 from basicsr.archs.srvgg_arch import SRVGGNetCompact
 from realesrgan.utils import RealESRGANer
+from glob import glob
+
+sys.path.insert(1, os.path.join(sys.path[0], '..'))
+from gradio_demo.RestoreFormer import RestoreFormer
+
+if not os.path.exists('experiments/pretrained_models'):
+    os.makedirs('experiments/pretrained_models')
+realesr_model_path = 'experiments/pretrained_models/RealESRGAN_x4plus.pth'
+if not os.path.exists(realesr_model_path):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -O experiments/pretrained_models/RealESRGAN_x4plus.pth")
+
+if not os.path.exists('experiments/RestoreFormer/'):
+    os.makedirs('experiments/RestoreFormer/')
+restoreformer_model_path = 'experiments/RestoreFormer/last.ckpt'
+if not os.path.exists(restoreformer_model_path):
+    os.system("wget https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer.ckpt -O experiments/RestoreFormer/last.ckpt")
 
-from RestoreFormer import RestoreFormer
-
-os.system("pip freeze")
-# download weights
-if not os.path.exists('realesr-general-x4v3.pth'):
-    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
-if not os.path.exists('RestoreFormer.ckpt'):
-    os.system("wget https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer.ckpt -P .")
-if not os.path.exists('RestoreFormer++.ckpt'):
-    os.system("wget https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer++.ckpt -P .")
-
-# torch.hub.download_url_to_file(
-#     'https://upload.wikimedia.org/wikipedia/commons/thumb/a/ab/Abraham_Lincoln_O-77_matte_collodion_print.jpg/1024px-Abraham_Lincoln_O-77_matte_collodion_print.jpg',
-#     'lincoln.jpg')
-# torch.hub.download_url_to_file(
-#     'https://user-images.githubusercontent.com/17445847/187400315-87a90ac9-d231-45d6-b377-38702bd1838f.jpg',
-#     'AI-generate.jpg')
-# torch.hub.download_url_to_file(
-#     'https://user-images.githubusercontent.com/17445847/187400981-8a58f7a4-ef61-42d9-af80-bc6234cef860.jpg',
-#     'Blake_Lively.jpg')
-# torch.hub.download_url_to_file(
-#     'https://user-images.githubusercontent.com/17445847/187401133-8a3bf269-5b4d-4432-b2f0-6d26ee1d3307.png',
-#     '10045.png')
+if not os.path.exists('experiments/RestoreFormerPlusPlus/'):
+    os.makedirs('experiments/RestoreFormerPlusPlus/')
+restoreformerplusplus_model_path = 'experiments/RestoreFormerPlusPlus/last.ckpt'
+if not os.path.exists(restoreformerplusplus_model_path):
+    os.system("wget https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer++.ckpt -O experiments/RestoreFormerPlusPlus/last.ckpt")
 
 # background enhancer with RealESRGAN
 model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-model_path = 'realesr-general-x4v3.pth'
 half = True if torch.cuda.is_available() else False
-upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+upsampler = RealESRGANer(scale=4, model_path=realesr_model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
 
 os.makedirs('output', exist_ok=True)
 
@@ -66,10 +63,10 @@ def inference(img, version, scale):
 
         if version == 'RestoreFormer':
             face_enhancer = RestoreFormer(
-            model_path='RestoreFormer.ckpt', upscale=2, arch='RestoreFormer', bg_upsampler=upsampler)
+            model_path=restoreformer_model_path, upscale=2, arch='RestoreFormer', bg_upsampler=upsampler)
         elif version == 'RestoreFormer++':
             face_enhancer = RestoreFormer(
-            model_path='RestoreFormer++.ckpt', upscale=2, arch='RestoreFormer++', bg_upsampler=upsampler)
+            model_path=restoreformerplusplus_model_path, upscale=2, arch='RestoreFormer++', bg_upsampler=upsampler)
 
         try:
             # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
@@ -98,16 +95,90 @@ def inference(img, version, scale):
         return None, None
 
 
-title = "RestoreFormer: Blind Face Restoration Algorithm"
-description = r"""Gradio demo for <a href='https://github.com/wzhouxiff/RestoreFormerPlusPlus' target='_blank'><b>RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Paris</b></a>.<br>
-It is used to restore your **old photos**.<br>
+title = "RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Paris"
+important_links=r'''
+<div align='center'>
+[![paper_RestroeForemer++](https://img.shields.io/badge/TPAMI-Restorformer%2B%2B-green
+)](https://openaccess.thecvf.com/content/CVPR2022/papers/Wang_RestoreFormer_High-Quality_Blind_Face_Restoration_From_Undegraded_Key-Value_Pairs_CVPR_2022_paper.pdf)
+&nbsp; 
+[![paere_RestroeForemer](https://img.shields.io/badge/CVPR22-Restorformer-green)](https://openaccess.thecvf.com/content/CVPR2022/papers/Wang_RestoreFormer_High-Quality_Blind_Face_Restoration_From_Undegraded_Key-Value_Pairs_CVPR_2022_paper.pdf)
+&nbsp;
+[![code_RestroeForemer++](https://img.shields.io/badge/GitHub-RestoreFormer%2B%2B-red
+
+)](https://github.com/wzhouxiff/RestoreFormerPlusPlus)
+&nbsp; 
+[![code_RestroeForemer](https://img.shields.io/badge/GitHub-RestoreFormer-red)](https://github.com/wzhouxiff/RestoreFormer)
+&nbsp;
+[![demo](https://img.shields.io/badge/Demo-Gradio-orange
+)](https://gradio.app/hub/wzhouxiff/RestoreFormerPlusPlus)
+</div>
+'''
+description = r"""
+<div align='center'>
+<a target='_blank' href='https://arxiv.org/pdf/2308.07228.pdf' style='float: left'>
+<img src='https://img.shields.io/badge/TPAMI-RestorFormer%2B%2B-green' alt='paper_RestroeForemer++'>
+</a>
+&ensp;&ensp;&ensp;&ensp;&ensp;&ensp;
+<a target='_blank' href='https://openaccess.thecvf.com/content/CVPR2022/papers/Wang_RestoreFormer_High-Quality_Blind_Face_Restoration_From_Undegraded_Key-Value_Pairs_CVPR_2022_paper.pdf' style='float: left'>
+<img src='https://img.shields.io/badge/CVPR22-RestorFormer-green' alt='paere_RestroeForemer' >
+</a>
+&ensp;&ensp;&ensp;&ensp;&ensp;&ensp;
+<a target='_blank' href='https://github.com/wzhouxiff/RestoreFormerPlusPlus' style='float: left'>
+<img src='https://img.shields.io/badge/GitHub-RestoreFormer%2B%2B-red' alt='code_RestroeForemer++'>
+</a>
+&ensp;&ensp;&ensp;&ensp;&ensp;&ensp;
+<a target='_blank' href='https://github.com/wzhouxiff/RestoreFormer' style='float: left'>
+<img src='https://img.shields.io/badge/GitHub-RestoreFormer-red' alt='code_RestroeForemer' >
+</a>
+&ensp;&ensp;&ensp;&ensp;&ensp;&ensp;
+<a target='_blank' href='https://huggingface.co/spaces/wzhouxiff/RestoreFormerPlusPlus' style='float: left' >
+<img src='https://img.shields.io/badge/Demo-Gradio-orange' alt='demo' >
+</a>
+&ensp;&ensp;&ensp;&ensp;&ensp;&ensp;
+</div>
+<br>
+Gradio demo for <a href='https://github.com/wzhouxiff/RestoreFormerPlusPlus' target='_blank'><b>RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Paris</b></a>.
+<br>
+It is used to restore your Old Photos.
+<br>
 To use it, simply upload your image.<br>
 """
+
 article = r"""
-[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/pdf/2308.07228.pdf)
-[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://openaccess.thecvf.com/content/CVPR2022/papers/Wang_RestoreFormer_High-Quality_Blind_Face_Restoration_From_Undegraded_Key-Value_Pairs_CVPR_2022_paper.pdf)
+If the proposed algorithm is helpful, please help to ⭐ the GitHub Repositories: <a href='https://github.com/wzhouxiff/RestoreFormer' target='_blank'>RestoreFormer</a> and
+<a href='https://github.com/wzhouxiff/RestoreFormerPlusPlus' target='_blank'>RestoreFormer++</a>. Thanks! 
+[![GitHub Stars](https://img.shields.io/github/stars/wzhouxiff%2FRestoreFormer
+)](https://github.com/wzhouxiff/RestoreFormer)
+[![GitHub Stars](https://img.shields.io/github/stars/wzhouxiff%2FRestoreFormerPlusPlus
+)](https://github.com/wzhouxiff/RestoreFormerPlusPlus)
+
+---
+
+📝 **Citation**
+<br>
+If our work is useful for your research, please consider citing:
+```bibtex
+@article{wang2023restoreformer++,
+    title={RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Paris},
+    author={Wang, Zhouxia and Zhang, Jiawei and Chen, Tianshui and Wang, Wenping and Luo, Ping},
+    booktitle={IEEE Transactions on Pattern Analysis and Machine Intelligence (T-PAMI)},
+    year={2023}
+}
+@article{wang2022restoreformer,
+    title={RestoreFormer: High-Quality Blind Face Restoration from Undegraded Key-Value Pairs},
+    author={Wang, Zhouxia and Zhang, Jiawei and Chen, Runjian and Wang, Wenping and Luo, Ping},
+    booktitle={The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
+    year={2022}
+}
+```
+
 If you have any question, please email 📧 `wzhoux@connect.hku.hk`.
 """
+
+css=r"""
+
+"""
+
 demo = gr.Interface(
     inference, [
         gr.Image(type="filepath", label="Input"),
@@ -120,9 +191,6 @@ demo = gr.Interface(
     title=title,
     description=description,
     article=article,
-    # examples=[['AI-generate.jpg', 'v1.4', 2, 50], ['lincoln.jpg', 'v1.4', 2, 50], ['Blake_Lively.jpg', 'v1.4', 2, 50],
-    #           ['10045.png', 'v1.4', 2, 50]]).launch()
-    # examples=[['AI-generate.jpg', 'v1.4', 2], ['lincoln.jpg', 'v1.4', 2], ['Blake_Lively.jpg', 'v1.4', 2],
-    #           ['10045.png', 'v1.4', 2]]
     )
+
 demo.queue().launch()