Spaces:
Sleeping
Sleeping
ariankhalfani
commited on
Commit
•
f49a800
1
Parent(s):
771b1be
Delete app3.py
Browse files
app3.py
DELETED
@@ -1,237 +0,0 @@
|
|
1 |
-
import gradio as gr
|
2 |
-
from ultralytics import YOLO
|
3 |
-
import cv2
|
4 |
-
import numpy as np
|
5 |
-
from PIL import Image, ImageDraw, ImageFont
|
6 |
-
import base64
|
7 |
-
from io import BytesIO
|
8 |
-
import tempfile
|
9 |
-
import os
|
10 |
-
from pathlib import Path
|
11 |
-
import shutil
|
12 |
-
|
13 |
-
# Load YOLOv8 model
|
14 |
-
model = YOLO("best.pt")
|
15 |
-
|
16 |
-
# Create directories if not present
|
17 |
-
uploaded_folder = Path('Uploaded_Picture')
|
18 |
-
predicted_folder = Path('Predicted_Picture')
|
19 |
-
uploaded_folder.mkdir(parents=True, exist_ok=True)
|
20 |
-
predicted_folder.mkdir(parents=True, exist_ok=True)
|
21 |
-
|
22 |
-
def predict_image(input_image, name, age, medical_record, sex):
|
23 |
-
if input_image is None:
|
24 |
-
return None, "Please Input The Image"
|
25 |
-
|
26 |
-
# Convert Gradio input image (PIL Image) to numpy array
|
27 |
-
image_np = np.array(input_image)
|
28 |
-
|
29 |
-
# Ensure the image is in the correct format
|
30 |
-
if len(image_np.shape) == 2: # grayscale to RGB
|
31 |
-
image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
|
32 |
-
elif image_np.shape[2] == 4: # RGBA to RGB
|
33 |
-
image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
|
34 |
-
|
35 |
-
# Perform prediction
|
36 |
-
results = model(image_np)
|
37 |
-
|
38 |
-
# Draw bounding boxes on the image
|
39 |
-
image_with_boxes = image_np.copy()
|
40 |
-
raw_predictions = []
|
41 |
-
|
42 |
-
if results[0].boxes:
|
43 |
-
# Sort the results by confidence and take the highest confidence one
|
44 |
-
highest_confidence_result = max(results[0].boxes, key=lambda x: x.conf.item())
|
45 |
-
|
46 |
-
# Determine the label based on the class index
|
47 |
-
class_index = highest_confidence_result.cls.item()
|
48 |
-
if class_index == 0:
|
49 |
-
label = "Immature"
|
50 |
-
color = (0, 255, 255) # Yellow for Immature
|
51 |
-
elif class_index == 1:
|
52 |
-
label = "Mature"
|
53 |
-
color = (255, 0, 0) # Red for Mature
|
54 |
-
else:
|
55 |
-
label = "Normal"
|
56 |
-
color = (0, 255, 0) # Green for Normal
|
57 |
-
|
58 |
-
confidence = highest_confidence_result.conf.item()
|
59 |
-
xmin, ymin, xmax, ymax = map(int, highest_confidence_result.xyxy[0])
|
60 |
-
|
61 |
-
# Draw the bounding box
|
62 |
-
cv2.rectangle(image_with_boxes, (xmin, ymin), (xmax, ymax), color, 2)
|
63 |
-
|
64 |
-
# Enlarge font scale and thickness
|
65 |
-
font_scale = 1.0
|
66 |
-
thickness = 2
|
67 |
-
|
68 |
-
# Calculate label background size
|
69 |
-
(text_width, text_height), baseline = cv2.getTextSize(f'{label} {confidence:.2f}', cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
|
70 |
-
cv2.rectangle(image_with_boxes, (xmin, ymin - text_height - baseline), (xmin + text_width, ymin), (0, 0, 0), cv2.FILLED)
|
71 |
-
|
72 |
-
# Put the label text with black background
|
73 |
-
cv2.putText(image_with_boxes, f'{label} {confidence:.2f}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness)
|
74 |
-
|
75 |
-
raw_predictions.append(f"Label: {label}, Confidence: {confidence:.2f}, Box: [{xmin}, {ymin}, {xmax}, {ymax}]")
|
76 |
-
|
77 |
-
raw_predictions_str = "\n".join(raw_predictions)
|
78 |
-
|
79 |
-
# Convert to PIL image for further processing
|
80 |
-
pil_image_with_boxes = Image.fromarray(image_with_boxes)
|
81 |
-
|
82 |
-
# Add text and watermark
|
83 |
-
pil_image_with_boxes = add_text_and_watermark(pil_image_with_boxes, name, age, medical_record, sex, label)
|
84 |
-
|
85 |
-
# Save images to directories
|
86 |
-
image_name = f"{name}-{age}-{sex}-{medical_record}.png"
|
87 |
-
input_image.save(uploaded_folder / image_name)
|
88 |
-
pil_image_with_boxes.save(predicted_folder / image_name)
|
89 |
-
|
90 |
-
return pil_image_with_boxes, raw_predictions_str
|
91 |
-
|
92 |
-
# Function to add watermark
|
93 |
-
def add_watermark(image):
|
94 |
-
try:
|
95 |
-
logo = Image.open('image-logo.png').convert("RGBA")
|
96 |
-
image = image.convert("RGBA")
|
97 |
-
|
98 |
-
# Resize logo
|
99 |
-
basewidth = 100
|
100 |
-
wpercent = (basewidth / float(logo.size[0]))
|
101 |
-
hsize = int((float(wpercent) * logo.size[1]))
|
102 |
-
logo = logo.resize((basewidth, hsize), Image.LANCZOS)
|
103 |
-
|
104 |
-
# Position logo
|
105 |
-
position = (image.width - logo.width - 10, image.height - logo.height - 10)
|
106 |
-
|
107 |
-
# Composite image
|
108 |
-
transparent = Image.new('RGBA', (image.width, image.height), (0, 0, 0, 0))
|
109 |
-
transparent.paste(image, (0, 0))
|
110 |
-
transparent.paste(logo, position, mask=logo)
|
111 |
-
|
112 |
-
return transparent.convert("RGB")
|
113 |
-
except Exception as e:
|
114 |
-
print(f"Error adding watermark: {e}")
|
115 |
-
return image
|
116 |
-
|
117 |
-
# Function to add text and watermark
|
118 |
-
def add_text_and_watermark(image, name, age, medical_record, sex, label):
|
119 |
-
draw = ImageDraw.Draw(image)
|
120 |
-
|
121 |
-
# Load a larger font (adjust the size as needed)
|
122 |
-
font_size = 24 # Example font size
|
123 |
-
try:
|
124 |
-
font = ImageFont.truetype("font.ttf", size=font_size)
|
125 |
-
except IOError:
|
126 |
-
font = ImageFont.load_default()
|
127 |
-
print("Error: cannot open resource, using default font.")
|
128 |
-
|
129 |
-
text = f"Name: {name}, Age: {age}, Medical Record: {medical_record}, Sex: {sex}, Result: {label}"
|
130 |
-
|
131 |
-
# Calculate text bounding box
|
132 |
-
text_bbox = draw.textbbox((0, 0), text, font=font)
|
133 |
-
text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
|
134 |
-
text_x = 20
|
135 |
-
text_y = 40
|
136 |
-
padding = 10
|
137 |
-
|
138 |
-
# Draw a filled rectangle for the background
|
139 |
-
draw.rectangle(
|
140 |
-
[text_x - padding, text_y - padding, text_x + text_width + padding, text_y + text_height + padding],
|
141 |
-
fill="black"
|
142 |
-
)
|
143 |
-
|
144 |
-
# Draw text on top of the rectangle
|
145 |
-
draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font)
|
146 |
-
|
147 |
-
# Add watermark to the image
|
148 |
-
image_with_watermark = add_watermark(image)
|
149 |
-
|
150 |
-
return image_with_watermark
|
151 |
-
|
152 |
-
# Function to save patient info in HTML
|
153 |
-
def save_patient_info_to_html(name, age, medical_record, sex, result):
|
154 |
-
html_content = f"""
|
155 |
-
<html>
|
156 |
-
<body>
|
157 |
-
<h1>Patient Information</h1>
|
158 |
-
<p><strong>Name:</strong> {name}</p>
|
159 |
-
<p><strong>Age:</strong> {age}</p>
|
160 |
-
<p><strong>Medical Record:</strong> {medical_record}</p>
|
161 |
-
<p><strong>Sex:</strong> {sex}</p>
|
162 |
-
<p><strong>Result:</strong> {result}</p>
|
163 |
-
</body>
|
164 |
-
</html>
|
165 |
-
"""
|
166 |
-
|
167 |
-
# Save HTML content to file
|
168 |
-
html_file_path = os.path.join(tempfile.gettempdir(), f'{name}-{age}-{sex}-{medical_record}.html')
|
169 |
-
with open(html_file_path, 'w') as f:
|
170 |
-
f.write(html_content)
|
171 |
-
|
172 |
-
return html_file_path
|
173 |
-
|
174 |
-
# Function to download the folders
|
175 |
-
def download_folder(folder):
|
176 |
-
zip_path = os.path.join(tempfile.gettempdir(), f"{folder}.zip")
|
177 |
-
|
178 |
-
# Zip the folder
|
179 |
-
shutil.make_archive(zip_path.replace('.zip', ''), 'zip', folder)
|
180 |
-
|
181 |
-
return zip_path
|
182 |
-
|
183 |
-
# Gradio Interface
|
184 |
-
def interface(name, age, medical_record, sex, input_image):
|
185 |
-
if input_image is None:
|
186 |
-
return None, "Please upload an image.", None
|
187 |
-
|
188 |
-
output_image, raw_result = predict_image(input_image, name, age, medical_record, sex)
|
189 |
-
html_file = save_patient_info_to_html(name, age, medical_record, sex, raw_result)
|
190 |
-
|
191 |
-
return output_image, raw_result, html_file
|
192 |
-
|
193 |
-
# Download Functions
|
194 |
-
def download_predicted_folder():
|
195 |
-
return download_folder(predicted_folder)
|
196 |
-
|
197 |
-
def download_uploaded_folder():
|
198 |
-
return download_folder(uploaded_folder)
|
199 |
-
|
200 |
-
# Gradio Blocks
|
201 |
-
with gr.Blocks() as demo:
|
202 |
-
with gr.Column():
|
203 |
-
gr.Markdown("# Cataract Detection System")
|
204 |
-
gr.Markdown("Upload an image to detect cataract and add patient details.")
|
205 |
-
gr.Markdown("This application uses YOLOv8 with mAP=0.981")
|
206 |
-
|
207 |
-
with gr.Column():
|
208 |
-
name = gr.Textbox(label="Name")
|
209 |
-
age = gr.Number(label="Age")
|
210 |
-
medical_record = gr.Number(label="Medical Record")
|
211 |
-
sex = gr.Radio(["Male", "Female"], label="Sex")
|
212 |
-
input_image = gr.Image(type="pil", label="Upload an Image", image_mode="RGB")
|
213 |
-
|
214 |
-
with gr.Column():
|
215 |
-
submit_btn = gr.Button("Submit")
|
216 |
-
output_image = gr.Image(type="pil", label="Predicted Image")
|
217 |
-
|
218 |
-
with gr.Row():
|
219 |
-
raw_result = gr.Textbox(label="Prediction Result")
|
220 |
-
|
221 |
-
with gr.Row():
|
222 |
-
download_html_btn = gr.Button("Download Patient Information (HTML)")
|
223 |
-
download_uploaded_btn = gr.Button("Download Uploaded Images")
|
224 |
-
download_predicted_btn = gr.Button("Download Predicted Images")
|
225 |
-
|
226 |
-
# Add file download output components for the uploaded and predicted images
|
227 |
-
patient_info_file = gr.File(label="Patient Information HTML File")
|
228 |
-
uploaded_folder_file = gr.File(label="Uploaded Images Zip File")
|
229 |
-
predicted_folder_file = gr.File(label="Predicted Images Zip File")
|
230 |
-
|
231 |
-
submit_btn.click(fn=interface, inputs=[name, age, medical_record, sex, input_image], outputs=[output_image, raw_result])
|
232 |
-
download_html_btn.click(fn=save_patient_info_to_html, inputs=[name, age, medical_record, sex, raw_result], outputs=patient_info_file)
|
233 |
-
download_uploaded_btn.click(fn=download_uploaded_folder, outputs=uploaded_folder_file)
|
234 |
-
download_predicted_btn.click(fn=download_predicted_folder, outputs=predicted_folder_file)
|
235 |
-
|
236 |
-
# Launch Gradio app
|
237 |
-
demo.launch()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|