Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image, ImageOps, ImageEnhance | |
| import numpy as np | |
| import tempfile, io, base64 | |
| # Load models | |
| model_swelling = YOLO("model/swelling/best.pt") | |
| model_redness = YOLO("model/redness/best.pt") | |
| model_bleeding = YOLO("model/bleeding/best.pt") | |
| # --- Helpers --- | |
| def preprocess(image): | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| image = ImageOps.exif_transpose(image).convert("RGB") | |
| # Resize if needed | |
| w, h = image.size | |
| max_dim = max(w, h) | |
| if max_dim > 1024: | |
| scale = 1024 / max_dim | |
| image = image.resize((int(w * scale), int(h * scale)), Image.LANCZOS) | |
| # Light contrast boost | |
| image = ImageEnhance.Contrast(image).enhance(1.05) | |
| return image | |
| def np_to_base64(img_np, format="JPEG"): | |
| """Convert numpy RGB image to Base64 string.""" | |
| pil_img = Image.fromarray(img_np) | |
| buffer = io.BytesIO() | |
| pil_img.save(buffer, format=format) | |
| return base64.b64encode(buffer.getvalue()).decode("utf-8") | |
| def base64_to_pil(b64_str): | |
| """Convert Base64 string back to PIL image (for Gradio display).""" | |
| return Image.open(io.BytesIO(base64.b64decode(b64_str))) | |
| # --- Main detection --- | |
| def detect_gingivitis(image, conf=0.4, iou=0.5): | |
| image = preprocess(image) | |
| sw_res = model_swelling.predict(image, conf=conf, iou=iou) | |
| rd_res = model_redness.predict(image, conf=conf, iou=iou) | |
| bl_res = model_bleeding.predict(image, conf=conf, iou=iou) | |
| # Convert YOLO output → numpy RGB | |
| img_sw = sw_res[0].plot()[:, :, ::-1] # BGR → RGB | |
| img_rd = rd_res[0].plot()[:, :, ::-1] | |
| img_bl = bl_res[0].plot()[:, :, ::-1] | |
| # Encode images to Base64 (for backend API consumption) | |
| sw_b64 = np_to_base64(img_sw) | |
| rd_b64 = np_to_base64(img_rd) | |
| bl_b64 = np_to_base64(img_bl) | |
| # Convert Base64 back to PIL for Gradio display | |
| sw_pil = base64_to_pil(sw_b64) | |
| rd_pil = base64_to_pil(rd_b64) | |
| bl_pil = base64_to_pil(bl_b64) | |
| # Determine diagnosis | |
| has_sw = len(sw_res[0].boxes) > 0 | |
| has_rd = len(rd_res[0].boxes) > 0 | |
| has_bl = len(bl_res[0].boxes) > 0 | |
| if has_sw and has_rd and has_bl: | |
| diagnosis = ( | |
| "🦷 You have gingivitis.\n\n" | |
| "For accurate assessment and guidance, we recommend visiting your dentist.\n" | |
| "If you have a periapical X-ray, you may try the *Detect Periodontitis* tool." | |
| ) | |
| else: | |
| diagnosis = "🟢 You don't have gingivitis." | |
| # Return PIL for Gradio + Base64 is available for backend | |
| return [sw_pil, rd_pil, bl_pil, diagnosis] | |
| # --- Gradio Interface --- | |
| interface = gr.Interface( | |
| fn=detect_gingivitis, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Image"), | |
| gr.Slider(0, 1, value=0.4, step=0.05, label="Confidence Threshold"), | |
| gr.Slider(0, 1, value=0.5, step=0.05, label="NMS IoU Threshold"), | |
| ], | |
| outputs=[ | |
| gr.Image(label="Swelling Detection", type="pil"), | |
| gr.Image(label="Redness Detection", type="pil"), | |
| gr.Image(label="Bleeding Detection", type="pil"), | |
| gr.Textbox(label="Diagnosis") | |
| ], | |
| title="Gingivitis Detection" | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch(server_name="0.0.0.0", server_port=7860, show_error=True) |