Update app.py

app.py

@@ -1,11 +1,15 @@
 from PIL import Image
 from io import BytesIO
+from matplotlib.figure import Figure
 from torchvision import transforms
+from tqdm import tqdm
 from typing import Literal, Any
+from urllib.request import urlopen
 import gradio as gr
-from matplotlib.figure import Figure
 import matplotlib.pyplot as plt
+import os
 import spaces
+import sys
 import torch
 import torch.nn.functional as F
 
@@ -19,11 +23,28 @@ LABELS = [
     "Ephemeral",
     "Canopied",
 ]
+MODELFILE = "Litton-7type-visual-landscape-model.pth"
+
 
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
+if not os.path.exists(MODELFILE):
+    model_url = f"https://lclab.thu.edu.tw/modelzoo/{MODELFILE}"
+
+    print(f"fetch model from {model_url}...", file=sys.stderr)
+
+    with urlopen(model_url) as resp:
+        progress = tqdm(total=int(resp["Content-Length"]), desc="Downloading")
+        with open(MODELFILE, "wb") as modelfile:
+            while True:
+                chunk = resp.read(1024)
+                if len(chunk) == 0:
+                    break
+                modelfile.write(chunk)
+                progress.update(len(chunk))
+
 model = torch.load(
-    "Litton-7type-visual-landscape-model.pth", map_location=device, weights_only=False
+    MODELFILE, map_location=device, weights_only=False
 ).module
 model.eval()
 preprocess = transforms.Compose(
@@ -73,6 +94,14 @@ def draw_bar_chart(data: dict[str, list[str | float]]):
     return fig
 
 
+def choose_example(imgpath: str) -> gr.Image:
+    img = Image.open(imgpath)
+    width, height = img.size
+    ratio = 512 / max(width, height)
+    img = img.resize((int(width * ratio), int(height * ratio)))
+    return gr.Image(value=img, label="輸入影像（不支援 SVG 格式）", type="pil")
+
+
 def get_layout():
     css = """
     .main-title {
@@ -106,6 +135,10 @@ def get_layout():
         color: #d1d5db;
         font-size: 14px;
     }
+    .example-image {
+        height: 220px;
+        padding: 25px;
+    }
     """
     theme = gr.themes.Base(
         primary_hue="orange",
@@ -153,22 +186,56 @@ def get_layout():
             )
 
             with gr.Row(equal_height=True):
-                image_input = gr.Image(label="上傳影像", type="pil")
+                with gr.Group():
+                    img = gr.Image(label="上傳影像", type="pil", height="256px")
+                    gr.Label("範例影像", show_label=False)
+                    with gr.Row():
+                        ex1 = gr.Image(
+                            value="examples/beach.jpg",
+                            show_label=False,
+                            type="filepath",
+                            elem_classes="example-image",
+                            interactive=False,
+                            show_download_button=False,
+                            show_fullscreen_button=False,
+                        )
+                        ex2 = gr.Image(
+                            value="examples/field.jpg",
+                            show_label=False,
+                            type="filepath",
+                            elem_classes="example-image",
+                            interactive=False,
+                            show_download_button=False,
+                            show_fullscreen_button=False,
+                        )
+                        ex3 = gr.Image(
+                            value="examples/sky.jpg",
+                            show_label=False,
+                            type="filepath",
+                            elem_classes="example-image",
+                            interactive=False,
+                            show_download_button=False,
+                            show_fullscreen_button=False,
+                        )
                 chart = gr.Plot(label="分類結果")
 
-            start_button = gr.Button("開始分類", variant="primary")
+            start_button = gr.Button("開始", variant="primary")
             gr.HTML(
                 '<div class="footer">© 2024 LCL 版權所有<br>開發者：何立智、楊哲睿</div>',
             )
 
-            start_button.click(
-                fn=predict,
-                inputs=image_input,
-                outputs=chart,
-            )
+        start_button.click(
+            fn=predict,
+            inputs=img,
+            outputs=chart,
+        )
+
+        ex1.select(fn=choose_example, inputs=ex1, outputs=img)
+        ex2.select(fn=choose_example, inputs=ex2, outputs=img)
+        ex3.select(fn=choose_example, inputs=ex3, outputs=img)
 
     return demo
 
 
 if __name__ == "__main__":
-    get_layout().queue().launch()
+    get_layout().launch()