Update main.py

main.py

@@ -9,6 +9,10 @@ import os
 
 app = FastAPI()
 
+# Ensure the db directory exists
+DB_DIR = "/app/db"
+os.makedirs(DB_DIR, exist_ok=True)
+
 # Load the EfficientNet-B0 model
 model_name = "google/efficientnet-b0"
 feature_extractor = EfficientNetImageProcessor.from_pretrained(model_name)
@@ -113,7 +117,7 @@ benefits_mapping = {
 
 # SQLite database initialization for nutrition
 def init_db():
-    conn = sqlite3.connect("fruit_nutrition.db")
+    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_nutrition.db"))
     c = conn.cursor()
     c.execute('''CREATE TABLE IF NOT EXISTS fruits (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -139,7 +143,7 @@ def init_db():
 
 # New database for diabetic-specific information
 def init_diabetes_db():
-    conn = sqlite3.connect("fruit_diabetes.db")
+    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_diabetes.db"))
     c = conn.cursor()
     c.execute('''CREATE TABLE IF NOT EXISTS diabetes_info (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -186,7 +190,7 @@ async def predict(file: UploadFile = File(...)):
     bengali_fruit = fruit_mapping.get(predicted_label, predicted_label)
 
     # Fetch nutrition from database
-    conn = sqlite3.connect("fruit_nutrition.db")
+    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_nutrition.db"))
     c = conn.cursor()
     c.execute("SELECT vitamin, mineral, carbohydrate, protein, amino_acid FROM fruits WHERE name = ?", (bengali_fruit,))
     nutrition = c.fetchone()
@@ -202,7 +206,7 @@ async def predict(file: UploadFile = File(...)):
     })
 
     # Fetch diabetes info
-    conn = sqlite3.connect("fruit_diabetes.db")
+    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_diabetes.db"))
     c = conn.cursor()
     c.execute("SELECT glycemic_index, portion_size, best_time, precautions FROM diabetes_info WHERE fruit_name = ?", (bengali_fruit,))
     diabetes_info = c.fetchone()