Initialization

app.py

@@ -0,0 +1,260 @@
+import numpy as np
+import cv2
+import streamlit as st
+import pandas as pd
+import base64
+from openai import OpenAI
+import copy
+import time
+import io
+
+client = OpenAI(api_key='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')
+
+st.set_page_config(
+    page_title="Image to Analyze",
+    page_icon="📝",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+col1, col2 = st.columns(2)
+
+wl1, wl2 = st.columns(2)
+
+def to_base64(uploaded_file):
+    file_buffer = uploaded_file.read()
+    b64 = base64.b64encode(file_buffer).decode()
+    return f"data:image/png;base64,{b64}"
+
+def generate_df(new_img_list):
+
+    # if st.session_state.buttondo:
+    #     for img in new_img_list:
+    #         with open(img, mode="rb") as f:
+    #             uploaded_file = f.read()
+            
+    #         current_df = pd.DataFrame(
+    #             {
+    #                 "image_id": img,
+    #                 "image": f"data:image/png;base64,{base64.b64encode(uploaded_file).decode()}",
+    #                 "name": img,
+    #                 "defect_type": "",
+    #                 "description": "",
+    #             }
+    #         )
+    
+    # else:
+    current_df = pd.DataFrame(
+        {
+            "image_id": [img.file_id for img in st.session_state.images],
+            "image": [to_base64(img) for img in st.session_state.images],
+            "name": [img.name for img in st.session_state.images],
+            "defect_type": [""] * len(st.session_state.images),
+            "description": [""] * len(st.session_state.images),
+        }
+    )
+
+    if "df" not in st.session_state:
+        st.session_state.df = current_df
+        return
+
+    new_df = pd.merge(current_df, st.session_state.df, on=["image_id"], how="outer", indicator=True)
+    new_df = new_df[new_df["_merge"] != "right_only"].drop(columns=["_merge", "name_y", "image_y", "description_x"])
+    new_df = new_df.rename(columns={"name_x": "name", "image_x": "image", "description_y": "description"})
+    new_df["defect_type"] = new_df["defect_type"].fillna("")
+    new_df["description"] = new_df["description"].fillna("")
+
+    st.session_state.df = new_df
+
+
+def render_df():
+    def highlight_col(x):
+        r = 'background-color: red'
+        df1 = pd.DataFrame('', index=x.index, columns=x.columns)
+        df1.iloc[:, 2] = r
+        return df1    
+    
+    st.session_state.df.style.apply(highlight_col, axis=None)
+
+    st.data_editor(
+        st.session_state.df,
+        column_config={
+            "image": st.column_config.ImageColumn(
+                "Preview Image", help="Image preview", width=100
+            ),
+            "name": st.column_config.Column("Name", help="Image name", width=200),
+            "defect_type": st.column_config.Column(
+                "Defect Type", help="Defect description", width=400
+            ),
+            "description": st.column_config.Column(
+                "Description", help="Image description", width=800
+            ),
+        },
+        disabled=True,
+        hide_index=True,
+        height=500,
+        column_order=["image", "name", "defect_type", "description"],
+        use_container_width=True,
+    )
+
+
+def generate_description(image_base64):
+    response = client.chat.completions.create(
+        model="gpt-4-vision-preview",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": """Please answer at the following format:
+                     Wafer Defection Type: <put the defect type here only>
+                     Description: <Analyze how the defect type occur>
+                     Solution: <Suggest solution to solve this in the future>"""},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": image_base64,
+                        },
+                    },
+                ],
+            }
+        ],
+        max_tokens=50,
+    )
+    return response.choices[0].message.content
+
+
+def update_df():
+    indexes = st.session_state.df[st.session_state.df["description"] == ""].index
+    for idx in indexes:
+        description = generate_description(st.session_state.df.loc[idx, "image"])
+        st.session_state.df.loc[idx, "defect_type"] = description.split("\n")[0]
+        st.session_state.df.loc[idx, "description"] = description.split("\n")[-1]
+
+#Creating title for Streamlit app
+st.title("Wafer Defect Detection with LLM Classification and Analyze")
+
+with col1:
+    st.components.v1.iframe("http://localhost:7456/", width=None, height=500, scrolling=False)
+
+with col2:
+    st.components.v1.iframe("http://localhost:7457/", width=None, height=500, scrolling=False)
+
+default_img = ""
+st.session_state.buttondo = False
+
+with st.sidebar:
+    st.title("Upload Your Images")
+    st.session_state.images = st.file_uploader(label=" ", accept_multiple_files=True)
+
+    genre = st.radio(
+        "Baseline",
+        ["Propose Solution", "Baseline 1", "Baseline 0.2", "Baseline 0.1"])
+    
+    Center = st.checkbox('Center')
+    if Center:
+        st.image('44.png', width=100)
+    
+    Donut = st.checkbox('Donut')
+    if Donut:
+        st.image('7316.png', width=100)
+    
+    EdgeLoc = st.checkbox('Edge-Loc')
+    if EdgeLoc:
+        st.image('36.png', width=100)
+    
+    EdgeRing = st.checkbox('Edge-Ring')
+    if EdgeRing:
+        st.image('100.png', width=100)
+    
+    Loc = st.checkbox('Loc')
+    if Loc:
+        st.image('19.png', width=100)
+    
+    NearFull = st.checkbox('Near-Full')
+    if NearFull:
+        st.image('929.png', width=100)
+    
+    NoDefect = st.checkbox('No Defect')
+    if NoDefect:
+        st.image('0.png', width=100)
+    
+    Random = st.checkbox('Random')
+    if Random:
+        st.image('602.png', width=100)
+    
+    Scratch = st.checkbox('Scratch')
+    if Scratch:
+        st.image('134.png', width=100)
+
+    if st.button("Detect", type="primary"):
+        st.session_state.buttondo = True
+
+#uploading file for processing
+# uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
+    
+if st.session_state.images or st.session_state.buttondo:
+    imgs = []
+    new_img_list = copy.deepcopy(st.session_state.images)
+    # if st.session_state.buttondo:
+    #     new_img_list = [default_img]
+    # else:
+    #     new_img_list = copy.deepcopy(st.session_state.images)
+
+    for img in new_img_list:
+        image = cv2.imdecode(np.frombuffer(img.read(), np.uint8), 1)
+        # if st.session_state.buttondo:
+        #     image = cv2.imread(img)
+        # else:
+        #     image = cv2.imdecode(np.frombuffer(img.read(), np.uint8), 1)
+
+        img = image.copy()
+        gray=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
+
+        blur=cv2.blur(gray,(10,10))
+
+        dst=cv2.fastNlMeansDenoising(blur,None,10,7,21)
+
+        _,binary=cv2.threshold(dst,127,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
+
+        kernel=np.ones((5,5),np.uint8)
+
+        erosion=cv2.erode(binary,kernel,iterations=1)
+        dilation=cv2.dilate(binary,kernel,iterations=1)
+
+        if (dilation==0).sum()>1:
+            contours,_=cv2.findContours(dilation,cv2.RETR_TREE,cv2.CHAIN_APPROX_NONE)
+            for i in contours:
+                if cv2.contourArea(i)<261121.0:
+                    cv2.drawContours(img,i,-1,(0,0,255),3)
+                cv2.putText(img,"",(30,40),cv2.FONT_HERSHEY_SIMPLEX,1,(0,0,255),2)
+        else:
+            cv2.putText(img, "Good wafer", (30, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
+        # st.image(image,caption="original image",channels="BGR")
+        # st.image(blur,caption="blur")
+        # st.image(binary,caption="binary")
+        # st.image(erosion,caption="erosion")
+        # st.image(dilation,caption="dilation")
+        imgs.append(img)
+
+    # st.image(imgs)
+
+    with st.spinner('Wait for it...'):
+        time.sleep(10)
+    st.success('Done!')
+
+    generate_df(new_img_list)
+    # print(st.session_state.images)
+    update_df()
+    
+    st.download_button(
+        "Download descriptions as CSV",
+        st.session_state.df.drop(['image', "image_id"], axis=1).to_csv(index=False),
+        "descriptions.csv",
+        "text/csv",
+        use_container_width=True
+    )
+
+    render_df()
+# cv2.waitKey(0)
+# cv2.destroyAllWindows()