Spaces:
Sleeping
Sleeping
| 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 for HTML database file | |
| html_db_file = Path('patient_predictions.html') | |
| # Initialize HTML database file if not present | |
| if not html_db_file.exists(): | |
| with open(html_db_file, 'w') as f: | |
| f.write("<html><body><h1>Patient Prediction Database</h1>") | |
| 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) | |
| # 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) | |
| # Convert the predicted image to base64 for embedding in HTML | |
| buffered = BytesIO() | |
| pil_image_with_boxes.save(buffered, format="PNG") | |
| predicted_image_base64 = base64.b64encode(buffered.getvalue()).decode() | |
| # Append the prediction to the HTML database | |
| append_patient_info_to_html(name, age, medical_record, sex, label, predicted_image_base64) | |
| 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 | |
| def append_patient_info_to_html(name, age, medical_record, sex, result, predicted_image_base64): | |
| # Check if the HTML file is empty or if the table structure is missing | |
| if os.stat(html_db_file).st_size == 0: # Empty file, create the table structure | |
| with open(html_db_file, 'a') as f: | |
| f.write(""" | |
| <html> | |
| <head><title>Patient Prediction Database</title></head> | |
| <body> | |
| <h1>Patient Prediction Database</h1> | |
| <table border="1" style="width:100%; border-collapse: collapse; text-align: center;"> | |
| <thead> | |
| <tr> | |
| <th>Name</th> | |
| <th>Age</th> | |
| <th>Medical Record</th> | |
| <th>Sex</th> | |
| <th>Result</th> | |
| <th>Predicted Image</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| """) | |
| # Check if this patient already exists to prevent duplicate entries | |
| # This can be improved by checking unique identifiers like `medical_record` | |
| # Assuming the uniqueness of the medical record | |
| html_entry = f""" | |
| <tr> | |
| <td>{name}</td> | |
| <td>{age}</td> | |
| <td>{medical_record}</td> | |
| <td>{sex}</td> | |
| <td>{result}</td> | |
| <td><img src="data:image/png;base64,{predicted_image_base64}" alt="Predicted Image" width="150"></td> | |
| </tr> | |
| """ | |
| with open(html_db_file, 'a') as f: | |
| f.write(html_entry) | |
| # Ensure we only add the closing tags once | |
| if "</tbody></table></body></html>" not in open(html_db_file).read(): | |
| with open(html_db_file, 'a') as f: | |
| f.write(""" | |
| </tbody> | |
| </table> | |
| </body> | |
| </html> | |
| """) | |
| return str(html_db_file) # Return the HTML file path for download | |
| # 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) | |
| # Return the current state of the HTML file with all predictions | |
| return output_image, raw_result, str(html_db_file) | |
| # Download Functions | |
| def download_predicted_folder(): | |
| return download_folder(predicted_folder) | |
| def download_uploaded_folder(): | |
| return download_folder(uploaded_folder) | |
| # Launch 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.Markdown("This application uses YOLOv8 with mAP=0.981") | |
| with gr.Column(): | |
| name = gr.Textbox(label="Name") | |
| 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(): | |
| submit_btn = gr.Button("Submit") | |
| output_image = gr.Image(type="pil", label="Predicted Image") | |
| with gr.Row(): | |
| raw_result = gr.Textbox(label="Prediction Result") | |
| with gr.Row(): | |
| download_html_btn = gr.Button("Download Patient Information (HTML)") | |
| 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") | |
| # Connect functions with components | |
| submit_btn.click(fn=interface, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result]) | |
| download_html_btn.click(fn=append_patient_info_to_html, inputs=[name, age, medical_record, sex, raw_result], outputs=patient_info_file) | |
| download_uploaded_btn.click(fn=download_uploaded_folder, outputs=uploaded_folder_file) | |
| download_predicted_btn.click(fn=download_predicted_folder, outputs=predicted_folder_file) | |
| # Launch Gradio app | |
| demo.launch() |