Update app.py

use Mongo DB to collect feedback.

app.py

@@ -9,11 +9,21 @@ from PIL import Image
 import json
 import sqlite3
 from datetime import datetime
+from pymongo import MongoClient
 
 # Load the dataset
 ds = load_dataset("svjack/pokemon-blip-captions-en-zh")
 ds = ds["train"]
 
+# load MongoDB client
+MONGO_URI = os.environ.get('MONGO_URI')
+if not MONGO_URI:
+    raise ValueError("MONGO_URI is not set in the environment variables.")
+
+client = MongoClient(MONGO_URI)
+db = client['capimg']  # Choose a database name
+collection = db['feedback']  # Choose a collection name
+
 # Load environment variables
 api_key = os.environ.get('MISTRAL_API_KEY')
 
@@ -87,18 +97,21 @@ def generate_caption(image):
         return None
 
 # Initialize SQLite database
-def init_db():
-    conn = sqlite3.connect('feedback.db')
-    c = conn.cursor()
-    c.execute('''CREATE TABLE IF NOT EXISTS thumbs_up
-                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
-                  timestamp TEXT,
-                  input_data TEXT,
-                  output_data TEXT)''')
-    conn.commit()
-    conn.close()
+def save_feedback(image, caption):
+    # Convert image to base64 string for storage
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
 
-init_db()
+    feedback_entry = {
+        "timestamp": datetime.now(),
+        "input_data": img_str,
+        "output_data": caption
+    }
+    
+    result = collection.insert_one(feedback_entry)
+    print(f"Feedback saved with id: {result.inserted_id}")
+    return gr.Info("Thanks for your feedback!")
 
 def process_image(image):
     if image is None:
@@ -112,23 +125,12 @@ def process_image(image):
         return "Failed to generate caption. Please check the API call or network connectivity."
 
 def thumbs_up(image, caption):
-    # Convert image to base64 string for storage
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-
-    conn = sqlite3.connect('feedback.db')
-    c = conn.cursor()
-    c.execute("INSERT INTO thumbs_up (timestamp, input_data, output_data) VALUES (?, ?, ?)",
-              (datetime.now().isoformat(), img_str, caption))
-    conn.commit()
-    conn.close()
-    print(f"Thumbs up data saved to database.")
-    return gr.Info("Thank you for your feedback!")
+    return save_feedback(image, caption)
 
 with gr.Blocks() as iface:
     gr.Markdown("# Image Captioner")
-    gr.Markdown("Upload an image to generate captions in English and Chinese. Use the 'Thumbs Up' button if you like the result!")
+    gr.Markdown("Upload an image to generate captions in English and Chinese.")
+    gr.Markdown("Use the 'Thumbs Up' button if you like the result!!") 
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -146,4 +148,4 @@ with gr.Blocks() as iface:
     thumbs_up_btn.click(fn=thumbs_up, inputs=[input_image, output_text], outputs=None)
 
 # Launch the interface
-iface.launch(share=True)
+iface.launch()