import gradio as gr from ultralytics import YOLO import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont import base64 from io import BytesIO import tempfile import os from pathlib import Path import shutil # Load YOLOv8 model model = YOLO("best.pt") # Create directories if not present uploaded_folder = Path('Uploaded_Picture') predicted_folder = Path('Predicted_Picture') uploaded_folder.mkdir(parents=True, exist_ok=True) predicted_folder.mkdir(parents=True, exist_ok=True) # Path to store accumulated HTML data html_file_path = Path(tempfile.gettempdir()) / 'patient_data.html' # Function to predict image and add bounding box, text, circle, and watermark def predict_image(input_image, name, age, medical_record, sex): if input_image is None: return None, "Please Input The Image" # Convert Gradio input image (PIL Image) to numpy array image_np = np.array(input_image) # Ensure the image is in the correct format if len(image_np.shape) == 2: # grayscale to RGB image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB) elif image_np.shape[2] == 4: # RGBA to RGB image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB) # Perform prediction results = model(image_np) # Draw bounding boxes on the image image_with_boxes = image_np.copy() 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 elif class_index == 1: label = "Mature" color = (255, 0, 0) # Red for Mature else: label = "Normal" color = (0, 255, 0) # Green for Normal confidence = highest_confidence_result.conf.item() xmin, ymin, xmax, ymax = map(int, highest_confidence_result.xyxy[0]) # Draw the bounding box cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), color, 2) # Calculate the center of the bounding box center_x = (xmin + xmax) // 2 center_y = (ymin + ymax) // 2 # Calculate the radius (1/12 of the average of the width and height of the bounding box) box_width = xmax - xmin box_height = ymax - ymin radius = int((box_width + box_height) / 24) # Average of width and height divided by 12 # Draw a white circle at the center of the bounding box cv2.circle(image_with_boxes, (center_x, center_y), radius, (255, 255, 255), thickness=2) # Enlarge font scale and thickness 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) raw_predictions.append(f"Label: {label}, Confidence: {confidence:.2f}, Box: [{xmin}, {ymin}, {xmax}, {ymax}]") raw_predictions_str = "\n".join(raw_predictions) # Convert to PIL image for further processing 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, age, medical_record, sex, label) # Save images to directories image_name = f"{name}-{age}-{sex}-{medical_record}.png" input_image.save(uploaded_folder / image_name) pil_image_with_boxes.save(predicted_folder / image_name) return pil_image_with_boxes, raw_predictions_str # Function to add watermark def add_watermark(image): try: logo = Image.open('image-logo.png').convert("RGBA") image = image.convert("RGBA") # Resize logo basewidth = 100 wpercent = (basewidth / float(logo.size[0])) hsize = int((float(wpercent) * logo.size[1])) logo = logo.resize((basewidth, hsize), Image.LANCZOS) # Position logo position = (image.width - logo.width - 10, image.height - logo.height - 10) # Composite image transparent = Image.new('RGBA', (image.width, image.height), (0, 0, 0, 0)) transparent.paste(image, (0, 0)) transparent.paste(logo, position, mask=logo) return transparent.convert("RGB") except Exception as e: print(f"Error adding watermark: {e}") return image # Function to add text and watermark 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 try: font = ImageFont.truetype("font.ttf", size=font_size) except IOError: font = ImageFont.load_default() 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 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], fill="black" ) # Draw text on top of the rectangle draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font) # Add watermark to the image image_with_watermark = add_watermark(image) return image_with_watermark # Function to save patient info in HTML and accumulate data def save_patient_info_to_html(name, age, medical_record, sex, result): html_content = f"""

Patient Information

Name: {name}

Age: {age}

Medical Record: {medical_record}

Sex: {sex}

Result: {result}

""" # Check if the HTML file already exists if html_file_path.exists(): with open(html_file_path, 'a') as f: f.write(html_content) else: with open(html_file_path, 'w') as f: f.write(html_content) return str(html_file_path) # Function to download the folders def download_folder(folder): zip_path = os.path.join(tempfile.gettempdir(), f"{folder}.zip") # Zip the folder shutil.make_archive(zip_path.replace('.zip', ''), 'zip', folder) return zip_path # Gradio Interface def interface(name, age, medical_record, sex, input_image): if input_image is None: return None, "Please upload an image.", None output_image, raw_result = predict_image(input_image, name, age, medical_record, sex) if output_image is None: return None, raw_result, None # Save patient info to HTML html_file_path = save_patient_info_to_html(name, age, medical_record, sex, raw_result) # Encode the image to display in Gradio buffered = BytesIO() output_image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") # Provide the zip file path for download zip_file = download_folder(predicted_folder) return f'Processed Image', raw_result, zip_file # Define Gradio interface gr.Interface( fn=interface, inputs=[ gr.Textbox(label="Name"), gr.Textbox(label="Age"), gr.Textbox(label="Medical Record"), gr.Dropdown(label="Sex", choices=["Male", "Female", "Other"]), gr.Image(source="upload", tool="editor", label="Upload an Image") ], outputs=[ gr.HTML(label="Processed Image"), gr.Textbox(label="Raw Predictions"), gr.File(label="Download ZIP") ], title="Patient Image Analysis" ).launch()