Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,34 +1,34 @@
-import streamlit as st
-import torch
-import torchvision.transforms as transforms
-from PIL import Image
-
-# โหลดโมเดล
-model = torch.load('mobilenetv2.pth', map_location=torch.device('cpu'))
-model.eval()
-
-# โหลด label
-with open('classes.txt', 'r') as f:
-    class_names = [line.strip() for line in f]
-
-# UI
-st.title("🔥 MobileNetV2 Classifier")
-uploaded_file = st.file_uploader("อัปโหลดภาพ", type=["jpg", "png", "jpeg"])
-
-if uploaded_file:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-    ])
-    img_tensor = transform(image).unsqueeze(0)
-
-    with torch.no_grad():
-        outputs = model(img_tensor)
-        probs = torch.softmax(outputs, dim=1)
-        pred_class = class_names[probs.argmax().item()]
-        confidence = probs.max().item() * 100
-
+import streamlit as st
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+
+# โหลดโมเดล
+model = models.mobilenet_v2(pretrained=True)
+model.classifier[1] = torch.nn.Linear(model.last_channel, 3)  # 3 class
+model.load_state_dict(torch.load('mobilenetv2.pth', map_location=device))
+# โหลด label
+with open('classes.txt', 'r') as f:
+    class_names = [line.strip() for line in f]
+
+# UI
+st.title("🔥 MobileNetV2 Classifier")
+uploaded_file = st.file_uploader("อัปโหลดภาพ", type=["jpg", "png", "jpeg"])
+
+if uploaded_file:
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+    ])
+    img_tensor = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = torch.softmax(outputs, dim=1)
+        pred_class = class_names[probs.argmax().item()]
+        confidence = probs.max().item() * 100
+
     st.markdown(f"### 🔍 คำทำนาย: `{pred_class}` ({confidence:.2f}%)")