Update app.py

app.py

@@ -6,18 +6,29 @@ import os
 import cv2
 import tensorflow as tf
 from datetime import datetime
-import psycopg2
-from psycopg2.extras import DictCursor
+import sqlite3
 
 app = Flask(__name__)
 
-# ✅ 1. Connect to PostgreSQL (Supabase)
-POSTGRES_URL = "postgresql://postgres.otihqjwfqjwccsipzroy:Dhruvagr%[email protected]:5432/postgres"
+# ✅ 1. SQLite database setup
+DB_PATH = 'results.db'
+def init_db():
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS results (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            image_filename TEXT,
+            prediction TEXT,
+            confidence REAL,
+            gradcam_filename TEXT,
+            timestamp TEXT
+        )
+    """)
+    conn.commit()
+    conn.close()
 
-def get_db_connection():
-    """Get a new database connection."""
-    conn = psycopg2.connect(POSTGRES_URL)
-    return conn
+init_db()
 
 # ✅ 2. Load Model
 model = load_model('mobilenet_glaucoma_model.h5', compile=False)
@@ -86,7 +97,7 @@ def test_file():
 
 @app.route('/predict', methods=['POST'])
 def predict():
-    """Perform prediction and save results to PostgreSQL (Supabase)."""
+    """Perform prediction and save results to SQLite database."""
     if 'file' not in request.files:
         return jsonify({'error': 'No file uploaded'}), 400
 
@@ -109,15 +120,14 @@ def predict():
         gradcam_filename = f"gradcam_{int(datetime.now().timestamp())}.png"
         save_gradcam_image(img, heatmap, filename=gradcam_filename)
 
-        # Save results to PostgreSQL
-        conn = get_db_connection()
+        # Save results to SQLite
+        conn = sqlite3.connect(DB_PATH)
         cursor = conn.cursor()
         cursor.execute("""
             INSERT INTO results (image_filename, prediction, confidence, gradcam_filename, timestamp) 
-            VALUES (%s, %s, %s, %s, %s)
+            VALUES (?, ?, ?, ?, ?)
         """, (file.filename, result, confidence, gradcam_filename, datetime.now().isoformat()))
         conn.commit()
-        cursor.close()
         conn.close()
 
         return jsonify({
@@ -133,23 +143,22 @@ def predict():
 
 @app.route('/results', methods=['GET'])
 def results():
-    """List all results from the PostgreSQL database."""
-    conn = get_db_connection()
-    cursor = conn.cursor(cursor_factory=DictCursor)
+    """List all results from the SQLite database."""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
     cursor.execute("SELECT * FROM results ORDER BY timestamp DESC")
     results_data = cursor.fetchall()
-    cursor.close()
     conn.close()
 
     results_list = []
     for record in results_data:
         results_list.append({
-            'id': record['id'],  # Assumes 'id' is your primary key column
-            'image_filename': record['image_filename'],
-            'prediction': record['prediction'],
-            'confidence': record['confidence'],
-            'gradcam_filename': record['gradcam_filename'],
-            'timestamp': record['timestamp']
+            'id': record[0],
+            'image_filename': record[1],
+            'prediction': record[2],
+            'confidence': record[3],
+            'gradcam_filename': record[4],
+            'timestamp': record[5]
         })
 
     return jsonify(results_list)