Update app.py

app.py

@@ -6,46 +6,63 @@ from transformers import Wav2Vec2Processor, HubertModel
 from sklearn.preprocessing import StandardScaler
 
 import gradio as gr
+import joblib
 
-# โหลด processor, model HuBERT และ fine-tuned classifier
+# 👇 โหลด processor และโมเดล HuBERT
 processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
 model = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft")
 model.eval()
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# โหลดไฟล์ scaler และ model classifier
-import joblib
+# 👇 โหลด scaler
 scaler = joblib.load("scaler.joblib")
-from your_module import PlantSoundClassifier  # นำเข้า class ของคุณเอง
 
+# 👇 โหลด classifier
+from your_module import PlantSoundClassifier  # แก้เป็น module/class ที่คุณสร้างเอง
+
+# ✅ ให้ torch รู้จัก class (safe!)
 torch.serialization.add_safe_globals({"__main__.PlantSoundClassifier": PlantSoundClassifier})
 
-model_cls = torch.load("classifier.pth", map_location=device, weights_only=True)
+# ✅ โหลด state_dict ของโมเดลอย่างปลอดภัย
+model_cls_state_dict = torch.load("classifier.pth", map_location=device, weights_only=True)
+
+# ✅ สร้าง instance classifier แล้วโหลด state_dict
+model_cls = PlantSoundClassifier().to(device)
+model_cls.load_state_dict(model_cls_state_dict)
 model_cls.eval()
 
 def extract_mean_embedding(wav_path):
     waveform, sample_rate = torchaudio.load(wav_path)
     waveform = waveform.squeeze()
+
     inputs = processor(waveform, sampling_rate=sample_rate, return_tensors="pt", padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
+
     embedding = outputs.last_hidden_state
-    mean_embedding = embedding.mean(dim=1).squeeze().numpy()
+    mean_embedding = embedding.mean(dim=1).squeeze().cpu().numpy()
     return mean_embedding
 
 def predict_water_status(file):
     vec = extract_mean_embedding(file).reshape(1, -1)
     vec_scaled = scaler.transform(vec)
     vec_tensor = torch.tensor(vec_scaled, dtype=torch.float32).to(device)
+
     with torch.no_grad():
         outputs = model_cls(vec_tensor)
         pred = outputs.argmax(dim=1).item()
-    return "🌵 ขาดน้ำ" if pred == 0 else "💧 มีน้ำเพียงพอ"
+
+    if pred == 0:
+        return "🌵 ขาดน้ำ"
+    elif pred == 1:
+        return "💧 มีน้ำเพียงพอ"
+    else:
+        return "⚠️ ไม่ทราบสถานะ"
 
 with gr.Blocks() as interface:
-    gr.Markdown("## 🌱 Plant Sound Classifier (Fine-tuned)")
-    gr.Markdown("อัปโหลดเสียงพืชเพื่อทำนายสถานะ: ขาดน้ำ หรือ มีน้ำเพียงพอ")
+    gr.Markdown("## 🌱 Plant Sound Classifier (Fine-tuned HuBERT)")
+    gr.Markdown("อัปโหลดเสียงพืช (.wav) เพื่อทำนายสถานะ: ขาดน้ำ หรือ มีน้ำเพียงพอ")
 
     audio_input = gr.Audio(type="filepath", label="🎧 อัปโหลดเสียงพืช (.wav)")
     output_text = gr.Textbox(label="📋 ผลการทำนาย", lines=2)