Update app.py

app.py

@@ -41,20 +41,17 @@ def predict_image(input_image, name, age, medical_record, sex):
     raw_predictions = []
 
     if results[0].boxes:
-        # Sort the results by confidence and take the highest confidence one
         highest_confidence_result = max(results[0].boxes, key=lambda x: x.conf.item())
-
-        # Determine the label based on the class index
         class_index = highest_confidence_result.cls.item()
         if class_index == 0:
             label = "Immature"
-            color = (0, 255, 255)  # Yellow for Immature
+            color = (0, 255, 255)
         elif class_index == 1:
             label = "Mature"
-            color = (255, 0, 0)  # Red for Mature
+            color = (255, 0, 0)
         else:
             label = "Normal"
-            color = (0, 255, 0)  # Green for Normal
+            color = (0, 255, 0)
 
         confidence = highest_confidence_result.conf.item()
         xmin, ymin, xmax, ymax = map(int, highest_confidence_result.xyxy[0])
@@ -74,10 +71,6 @@ def predict_image(input_image, name, age, medical_record, sex):
         font_scale = 1.0
         thickness = 2
         
-        # Calculate label background size
-        (text_width, text_height), baseline = cv2.getTextSize(f'{label} {confidence:.2f}', cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
-        cv2.rectangle(image_with_boxes, (xmin, ymin - text_height - baseline), (xmin + text_width, ymin), (0, 0, 0), cv2.FILLED)
-        
         # Put the label text with black background
         cv2.putText(image_with_boxes, f'{label} {confidence:.2f}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness)
         
@@ -102,8 +95,7 @@ def predict_image(input_image, name, age, medical_record, sex):
 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)
-    font_size = 24  # Example font size
+    font_size = 24
     try:
         font = ImageFont.truetype("font.ttf", size=font_size)
     except IOError:
@@ -111,17 +103,12 @@ def add_text_and_watermark(image, name, age, medical_record, sex, label):
         print("Error: cannot open resource, using default font.")
     
     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)
-    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
-    text_x = 20
-    text_y = 40
+    text_x, text_y = 20, 40
     padding = 10
 
     # Draw a filled rectangle for the background
     draw.rectangle(
-        [text_x - padding, text_y - padding, text_x + text_width + padding, text_y + text_height + padding],
+        [text_x - padding, text_y - padding, text_x + 500, text_y + 30 + padding],
         fill="black"
     )
     
@@ -145,20 +132,16 @@ def save_patient_info_to_html(name, age, medical_record, sex, result):
     </html>
     """
     
-    # Save HTML content to file
     html_file_path = os.path.join(tempfile.gettempdir(), 'patient_info.html')
     with open(html_file_path, 'w') as f:
         f.write(html_content)
     
     return html_file_path
 
-# Function to download the folders (fix: pass string path instead of Path object)
+# Function to download the folders
 def download_folder(folder_path):
     zip_path = os.path.join(tempfile.gettempdir(), f"{Path(folder_path).name}.zip")
-    
-    # Zip the folder
     shutil.make_archive(zip_path.replace('.zip', ''), 'zip', folder_path)
-    
     return zip_path
 
 # Gradio Interface
@@ -187,16 +170,18 @@ with gr.Blocks() as demo:
         download_uploaded_btn = gr.Button("Download Uploaded Images")
         download_predicted_btn = gr.Button("Download Predicted Images")
     
-    # Add file download output components for the uploaded and predicted images
     patient_info_file = gr.File(label="Patient Information HTML File")
     uploaded_folder_file = gr.File(label="Uploaded Images Zip File")
     predicted_folder_file = gr.File(label="Predicted Images Zip File")
 
-    # Fix: pass string representation of Path objects
+    # Use gr.State to hold folder paths
+    uploaded_folder_state = gr.State(str(uploaded_folder))
+    predicted_folder_state = gr.State(str(predicted_folder))
+
     submit_btn.click(fn=predict_image, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result])
     download_html_btn.click(fn=save_patient_info_to_html, inputs=[name, age, medical_record, sex, raw_result], outputs=patient_info_file)
-    download_uploaded_btn.click(fn=download_folder, inputs=[str(uploaded_folder)], outputs=uploaded_folder_file)
-    download_predicted_btn.click(fn=download_folder, inputs=[str(predicted_folder)], outputs=predicted_folder_file)
+    download_uploaded_btn.click(fn=download_folder, inputs=[uploaded_folder_state], outputs=uploaded_folder_file)
+    download_predicted_btn.click(fn=download_folder, inputs=[predicted_folder_state], outputs=predicted_folder_file)
 
 # Launch Gradio app
 demo.launch()