Update app.py

app.py

@@ -151,18 +151,15 @@ 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()
@@ -173,12 +170,21 @@ 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, 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", "Result"])
+    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))
     
     return df
 
@@ -190,93 +196,44 @@ def download_file(choice):
     elif choice == "Database (.html)":
         conn = sqlite3.connect('results.db')
         c = conn.cursor()
-        c.execute("SELECT * FROM results")
+        c.execute("SELECT name, patient_id, input_image, predicted_image, result 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
+        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
     else:
-        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
-                return temp_file.name
-        else:
-            raise FileNotFoundError("No images found in the database.")
+        # Handle other download options if necessary
+        pass
 
 # Initialize the database
 init_db()
 
-# 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
+# Define the Gradio interface
 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.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])
     
-    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])
+    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)
+        
+        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 Gradio app
+# Launch the interface
 demo.launch()