Update app.py

app.py

@@ -170,12 +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, result FROM results")
+    c.execute("SELECT name, patient_id, input_image, predicted_image FROM results")
     rows = c.fetchall()
     conn.close()
     
     # Convert to pandas DataFrame for better display in Gradio
-    df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Result"])
+    df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Input Image", "Predicted Image"])
     
     return df
 
@@ -184,19 +184,6 @@ def download_file(choice):
     if choice == "Database (.db)":
         # 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, result FROM results")
-        rows = c.fetchall()
-        conn.close()
-        
-        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:
         conn = sqlite3.connect('results.db')
         c = conn.cursor()
@@ -207,35 +194,73 @@ def download_file(choice):
             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
+                temp_file.flush()  # Ensure all data is written before closing
                 return temp_file.name
         else:
-            return None
+            raise FileNotFoundError("No images found in the database.")
 
 # Initialize the database
 init_db()
 
-# Gradio interface
-with gr.Blocks() as interface:
-    gr.Markdown("## Image Prediction System")
+# Gradio Interface
+def interface(name, patient_id, input_image):
+    if input_image is None:
+        return "Please upload an image."
+
+    output_image, raw_result = predict_image(input_image, name, patient_id)
+    submit_status = submit_result(name, patient_id, input_image, output_image, raw_result)
+
+    return output_image, raw_result, submit_status
+
+# View Database Function
+def view_db_interface():
+    df = view_database()
+    return df
+
+# Download Function
+def download_interface(choice):
+    try:
+        file_path = download_file(choice)
+        with open(file_path, "rb") as file:
+            return file.read(), file_path.split('/')[-1]
+    except FileNotFoundError as e:
+        return str(e), None
+
+# Gradio Blocks
+with gr.Blocks() as demo:
+    with gr.Column():
+        gr.Markdown("# Cataract Detection System")
+        gr.Markdown("Upload an image to detect cataract and add patient details.")
+        gr.Image("PR_curve.png", label="Model PR Curve")
+        gr.Markdown("This application uses YOLOv8 with mAP=0.981")
     
+    with gr.Column():
+        name = gr.Textbox(label="Name")
+        patient_id = gr.Textbox(label="Patient ID")
+        input_image = gr.Image(type="pil", label="Upload an Image", image_mode="RGB")
+
+    with gr.Column():
+        submit_btn = gr.Button("Submit")
+        output_image = gr.Image(type="pil", label="Predicted Image")
+
     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")
-        
-        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()
+        raw_result = gr.Textbox(label="Raw Result", lines=5)
+        submit_status = gr.Textbox(label="Submission Status")
+
+    submit_btn.click(fn=interface, inputs=[name, patient_id, input_image], outputs=[output_image, raw_result, submit_status])
+
+    with gr.Column():
+        view_db_btn = gr.Button("View Database")
+        db_output = gr.Dataframe(label="Database Records")
+
+    view_db_btn.click(fn=view_db_interface, inputs=[], outputs=[db_output])
+
+    with gr.Column():
+        download_choice = gr.Radio(["Database (.db)", "Predicted Image (.png)"], label="Choose the file to download:")
+        download_btn = gr.Button("Download")
+        download_output = gr.File(label="Download File")
+
+    download_btn.click(fn=download_interface, inputs=[download_choice], outputs=[download_output])
+
+# Launch the Gradio app
+demo.launch()