johnwesley756/instance-segmentation

1import streamlit as st2from ultralytics import YOLO3import tempfile4import cv25import pandas as pd6import os7from datetime import datetime8from fpdf import FPDF9 10 11model = YOLO("best.pt")12 13 14st.set_page_config(page_title="Tooth Detection", layout="centered")15st.title("Tooth Detection with Severity, Summary & PDF")16 17def generate_summary(severity):18    if severity == "Cavity":19        return ("A cavity has been detected. This suggests advanced tooth decay "20                "that requires immediate dental attention to avoid complications.")21    elif severity == "Caries":22        return ("Caries (initial decay) detected. Early treatment such as fluoride "23                "application or filling may help prevent cavity formation.")24    elif severity == "Tooth":25        return ("No signs of decay detected. Maintain proper oral hygiene with regular "26                "brushing, flossing, and dental checkups.")27    else:28        return ("No objects detected. Please ensure the image is clear and properly focused. "29                "Try uploading a new image or consult a dental professional.")30 31def generate_pdf(severity, summary_text, image_path):32    pdf = FPDF()33    pdf.add_page()34    pdf.set_font("Arial", "B", 16)35    pdf.cell(0, 10, "Tooth Detection Report", ln=True, align='C')36    pdf.set_font("Arial", size=12)37    pdf.cell(0, 10, f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True)38    pdf.cell(0, 10, f"Severity Level: {severity}", ln=True)39    pdf.ln(5)40    pdf.multi_cell(0, 10, f"Summary:\n{summary_text}")41    if image_path:42        pdf.image(image_path, x=10, y=None, w=180)43    file_name = f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"44    pdf.output(file_name)45    return file_name46 47 48log_file = "severity_log.csv"49if not os.path.exists(log_file):50    pd.DataFrame(columns=["Timestamp", "Severity"]).to_csv(log_file, index=False)51 52uploaded_file = st.file_uploader("Upload a dental image", type=["jpg", "jpeg", "png"])53 54if uploaded_file:55    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:56        tmp.write(uploaded_file.read())57        image_path = tmp.name58 59    results = model.predict(image_path, conf=0.25)60    annotated_img = results[0].plot()61    img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)62 63    # Extract class names64    class_names = [results[0].names[int(cls.item())].lower() for cls in results[0].boxes.cls]65 66    # Severity logic67    if "cavity" in class_names:68        severity = "Cavity"69    elif "caries" in class_names:70        severity = "Caries"71    elif all(name == "tooth" for name in class_names) and class_names:72        severity = "Tooth"73    else:74        severity = "No Detection"75 76    # Generate summary77    summary_text = generate_summary(severity)78 79    # Display results80    st.image(img_rgb, caption="Detected Image", use_column_width=True)81    st.markdown(f"Severity Level: *{severity}*")82    st.markdown("GPT-style Dental Summary:")83    st.write(summary_text)84 85    # Save image86    img_save_path = f"result_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"87    cv2.imwrite(img_save_path, annotated_img)88 89    # Generate and download PDF90    pdf_file = generate_pdf(severity, summary_text, img_save_path)91    with open(pdf_file, "rb") as f:92        st.download_button("Download PDF Report", data=f, file_name=pdf_file, mime="application/pdf")93 94    # Log result95    log_df = pd.read_csv(log_file)96    log_df.loc[len(log_df)] = [datetime.now().strftime("%Y-%m-%d %H:%M:%S"), severity]97    log_df.to_csv(log_file, index=False)