Update app.py

app.py

@@ -170,21 +170,12 @@ def submit_result(name, patient_id, input_image, predicted_image, result):
 def view_database():
     conn = sqlite3.connect('results.db')
     c = conn.cursor()
-    c.execute("SELECT name, patient_id, input_image, predicted_image, result FROM results")
+    c.execute("SELECT name, patient_id, result FROM results")
     rows = c.fetchall()
     conn.close()
     
     # Convert to pandas DataFrame for better display in Gradio
-    df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Input Image", "Predicted Image", "Raw Result"])
-    
-    # Decode images from BLOB and convert to displayable format
-    def decode_image(image_blob):
-        image_np = np.frombuffer(image_blob, dtype=np.uint8)
-        image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
-        return image
-    
-    df["Input Image"] = df["Input Image"].apply(lambda x: decode_image(x))
-    df["Predicted Image"] = df["Predicted Image"].apply(lambda x: decode_image(x))
+    df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Result"])
     
     return df
 
@@ -194,46 +185,57 @@ def download_file(choice):
         # Provide the path to the database file
         return 'results.db'
     elif choice == "Database (.html)":
+        # Export the database as an HTML file
         conn = sqlite3.connect('results.db')
         c = conn.cursor()
-        c.execute("SELECT name, patient_id, input_image, predicted_image, result FROM results")
+        c.execute("SELECT name, patient_id, result FROM results")
         rows = c.fetchall()
         conn.close()
-        df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Input Image", "Predicted Image", "Raw Result"])
-        html = df.to_html()
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.html') as f:
-            f.write(html.encode())
-            return f.name
+        
+        df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Result"])
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.html') as temp_file:
+            df.to_html(temp_file.name, index=False)
+            return temp_file.name
     else:
-        # Handle other download options if necessary
-        pass
+        conn = sqlite3.connect('results.db')
+        c = conn.cursor()
+        c.execute("SELECT predicted_image FROM results ORDER BY id DESC LIMIT 1")
+        row = c.fetchone()
+        conn.close()
+        if row:
+            image_bytes = row[0]
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+                temp_file.write(image_bytes)
+                temp_file.flush()  # Ensure all data is written before closing the file
+                return temp_file.name
+        else:
+            return None
 
 # Initialize the database
 init_db()
 
-# Define the Gradio interface
-with gr.Blocks() as demo:
-    with gr.Tab("YOLOv8 Inference"):
-        with gr.Row():
-            input_image = gr.Image(label="Input Image", type="pil")
-        with gr.Row():
-            name = gr.Textbox(label="Patient Name")
-            patient_id = gr.Textbox(label="Patient ID")
-        with gr.Row():
-            submit_button = gr.Button("Submit")
-            predicted_image = gr.Image(label="Predicted Image")
-        with gr.Row():
-            result = gr.Textbox(label="Raw Result", lines=5)
-        submit_button.click(predict_image, inputs=[input_image, name, patient_id], outputs=[predicted_image, result])
+# Gradio interface
+with gr.Blocks() as interface:
+    gr.Markdown("## Image Prediction System")
     
-    with gr.Tab("View Database"):
-        view_button = gr.Button("View Database")
-        database_output = gr.DataFrame(label="Database Records")
-        view_button.click(view_database, outputs=database_output)
+    with gr.Row():
+        with gr.Column():
+            name_input = gr.Textbox(label="Name", placeholder="Enter patient name")
+            patient_id_input = gr.Textbox(label="Patient ID", placeholder="Enter patient ID")
+            input_image = gr.Image(source="upload", type="pil", label="Input Image")
+            submit_button = gr.Button("Submit")
+            database_view = gr.DataFrame(label="Database View")
         
-        download_choice = gr.Radio(["Database (.db)", "Database (.html)", "Predicted Image (.png)"], label="Choose the file to download:")
-        download_button = gr.Button("Download")
-        download_button.click(download_file, inputs=download_choice, outputs=gr.File())
-
-# Launch the interface
-demo.launch()
+        with gr.Column():
+            predicted_image_output = gr.Image(type="pil", label="Predicted Image")
+            result_text_output = gr.Textbox(label="Raw Result")
+            download_choice = gr.Radio(["Database (.db)", "Database (.html)", "Predicted Image (.png)"], label="Download Options")
+            download_button = gr.File(label="Download", interactive=True)
+
+    submit_button.click(predict_image, inputs=[input_image, name_input, patient_id_input], outputs=[predicted_image_output, result_text_output])
+    submit_button.click(submit_result, inputs=[name_input, patient_id_input, input_image, predicted_image_output, result_text_output], outputs=None)
+    gr.Button("View Database").click(view_database, outputs=database_view)
+    download_button.upload(download_file, inputs=download_choice, outputs=download_button)
+    
+interface.launch()