Piyapawashe/casting_defect_detection_app

1import gradio as gr2import numpy as np3import tensorflow as tf4import joblib5import pandas as pd6import datetime7import cv28from PIL import Image9import matplotlib.pyplot as plt10from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Input11 12 13# ----------------------------------------------------14# ✅ Define CNN architecture and load weights15inputs = tf.keras.Input(shape=(128, 128, 1))16x = Conv2D(32, (3, 3), activation='relu')(inputs)17x = MaxPooling2D()(x)18x = Conv2D(64, (3, 3), activation='relu')(x)19x = MaxPooling2D()(x)20x = Flatten()(x)21x = Dense(64, activation='relu', name="features")(x)22outputs = Dense(1, activation='sigmoid')(x)23 24model = tf.keras.Model(inputs, outputs)25model.load_weights("cnn_model.keras")26 27# ✅ Feature extractor for clustering28feature_model = tf.keras.Model(inputs, model.get_layer("features").output)29 30# ✅ Load KMeans model31kmeans = joblib.load("kmeans_model.pkl")32defect_types = ['Gas Porosity', 'Shrinkage', 'Metallurgical', 'Mold', 'Pouring', 'Shape']33 34# ---------------------------35# ✅ Logging defect count36def log_prediction(predicted_label, defect_type):37    today = datetime.date.today().isoformat()38    try:39        df = pd.read_csv("defect_logs.csv")40    except:41        df = pd.DataFrame(columns=["date", "ok", "defective", "defect_type"])42 43    if today not in df["date"].values:44        df = pd.concat([df, pd.DataFrame([{"date": today, "ok": 0, "defective": 0, "defect_type": ""}])], ignore_index=True)45 46    if predicted_label == 1:47        df.loc[df["date"] == today, "ok"] += 148    else:49        df.loc[df["date"] == today, "defective"] += 150        df.loc[df["date"] == today, "defect_type"] += f"{defect_type},"51 52    df.to_csv("defect_logs.csv", index=False)53 54 55 56 57# ✅ Plot side-by-side bar chart58def plot_defect_graph():59    try:60        df = pd.read_csv("defect_logs.csv")61        df["date"] = pd.to_datetime(df["date"])62        df["month"] = df["date"].dt.strftime("%b %Y")63        monthly_summary = df.groupby("month")[["ok", "defective"]].sum().reset_index()64 65        fig, ax = plt.subplots(figsize=(6, 3))66        x = np.arange(len(monthly_summary["month"]))67        width = 0.3568 69        ax.bar(x - width/2, monthly_summary["ok"], width, label="OK", color="#43a047")70        ax.bar(x + width/2, monthly_summary["defective"], width, label="Defective", color="#e53935")71 72        ax.set_xticks(x)73        ax.set_xticklabels(monthly_summary["month"], rotation=45)74        ax.set_title("Monthly Casting Summary")75        ax.set_ylabel("Count")76        ax.legend()77        plt.tight_layout()78        return fig79 80    except Exception as e:81        print("Graph error:", e)82        return plt.figure()83 84# ---------------------------85# ✅ Main Prediction Function86def classify(image):87    try:88        img = image.convert("L").resize((128, 128))89        img_array = np.array(img).reshape(1, 128, 128, 1) / 255.090 91        pred = model.predict(img_array)[0][0]92        predicted_label = 1 if pred > 0.5 else 093 94        features = feature_model.predict(img_array).reshape(1, -1)95        cluster = kmeans.predict(features)[0]96        defect_type = defect_types[cluster] if predicted_label == 0 else "None"97 98        log_prediction(predicted_label, defect_type)99 100        label = "✅ OK Casting" if predicted_label == 1 else "❌ Defective Casting"101        overlay_img = image.convert("RGB")102        graph = plot_defect_graph()103 104        return (105            label,106            f"Confidence: {pred:.2f}",107            defect_type,108            graph109        )110 111    except Exception as e:112        print("Prediction error:", e)113        return "Error", "Error", "Error", plt.figure()114 115# ---------------------------116# ✅ Gradio Interface117with gr.Blocks(css="""118    body { background-color: #ffffff; color: #000000; font-family: 'Segoe UI', sans-serif; }119    .gradio-container { padding: 30px; }120    .gr-box, .gr-image, .gr-textbox, .gr-plot {121        border: 2px solid #cccccc;122        border-radius: 8px;123        padding: 10px;124        background-color: #ffffff;125        color: #000000;126    }127    .gr-textbox input, .gr-textbox textarea {128        color: #000000 !important;129        background-color: #ffffff !important;130        border: 1px solid #cccccc;131    }132    .gr-markdown {133        color: #000000 !important;134    }135    .gr-textbox .label {136        color: #000000 !important;137    }138    .gr-button {139        background-color: #dddddd;140        color: #000000;141        font-weight: bold;142        border: 1px solid #999999;143    }144    #watermark {145        text-align: center;146        font-size: 12px;147        color: #000000 !important;148        margin-top: 30px;149    }150""") as demo:151    gr.Markdown("## **CASTING DEFECT DETECTOR**")152    gr.Markdown("""153    Upload a casting image to detect defects using machine learning.  154    🔍 **Grad-CAM** highlights regions influencing the prediction.  155    📊 Track daily defect trends and classification confidence.156    """)157 158    with gr.Row():159        image_input = gr.Image(type="pil", label="📷 Upload or Capture Casting Image")160        with gr.Column():161            prediction = gr.Text(label="🟢 Prediction")162            confidence = gr.Text(label="📊 Confidence Score")163            defect_type = gr.Text(label="🔍 Defect Type")164            submit_btn = gr.Button("🔍 Analyze Casting Image")165 166    graph_output = gr.Plot(label="📊 Monthly Casting Summary")167 168    submit_btn.click(169        fn=classify,170        inputs=image_input,171        outputs=[prediction, confidence, defect_type, graph_output]172    )173 174    gr.Markdown("© Priyanka_Pawashe | Powered by ML & Vision", elem_id="watermark")175 176demo.launch()177