Update app.py

app.py

@@ -12,7 +12,7 @@ import pandas as pd
 # Load YOLOv8 model
 model = YOLO("best.pt")
 
-def predict_image(input_image, name, patient_id):
+def predict_image(input_image, name, age, medical_record, sex):
     if input_image is None:
         return None, "Please Input The Image"
 
@@ -73,7 +73,7 @@ def predict_image(input_image, name, patient_id):
     pil_image_with_boxes = Image.fromarray(image_with_boxes)
 
     # Add text and watermark
-    pil_image_with_boxes = add_text_and_watermark(pil_image_with_boxes, name, patient_id, label)
+    pil_image_with_boxes = add_text_and_watermark(pil_image_with_boxes, name, age, medical_record, sex, label)
     
     return pil_image_with_boxes, raw_predictions_str
 
@@ -103,7 +103,7 @@ def add_watermark(image):
         return image
 
 # Function to add text and watermark
-def add_text_and_watermark(image, name, patient_id, label):
+def add_text_and_watermark(image, name, age, medical_record, sex, label):
     draw = ImageDraw.Draw(image)
     
     # Load a larger font (adjust the size as needed)
@@ -114,7 +114,7 @@ def add_text_and_watermark(image, name, patient_id, label):
         font = ImageFont.load_default()
         print("Error: cannot open resource, using default font.")
     
-    text = f"Name: {name}, ID: {patient_id}, Result: {label}"
+    text = f"Name: {name}, Age: {age}, Medical Record: {medical_record}, Sex: {sex}, Result: {label}"
     
     # Calculate text bounding box
     text_bbox = draw.textbbox((0, 0), text, font=font)
@@ -142,12 +142,12 @@ def init_db():
     conn = sqlite3.connect('results.db')
     c = conn.cursor()
     c.execute('''CREATE TABLE IF NOT EXISTS results
-                 (id INTEGER PRIMARY KEY, name TEXT, patient_id TEXT, input_image BLOB, predicted_image BLOB, result TEXT)''')
+                 (id INTEGER PRIMARY KEY, name TEXT, age INTEGER, medical_record INTEGER, sex TEXT, input_image BLOB, predicted_image BLOB, result TEXT)''')
     conn.commit()
     conn.close()
 
 # Function to submit result to the database
-def submit_result(name, patient_id, input_image, predicted_image, result):
+def submit_result(name, age, medical_record, sex, input_image, predicted_image, result):
     conn = sqlite3.connect('results.db')
     c = conn.cursor()
     
@@ -160,8 +160,8 @@ def submit_result(name, patient_id, input_image, predicted_image, result):
     _, predicted_buffer = cv2.imencode('.png', predicted_image_rgb)
     predicted_image_bytes = predicted_buffer.tobytes()
     
-    c.execute("INSERT INTO results (name, patient_id, input_image, predicted_image, result) VALUES (?, ?, ?, ?, ?)", 
-              (name, patient_id, input_image_bytes, predicted_image_bytes, result))
+    c.execute("INSERT INTO results (name, age, medical_record, sex, input_image, predicted_image, result) VALUES (?, ?, ?, ?, ?, ?, ?)", 
+              (name, age, medical_record, sex, input_image_bytes, predicted_image_bytes, result))
     conn.commit()
     conn.close()
     return "Result submitted to database."
@@ -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, input_image, predicted_image FROM results")
+    c.execute("SELECT name, age, medical_record, sex, 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", "Input Image", "Predicted Image"])
+    df = pd.DataFrame(rows, columns=["Name", "Age", "Medical Record", "Sex", "Input Image", "Predicted Image", "Result"])
     
     return df
 
@@ -203,12 +203,12 @@ def download_file(choice):
 init_db()
 
 # Gradio Interface
-def interface(name, patient_id, input_image):
+def interface(name, age, medical_record, sex, 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)
+    output_image, raw_result = predict_image(input_image, name, age, medical_record, sex)
+    submit_status = submit_result(name, age, medical_record, sex, input_image, output_image, raw_result)
 
     return output_image, raw_result, submit_status
 
@@ -236,7 +236,9 @@ with gr.Blocks() as demo:
     
     with gr.Column():
         name = gr.Textbox(label="Name")
-        patient_id = gr.Textbox(label="Patient ID")
+        age = gr.Number(label="Age")
+        medical_record = gr.Number(label="Medical Record")
+        sex = gr.Radio(["Male", "Female"], label="Sex")
         input_image = gr.Image(type="pil", label="Upload an Image", image_mode="RGB")
 
     with gr.Column():
@@ -247,7 +249,7 @@ with gr.Blocks() as demo:
         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])
+    submit_btn.click(fn=interface, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result, submit_status])
 
     with gr.Column():
         view_db_btn = gr.Button("View Database")