Update app.py

app.py

@@ -1,15 +1,16 @@
-import matplotlib.pyplot as plt
-import torch
 from PIL import Image
+from io import BytesIO
 from torchvision import transforms
-import torch.nn.functional as F
 from typing import Literal, Any
 import gradio as gr
+from matplotlib.figure import Figure
+import matplotlib.pyplot as plt
 import spaces
-from io import BytesIO
+import torch
+import torch.nn.functional as F
 
 
-class Litton7Classifier:
+class Classifier:
     LABELS = [
         "Panoramic",
         "Feature",
@@ -47,7 +48,7 @@ class Litton7Classifier:
         )
 
     @spaces.GPU(duration=60)
-    def predict(self, image: Image.Image) -> tuple[Literal["Failed", "Success"], Any]:
+    def predict(self, image: Image.Image) -> Figure:
         image = image.convert("RGB")
         input_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
 
@@ -67,19 +68,20 @@ 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")
+    #fig = plt.figure()
+    fig, ax = plt.subplots(figsize=(8, 6))  
+    ax.bar(classes, probabilities, color="skyblue")
 
-    plt.xlabel("Class")
-    plt.ylabel("Probability (%)")
-    plt.title("Class Probabilities")
+    ax.set_xlabel("Class")
+    ax.set_ylabel("Probability (%)")
+    ax.set_title("Class Probabilities")
 
     for i, prob in enumerate(probabilities):
-        plt.text(i, prob + 0.01, f"{prob:.2f}", ha="center", va="bottom")
+        ax.text(i, prob + 0.01, f"{prob:.2f}", ha="center", va="bottom")
 
-    plt.tight_layout()
+    fig.tight_layout()
 
-    return plt
+    return fig
 
 
 def get_layout():
@@ -138,17 +140,17 @@ def get_layout():
                     "</div>"
                 ),
             )
-    
+
             with gr.Row(equal_height=True):
                 image_input = gr.Image(label="上傳影像", type="pil")
-                chart = gr.Image(label="分類結果")
+                chart = gr.Plot(label="分類結果")
 
             start_button = gr.Button("開始分類", variant="primary")
             gr.HTML(
                 '<div class="footer">© 2024 LCL 版權所有<br>開發者：何立智、楊哲睿</div>',
             )
             start_button.click(
-                fn=Litton7Classifier().predict,
+                fn=Classifier().predict,
                 inputs=image_input,
                 outputs=chart,
             )
@@ -157,4 +159,4 @@ def get_layout():
 
 
 if __name__ == "__main__":
-    get_layout().launch()
+    get_layout().launch()