Auto deploy

stCompressService.py

@@ -1,8 +1,11 @@
+import os
 import pathlib
 import torch
 import torch.hub
 from torchvision.transforms.functional import convert_image_dtype
 from torchvision.io.image import ImageReadMode, encode_png, decode_image
+from PIL import Image
+import PIL
 
 from mcquic import Config
 from mcquic.modules.compressor import BaseCompressor, Compressor
@@ -18,6 +21,8 @@ except:
 
 MODELS_URL = "https://github.com/xiaosu-zhu/McQuic/releases/download/generic/qp_3_msssim_fcc58b73.mcquic"
 
+HF_SPACE = "HF_SPACE" in os.environ
+
 
 @st.experimental_singleton
 def loadModel(qp: int, local: pathlib.Path, device, mse: bool):
@@ -60,7 +65,7 @@ def decompressImage(sourceFile: File, model: BaseCompressor) -> torch.ByteTensor
 
 
 
-def main(debug: bool, quiet: bool, disable_gpu: bool):
+def main(debug: bool, quiet: bool, qp: int, disable_gpu: bool):
     if disable_gpu or not torch.cuda.is_available():
         device = torch.device("cpu")
     else:
@@ -84,23 +89,38 @@ def main(debug: bool, quiet: bool, disable_gpu: bool):
 </p>
 
 
-<a href="#">
-  <image src="https://img.shields.io/badge/NOTE-yellow?style=for-the-badge" alt="NOTE"/>
-</a>
+<image src="https://img.shields.io/badge/NOTE-yellow?style=for-the-badge" alt="NOTE"/>
 
 > Due to resources limitation, I only provide compression service with model `qp = 3`.
 
+<br/>
+<br/>
+<br/>
+<br/>
+<br/>
+<br/>
+<br/>
 
-<a href="#">
+<a href="https://github.com/xiaosu-zhu/McQuic">
+  <image src="https://raw.githubusercontent.com/xiaosu-zhu/McQuic/main/assets/GitHub-Mark-120px-plus.png" alt="Github"/>
   <image src="https://img.shields.io/github/stars/xiaosu-zhu/McQuic?style=social" alt="Github"/>
 </a>
 
 """, unsafe_allow_html=True)
 
+    if HF_SPACE:
+        st.markdown("""
+<image src="https://img.shields.io/badge/NOTE-yellow?style=for-the-badge" alt="NOTE"/>
+
+> Due to resources limitation of HF spaces, the upload image size is restricted to lower than `3000 x 3000`.
+
+<image src="https://img.shields.io/badge/NOTE-yellow?style=for-the-badge" alt="NOTE"/>
 
+> Also, this demo running on HF space is GPU-disabled. So it may be slow.
+""")
     with st.form("SubmitForm"):
         uploadedFile = st.file_uploader("Try running McQuic to compress or restore images!", type=["png", "jpg", "jpeg", "mcq"], help="Upload your image or compressed `.mcq` file here.")
-        cropping = st.checkbox("Cropping image to align grids.", help="If checked, the image is cropped to align to feature map grids. This makes output smaller.")
+        cropping = st.checkbox("Cropping image to align grids.", help="If checked, the image is cropped to align feature map grids. This will make compressed file smaller.")
         submitted = st.form_submit_button("Submit", help="Click to start compress/restore.")
     if submitted and uploadedFile is not None:
         if uploadedFile.name.endswith(".mcq"):
@@ -114,16 +134,66 @@ def main(debug: bool, quiet: bool, disable_gpu: bool):
             st.image(result.cpu().permute(1, 2, 0).numpy())
             st.download_button("Click to download restored image", data=bytes(encode_png(result.cpu()).tolist()), file_name=".".join(uploadedFile.name.split(".")[:-1] + ["png"]), mime="image/png")
         else:
+            try:
+                a = Image.open(uploadedFile)
+            except PIL.UnidentifiedImageError:
+                st.markdown("""
+<image src="https://img.shields.io/badge/ERROR-red?style=for-the-badge" alt="ERROR"/>
+
+> Image open failed. Please try other images.
+""")
+                return
+            w, h = a.size
+            if HF_SPACE and (h > 3000 or w > 3000):
+                st.markdown("""
+<image src="https://img.shields.io/badge/ERROR-red?style=for-the-badge" alt="ERROR"/>
+
+> Image is too large. Please try other images.
+""")
+                return
             raw = torch.ByteTensor(torch.ByteStorage.from_buffer(uploadedFile.read())) # type: ignore
             image = decode_image(raw, ImageReadMode.RGB).to(device)
-            st.image(image.cpu().permute(1, 2, 0).numpy())
+            # st.image(image.cpu().permute(1, 2, 0).numpy())
             result = compressImage(image, model, cropping)
 
             st.text(str(result))
 
             st.download_button("Click to download compressed file", data=result.serialize(), file_name=".".join(uploadedFile.name.split(".")[:-1] + ["mcq"]), mime="image/mcq")
 
+    st.markdown("""
+<br/>
+<br/>
+<br/>
+<br/>
+<br/>
+
+
+<p align="center">
+  <a href="https://www.python.org/" target="_blank">
+    <image src="https://img.shields.io/badge/python-3670A0?style=for-the-badge&logo=python&logoColor=ffdd54" alt="Python"/>
+  </a>
+  <a href="https://pytorch.org/" target="_blank">
+    <image src="https://img.shields.io/badge/PyTorch-%23EE4C2C.svg?style=for-the-badge&logo=PyTorch&logoColor=white" alt="PyTorch"/>
+  </a>
+  <a href="https://github.com/xiaosu-zhu/McQuic/stargazers" target="_blank">
+    <image src="https://img.shields.io/github/stars/xiaosu-zhu/McQuic?logo=github&style=for-the-badge" alt="Github stars"/>
+  </a>
+  <a href="https://github.com/xiaosu-zhu/McQuic/network/members" target="_blank">
+    <image src="https://img.shields.io/github/forks/xiaosu-zhu/McQuic?logo=github&style=for-the-badge" alt="Github forks"/>
+  </a>
+  <a href="https://github.com/xiaosu-zhu/McQuic/blob/main/LICENSE" target="_blank">
+    <image src="https://img.shields.io/github/license/xiaosu-zhu/McQuic?logo=github&style=for-the-badge" alt="Github license"/>
+  </a>
+</p>
+
+<br/>
+<br/>
+<br/>
+
+<p align="center"><a href="localhost" target="_blank">CVF Open Access</a> | <a href="localhost" target="_blank">arXiv</a> | <a href="https://github.com/xiaosu-zhu/McQuic#citation" target="_blank">BibTex</a> | <a href="https://huggingface.co/spaces/xiaosu-zhu/McQuic">Demo</a></p>
+
+""")
 
 if __name__ == "__main__":
     with torch.inference_mode():
-        main(False, False, False)
+        main(False, False, 3, False)