import streamlit as st
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from detect import Detection
from classify import Classification
from compare import Compare

# Streamlit app
def main():
    st.title("Metal Defect Detection and Classification App")
    det_res = Detection()
    cls_res = Classification()
    comp = Compare()
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

    if uploaded_file is not None:
        # Display uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        df = det_res.detect_defect(image)
        df1 = cls_res.classify_defect(image)

        # Perform comparison between scores of detection and classification
        detection_results = comp.comparison(df, df1)

        # Display results
        fig, ax = plt.subplots(1)
        ax.imshow(image)

        for index, row in detection_results.iterrows():
            x1, y1, x2, y2 = row['x1'], row['y1'], row['x2'], row['y2']
            width, height = x2 - x1, y2 - y1
            rect = patches.Rectangle((x1, y1), width, height, linewidth=1, edgecolor='r', facecolor='none')
            ax.add_patch(rect)
            ax.text(x1, y1 - 5, f"{row['fnl_cls']}: {row['fnl_pred']:.2f}", color='r')

        ax.axis('off')
        st.pyplot(fig)

# Run the app
if __name__ == "__main__":
    main()