Update app.py

app.py

@@ -1,159 +1,162 @@
-import matplotlib.pyplot as plt
-import torch
-from PIL import Image
-from torchvision import transforms
-import torch.nn.functional as F
-from typing import Literal, Any
-import gradio as gr
-
-
-class Classifier:
-    LABELS = [
-        "Panoramic",
-        "Feature",
-        "Detail",
-        "Enclosed",
-        "Focal",
-        "Ephemeral",
-        "Canopied",
-    ]
-
-    def __init__(
-        self, model_path="Litton-7type-visual-landscape-model.pth", device="cuda:0"
-    ):
-        self.device = device
-        self.model = torch.load(
-            model_path, map_location=self.device, weights_only=False
-        )
-        if hasattr(self.model, "module"):
-            self.model = self.model.module
-        self.model.eval()
-        self.preprocess = transforms.Compose(
-            [
-                transforms.Resize(256),
-                transforms.CenterCrop(224),
-                transforms.ToTensor(),
-                transforms.Normalize(
-                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-                ),
-            ]
-        )
-
-    def predict(self, image: Image.Image) -> tuple[Literal["Failed", "Success"], Any]:
-        image = image.convert("RGB")
-        input_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            logits = self.model(input_tensor)
-            probs = F.softmax(logits[:, :7], dim=1).cpu()
-
-        # probs = pd.DataFrame(
-        #     {
-        #         "class": self.LABELS,
-        #         "probs": probs[0] * 100,
-        #     }
-        # )
-        return draw_bar_chart(
-            {
-                "class": self.LABELS,
-                "probs": probs[0] * 100,
-            }
-        )
-
-
-def draw_bar_chart(data: dict[str, list[str | float]]):
-    classes = data["class"]
-    probabilities = data["probs"]
-
-    plt.figure(figsize=(8, 6))
-    plt.bar(classes, probabilities, color="skyblue")
-
-    plt.xlabel("Class")
-    plt.ylabel("Probability (%)")
-    plt.title("Class Probabilities")
-
-    for i, prob in enumerate(probabilities):
-        plt.text(i, prob + 0.01, f"{prob:.2f}", ha="center", va="bottom")
-
-    plt.tight_layout()
-
-    return plt
-
-
-def get_layout():
-    css = """
-   .main-title {
-        font-size: 24px;
-        font-weight: bold;
-        text-align: center;
-        margin-bottom: 20px;
-    }
-    .reference {
-        text-align: center;
-        font-size: 1.2em;
-        color: #d1d5db;
-        margin-bottom: 20px;
-    }
-    .reference a {
-        color: #FB923C;
-        text-decoration: none;
-    }
-    .reference a:hover {
-        text-decoration: underline;
-        color: #FB923C;
-    }
-    .title {
-        border-bottom: 1px solid;
-    }
-    .footer {
-        text-align: center;
-        margin-top: 30px;
-        padding-top: 20px;
-        border-top: 1px solid #ddd;
-        color: #d1d5db;
-        font-size: 14px;
-    }
-    """
-    theme = gr.themes.Base(
-        primary_hue="orange",
-        secondary_hue="orange",
-        neutral_hue="gray",
-        font=gr.themes.GoogleFont("Source Sans Pro"),
-    ).set(
-        background_fill_primary="*neutral_950",  # 主背景色（深黑）
-        button_primary_background_fill="*primary_500",  # 按鈕顏色（橘色）
-        body_text_color="*neutral_200",  # 文字顏色（淺色）
-    )
-    with gr.Blocks(css=css, theme=theme) as demo:
-        gr.HTML(
-            value=(
-                '<div class="main-title">Litton7景觀分類模型</div>'
-                '<div class="reference">引用資料：'
-                '<a href="https://www.airitilibrary.com/Article/Detail/10125434-N202406210003-00003" target="_blank">'
-                "何立智、李沁築、邱浩修(2024)。Litton7：Litton視覺景觀分類深度學習模型。戶外遊憩研究，37(2)"
-                "</a>"
-                "</div>"
-            ),
-        )
-
-        with gr.Row():
-            image_input = gr.Image(label="上傳影像", type="pil")
-            bar_chart = gr.Plot(label="分類結果")
-        start_button = gr.Button("開始分類", variant="primary")
-
-        start_button.click(
-            fn=Classifier().predict,
-            inputs=image_input,
-            outputs=bar_chart,
-        )
-
-        gr.HTML(
-            '<div class="footer">© 2024 LCL 版權所有<br>開發者：何立智、楊哲睿</div>',
-        )
-
-    return demo
-
-
-if __name__ == "__main__":
-    app = get_layout()
-    app.queue().launch(server_name="0.0.0.0")
+import matplotlib.pyplot as plt
+import torch
+from PIL import Image
+from torchvision import transforms
+import torch.nn.functional as F
+from typing import Literal, Any
+import gradio as gr
+import space
+
+
+class Classifier:
+    LABELS = [
+        "Panoramic",
+        "Feature",
+        "Detail",
+        "Enclosed",
+        "Focal",
+        "Ephemeral",
+        "Canopied",
+    ]
+
+    @space.GPU(duration=60)
+    def __init__(
+        self, model_path="Litton-7type-visual-landscape-model.pth", device="cuda:0"
+    ):
+        self.device = device
+        self.model = torch.load(
+            model_path, map_location=self.device, weights_only=False
+        )
+        if hasattr(self.model, "module"):
+            self.model = self.model.module
+        self.model.eval()
+        self.preprocess = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+            ]
+        )
+
+    @space.GPU(duration=60)
+    def predict(self, image: Image.Image) -> tuple[Literal["Failed", "Success"], Any]:
+        image = image.convert("RGB")
+        input_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            logits = self.model(input_tensor)
+            probs = F.softmax(logits[:, :7], dim=1).cpu()
+
+        # probs = pd.DataFrame(
+        #     {
+        #         "class": self.LABELS,
+        #         "probs": probs[0] * 100,
+        #     }
+        # )
+        return draw_bar_chart(
+            {
+                "class": self.LABELS,
+                "probs": probs[0] * 100,
+            }
+        )
+
+
+def draw_bar_chart(data: dict[str, list[str | float]]):
+    classes = data["class"]
+    probabilities = data["probs"]
+
+    plt.figure(figsize=(8, 6))
+    plt.bar(classes, probabilities, color="skyblue")
+
+    plt.xlabel("Class")
+    plt.ylabel("Probability (%)")
+    plt.title("Class Probabilities")
+
+    for i, prob in enumerate(probabilities):
+        plt.text(i, prob + 0.01, f"{prob:.2f}", ha="center", va="bottom")
+
+    plt.tight_layout()
+
+    return plt
+
+
+def get_layout():
+    css = """
+   .main-title {
+        font-size: 24px;
+        font-weight: bold;
+        text-align: center;
+        margin-bottom: 20px;
+    }
+    .reference {
+        text-align: center;
+        font-size: 1.2em;
+        color: #d1d5db;
+        margin-bottom: 20px;
+    }
+    .reference a {
+        color: #FB923C;
+        text-decoration: none;
+    }
+    .reference a:hover {
+        text-decoration: underline;
+        color: #FB923C;
+    }
+    .title {
+        border-bottom: 1px solid;
+    }
+    .footer {
+        text-align: center;
+        margin-top: 30px;
+        padding-top: 20px;
+        border-top: 1px solid #ddd;
+        color: #d1d5db;
+        font-size: 14px;
+    }
+    """
+    theme = gr.themes.Base(
+        primary_hue="orange",
+        secondary_hue="orange",
+        neutral_hue="gray",
+        font=gr.themes.GoogleFont("Source Sans Pro"),
+    ).set(
+        background_fill_primary="*neutral_950",  # 主背景色（深黑）
+        button_primary_background_fill="*primary_500",  # 按鈕顏色（橘色）
+        body_text_color="*neutral_200",  # 文字顏色（淺色）
+    )
+    with gr.Blocks(css=css, theme=theme) as demo:
+        gr.HTML(
+            value=(
+                '<div class="main-title">Litton7景觀分類模型</div>'
+                '<div class="reference">引用資料：'
+                '<a href="https://www.airitilibrary.com/Article/Detail/10125434-N202406210003-00003" target="_blank">'
+                "何立智、李沁築、邱浩修(2024)。Litton7：Litton視覺景觀分類深度學習模型。戶外遊憩研究，37(2)"
+                "</a>"
+                "</div>"
+            ),
+        )
+
+        with gr.Row():
+            image_input = gr.Image(label="上傳影像", type="pil")
+            bar_chart = gr.Plot(label="分類結果")
+        start_button = gr.Button("開始分類", variant="primary")
+
+        start_button.click(
+            fn=Classifier().predict,
+            inputs=image_input,
+            outputs=bar_chart,
+        )
+
+        gr.HTML(
+            '<div class="footer">© 2024 LCL 版權所有<br>開發者：何立智、楊哲睿</div>',
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    app = get_layout()
+    app.queue().launch(server_name="0.0.0.0")