Update app.py

app.py

@@ -151,15 +151,18 @@ def submit_result(name, patient_id, input_image, predicted_image, result):
     conn = sqlite3.connect('results.db')
     c = conn.cursor()
     
+    # Encode input image
     input_image_np = np.array(input_image)
     _, input_buffer = cv2.imencode('.png', cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR))
     input_image_bytes = input_buffer.tobytes()
     
+    # Encode predicted image
     predicted_image_np = np.array(predicted_image)
     predicted_image_rgb = cv2.cvtColor(predicted_image_np, cv2.COLOR_RGB2BGR)  # Ensure correct color conversion
     _, predicted_buffer = cv2.imencode('.png', predicted_image_rgb)
     predicted_image_bytes = predicted_buffer.tobytes()
     
+    # Insert into database
     c.execute("INSERT INTO results (name, patient_id, input_image, predicted_image, result) VALUES (?, ?, ?, ?, ?)", 
               (name, patient_id, input_image_bytes, predicted_image_bytes, result))
     conn.commit()
@@ -170,12 +173,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 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"])
+    df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Result"])
     
     return df
 
@@ -185,7 +188,15 @@ def download_file(choice):
         # Provide the path to the database file
         return 'results.db'
     elif choice == "Database (.html)":
-        df = view_database()
+        conn = sqlite3.connect('results.db')
+        c = conn.cursor()
+        c.execute("SELECT * FROM results")
+        rows = c.fetchall()
+        conn.close()
+        
+        df = pd.DataFrame(rows, columns=["ID", "Name", "Patient ID", "Input Image", "Predicted Image", "Result"])
+        
+        # Convert DataFrame to HTML
         with tempfile.NamedTemporaryFile(delete=False, suffix='.html') as temp_file:
             df.to_html(temp_file.name)
             return temp_file.name
@@ -214,37 +225,58 @@ def interface(name, patient_id, input_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, submit_status
-
-inputs = [
-    gr.Textbox(label="Name"),
-    gr.Textbox(label="Patient ID"),
-    gr.Image(type="pil", label="Input Image")
-]
-
-outputs = [
-    gr.Image(label="Output Image"),
-    gr.Textbox(label="Status")
-]
-
-# File download interface
-download_inputs = gr.Radio(["Database (.db)", "Database (.html)", "Image (.png)"], label="Download Type")
-download_output = gr.File(label="Download File")
-
-app = gr.Interface(
-    fn=interface,
-    inputs=inputs,
-    outputs=outputs,
-    title="AI Cataract Detector",
-    description="Upload an image, enter the patient's name and ID, and receive a prediction."
-)
-
-download_app = gr.Interface(
-    fn=download_file,
-    inputs=download_inputs,
-    outputs=download_output,
-    title="Download Results"
-)
-
-# Combine both interfaces in one layout
-gr.TabbedInterface([app, download_app], ["Prediction", "Download"]).launch()
+
+    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():
+        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)", "Database (.html)"], 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()