Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image, ImageOps, ImageEnhance | |
| import numpy as np | |
| import io, base64 | |
| # ========================================================= | |
| # Lazy-loaded global models (LOAD ONLY ON FIRST REQUEST) | |
| # ========================================================= | |
| model_swelling = None | |
| model_redness = None | |
| model_bleeding = None | |
| def get_models(): | |
| """Load YOLO models only once (lazy loading).""" | |
| global model_swelling, model_redness, model_bleeding | |
| if model_swelling is None: | |
| model_swelling = YOLO("models/swelling/swelling_final.pt") | |
| if model_redness is None: | |
| model_redness = YOLO("models/redness/redness_final.pt") | |
| if model_bleeding is None: | |
| model_bleeding = YOLO("models/bleeding/bleeding_final.pt") | |
| return model_swelling, model_redness, model_bleeding | |
| # ========================================================= | |
| # Helper functions | |
| # ========================================================= | |
| def preprocess(image): | |
| """Resize, fix orientation, improve contrast.""" | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| image = ImageOps.exif_transpose(image).convert("RGB") | |
| # Resize if too large | |
| 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) | |
| # Slight contrast enhancement | |
| image = ImageEnhance.Contrast(image).enhance(1.05) | |
| return image | |
| def np_to_base64(img_np, format="JPEG"): | |
| """Convert numpy RGB image to Base64.""" | |
| 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 to PIL image.""" | |
| return Image.open(io.BytesIO(base64.b64decode(b64_str))) | |
| # ========================================================= | |
| # Main detection function | |
| # ========================================================= | |
| def detect_gingivitis(image, conf=0.25, iou=0.5): | |
| try: | |
| if image is None: | |
| return [None, None, None, "❌ No image uploaded"] | |
| # Load models (only once) | |
| sw_model, rd_model, bl_model = get_models() | |
| # Preprocess | |
| image = preprocess(image) | |
| # Run detections | |
| sw_res = sw_model.predict(image, conf=conf, iou=iou) | |
| rd_res = rd_model.predict(image, conf=conf, iou=iou) | |
| bl_res = bl_model.predict(image, conf=conf, iou=iou) | |
| # Convert YOLO output → numpy → PIL | |
| img_sw = sw_res[0].plot(labels=False)[:, :, ::-1] | |
| img_rd = rd_res[0].plot(labels=False)[:, :, ::-1] | |
| img_bl = bl_res[0].plot(labels=False)[:, :, ::-1] | |
| sw_pil = base64_to_pil(np_to_base64(img_sw)) | |
| rd_pil = base64_to_pil(np_to_base64(img_rd)) | |
| bl_pil = base64_to_pil(np_to_base64(img_bl)) | |
| # Diagnosis logic | |
| 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\n" | |
| "If you have a periapical X-ray, you may try the Detect Periodontitis tool." | |
| ) | |
| else: | |
| diagnosis = "🟢 You don't have gingivitis." | |
| return [sw_pil, rd_pil, bl_pil, diagnosis] | |
| except Exception as e: | |
| return [None, None, None, f"❌ Error during processing: {str(e)}"] | |
| # ========================================================= | |
| # Gradio Interface | |
| # ========================================================= | |
| interface = gr.Interface( | |
| fn=detect_gingivitis, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Image"), | |
| gr.Slider(0, 1, value=0.5, 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" | |
| ) | |
| # ========================================================= | |
| # Warm-start: preload models on startup | |
| # ========================================================= | |
| print("🔥 Preloading models to reduce Render cold start...") | |
| get_models() | |
| print("✅ Gingivitis models ready") | |
| # ========================================================= | |
| # Start server | |
| # ========================================================= | |
| if __name__ == "__main__": | |
| interface.launch(server_name="0.0.0.0", server_port=7860, show_error=True) |