Update app.py

app.py

@@ -4,6 +4,20 @@ import torch.nn as nn
 from torchvision import transforms, models
 import pickle
 from resnest.torch import resnest50
+import mysql.connector
+from datetime import datetime
+import os
+
+
+def get_db_connection():
+    return mysql.connector.connect(
+        host=os.environ['DB_HOST'],
+        port=os.environ['DB_PORT'],
+        user=os.environ['DB_USER'],
+        password=os.environ['DB_PASSWORD'],
+        database=os.environ['DB_NAME']
+    )
+
 
 with open('class_names.pkl', 'rb') as f:
     class_names = pickle.load(f)
@@ -66,6 +80,42 @@ def predict_image(img):
                 f"Top 3: {results}\n"
                 f"------------------------\n")
 
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor()
+        
+        # 创建表（如果尚未创建）
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS predictions (
+                id INT AUTO_INCREMENT PRIMARY KEY,
+                predicted_class VARCHAR(255),
+                confidence FLOAT,
+                top3_results TEXT,
+                timestamp DATETIME
+            )
+        ''')
+        
+        # 插入数据
+        insert_query = '''
+            INSERT INTO predictions 
+            (predicted_class, confidence, top3_results, timestamp)
+            VALUES (%s, %s, %s, %s)
+        '''
+        cursor.execute(insert_query, (
+            best_class,
+            best_conf,
+            str(results),
+            datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+        ))
+        
+        conn.commit()
+    except Exception as e:
+        print(f"Database error: {str(e)}")
+    finally:
+        if conn.is_connected():
+            cursor.close()
+            conn.close()
+    
     return best_class, best_conf, results
 
 # 创建Gradio界面