ariankhalfani's picture
Update app.py
99de52c verified
raw
history blame
9.09 kB
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import sqlite3
import base64
from io import BytesIO
import tempfile
import pandas as pd
# Load YOLOv8 model
model = YOLO("best.pt")
def predict_image(input_image, name, patient_id):
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, patient_id, label)
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, patient_id, label):
draw = ImageDraw.Draw(image)
# Load a larger font (adjust the size as needed)
font_size = 48 # 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}, ID: {patient_id}, 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 initialize the database
def init_db():
conn = sqlite3.connect('results.db')
c = conn.cursor()
c.execute('''CREATE TABLE IF NOT EXISTS results
(id INTEGER PRIMARY KEY, name TEXT, patient_id TEXT, input_image BLOB, predicted_image BLOB, result TEXT)''')
conn.commit()
conn.close()
# Function to submit result to the database
def submit_result(name, patient_id, input_image, predicted_image, result):
conn = sqlite3.connect('results.db')
c = conn.cursor()
input_image_np = np.array(input_image)
_, input_buffer = cv2.imencode('.png', cv2.cvtColor(input_image_np, cv2.COLOR_RGB2BGR))
input_image_bytes = input_buffer.tobytes()
predicted_image_np = np.array(predicted_image)
predicted_image_rgb = cv2.cvtColor(predicted_image_np, cv2.COLOR_RGB2BGR) # Ensure correct color conversion
_, predicted_buffer = cv2.imencode('.png', predicted_image_rgb)
predicted_image_bytes = predicted_buffer.tobytes()
c.execute("INSERT INTO results (name, patient_id, input_image, predicted_image, result) VALUES (?, ?, ?, ?, ?)",
(name, patient_id, input_image_bytes, predicted_image_bytes, result))
conn.commit()
conn.close()
return "Result submitted to database."
# Function to load and view database
def view_database():
conn = sqlite3.connect('results.db')
c = conn.cursor()
c.execute("SELECT name, patient_id, input_image, predicted_image, result FROM results")
rows = c.fetchall()
conn.close()
# Convert to pandas DataFrame for better display in Gradio
df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Input Image", "Predicted Image", "Raw Result"])
# Decode images from BLOB and convert to displayable format
def decode_image(image_blob):
image_np = np.frombuffer(image_blob, dtype=np.uint8)
image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
return image
df["Input Image"] = df["Input Image"].apply(lambda x: decode_image(x))
df["Predicted Image"] = df["Predicted Image"].apply(lambda x: decode_image(x))
return df
# Function to download database or image
def download_file(choice):
if choice == "Database (.db)":
# Provide the path to the database file
return 'results.db'
elif choice == "Database (.html)":
conn = sqlite3.connect('results.db')
c = conn.cursor()
c.execute("SELECT name, patient_id, input_image, predicted_image, result FROM results")
rows = c.fetchall()
conn.close()
df = pd.DataFrame(rows, columns=["Name", "Patient ID", "Input Image", "Predicted Image", "Raw Result"])
html = df.to_html()
with tempfile.NamedTemporaryFile(delete=False, suffix='.html') as f:
f.write(html.encode())
return f.name
else:
# Handle other download options if necessary
pass
# Initialize the database
init_db()
# Define the Gradio interface
with gr.Blocks() as demo:
with gr.Tab("YOLOv8 Inference"):
with gr.Row():
input_image = gr.Image(label="Input Image", type="pil")
with gr.Row():
name = gr.Textbox(label="Patient Name")
patient_id = gr.Textbox(label="Patient ID")
with gr.Row():
submit_button = gr.Button("Submit")
predicted_image = gr.Image(label="Predicted Image")
with gr.Row():
result = gr.Textbox(label="Raw Result", lines=5)
submit_button.click(predict_image, inputs=[input_image, name, patient_id], outputs=[predicted_image, result])
with gr.Tab("View Database"):
view_button = gr.Button("View Database")
database_output = gr.DataFrame(label="Database Records")
view_button.click(view_database, outputs=database_output)
download_choice = gr.Radio(["Database (.db)", "Database (.html)", "Predicted Image (.png)"], label="Choose the file to download:")
download_button = gr.Button("Download")
download_button.click(download_file, inputs=download_choice, outputs=gr.File())
# Launch the interface
demo.launch()