Spaces:
Sleeping
Sleeping
File size: 10,413 Bytes
4cd3714 22a68d2 714cfbf d781ad4 7008805 4cd3714 dc36253 4cd3714 7008805 22a68d2 7008805 5a203eb 4cd3714 dc36253 4cd3714 22a68d2 4cd3714 22a68d2 925fada 22a68d2 925fada 22a68d2 925fada 22a68d2 925fada 22a68d2 925fada 22a68d2 925fada dc36253 925fada dc36253 4cd3714 dc36253 5a203eb 925fada 7008805 22a68d2 dc36253 22a68d2 dc36253 5a203eb dc36253 925fada 22a68d2 dc36253 925fada 5a203eb 22a68d2 dc36253 22a68d2 dc36253 925fada dc36253 22a68d2 925fada 22a68d2 76b9e1c 22a68d2 9bb0397 0f7ace5 9bb0397 fbd87f3 0f7ace5 9bb0397 0f7ace5 22a68d2 fbd87f3 0f7ace5 22a68d2 925fada 22a68d2 0f7ace5 5e6a1f4 59adf16 b05b0e9 22a68d2 7008805 4cd3714 e65f39b 22a68d2 749de3f 22a68d2 749de3f 22a68d2 5e6a1f4 22a68d2 5e6a1f4 749de3f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 |
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
# Function to append patient info and image to HTML database
import os
# Function to append patient info and image to HTML database
def append_patient_info_to_html(name, age, medical_record, sex, result, predicted_image_base64):
# Check if the table header is already present
if os.stat(html_db_file).st_size == 0: # Empty file, need to 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>
""")
# Append patient data with a horizontal line and descriptive headers
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>
<tr><td colspan="6"><hr></td></tr> <!-- Horizontal line across all columns -->
"""
with open(html_db_file, 'a') as f:
f.write(html_entry)
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() |