fix model/tensor device difference

app.py

@@ -10,65 +10,55 @@ import torch
 import torch.nn.functional as F
 
 
-class Litton7Classifier:
-    LABELS = [
-        "Panoramic",
-        "Feature",
-        "Detail",
-        "Enclosed",
-        "Focal",
-        "Ephemeral",
-        "Canopied",
+LABELS = [
+    "Panoramic",
+    "Feature",
+    "Detail",
+    "Enclosed",
+    "Focal",
+    "Ephemeral",
+    "Canopied",
+]
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+model = torch.load(
+    "Litton-7type-visual-landscape-model.pth", map_location=device, weights_only=False
+).module
+model.eval()
+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 __init__(
-        self, model_path="Litton-7type-visual-landscape-model.pth", device=None
-    ):
-        if device is None:
-            self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
-        else:
-            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]
-                ),
-            ]
-        )
-
-    @spaces.GPU(duration=60)
-    def predict(self, image: Image.Image) -> Figure:
-        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()
-
-        return draw_bar_chart(
-            {
-                "class": self.LABELS,
-                "probs": probs[0] * 100,
-            }
-        )
+)
+
+@spaces.GPU
+def predict(image: Image.Image) -> Figure:
+    image = image.convert("RGB")
+    input_tensor = preprocess(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        logits = model(input_tensor)
+        probs = F.softmax(logits[:, :7], dim=1).cpu()
+
+    return draw_bar_chart(
+        {
+            "class": LABELS,
+            "probs": probs[0] * 100,
+        }
+    )
 
 
 def draw_bar_chart(data: dict[str, list[str | float]]):
     classes = data["class"]
     probabilities = data["probs"]
 
-    #fig = plt.figure()
-    fig, ax = plt.subplots(figsize=(8, 6))  
+    fig, ax = plt.subplots(figsize=(8, 6))
     ax.bar(classes, probabilities, color="skyblue")
 
     ax.set_xlabel("Class")
@@ -149,7 +139,7 @@ def get_layout():
                 '<div class="footer">© 2024 LCL 版權所有<br>開發者：何立智、楊哲睿</div>',
             )
             start_button.click(
-                fn=Litton7Classifier().predict,
+                fn=predict,
                 inputs=image_input,
                 outputs=chart,
             )