Update app.py

app.py

@@ -9,6 +9,7 @@ from io import BytesIO
 import tempfile
 import pandas as pd
 import os
+from pathlib import Path
 
 # Load YOLOv8 model
 model = YOLO("best.pt")
@@ -203,30 +204,37 @@ def view_database():
 
 # Function to download database or HTML file
 def download_file(choice):
+    # Define a temporary directory to store the database file
     directory = tempfile.gettempdir()
-    
-    # Ensure the directory exists
-    if not os.path.exists(directory):
-        os.makedirs(directory)
-
     db_file_path = os.path.join(directory, 'results.db')
 
     if choice == "Database (.db)":
-        # Return the correct database file path
+        # Ensure the database file exists
+        if not os.path.exists(db_file_path):
+            # Create and populate the database with dummy data if needed
+            conn = sqlite3.connect(db_file_path)
+            # Example table creation (if not already exists)
+            conn.execute('''CREATE TABLE IF NOT EXISTS results (
+                            id INTEGER PRIMARY KEY,
+                            name TEXT NOT NULL,
+                            value TEXT NOT NULL
+                        );''')
+            conn.execute("INSERT INTO results (name, value) VALUES ('example', '123')")
+            conn.commit()
+            conn.close()
+
+        # Return the database file path for download
         return db_file_path
     
     elif choice == "Database (.html)":
-        # Check if the database file exists
-        if not os.path.isfile(db_file_path):
-            raise FileNotFoundError(f"Database file not found at path: {db_file_path}")
-
         # Connect to the SQLite database
         conn = sqlite3.connect(db_file_path)
-        
+
         try:
             # Attempt to read the results table into a DataFrame
             df = pd.read_sql_query("SELECT * FROM results", conn)
         except pd.errors.DatabaseError as e:
+            # Handle the case where the table doesn't exist
             conn.close()
             raise ValueError("Table 'results' does not exist in the database.") from e
         
@@ -243,7 +251,7 @@ def download_file(choice):
         return html_file_path
     else:
         raise ValueError("Invalid choice. Please select a valid format.")
-        
+
 # Initialize the database
 init_db()
 
@@ -309,7 +317,7 @@ with gr.Blocks() as demo:
         download_btn = gr.Button("Download")
         download_output = gr.File(label="Download File")
 
-    download_btn.click(fn=download_interface, inputs=[download_choice], outputs=gr.File())
+    download_btn.click(download_file, inputs=choice, outputs=download_file_output)
 
 # Launch the Gradio app
 demo.launch()