Upload 7 files

app.py

@@ -0,0 +1,44 @@
+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()

classify.py

@@ -0,0 +1,24 @@
+import pandas as pd
+from ultralytics import YOLO
+
+class Classification:
+    def __init__(self) -> None:
+        self.__cls_model = YOLO('/home/pratham/MDD/bin/model/cls_best.pt')
+
+    def classify_defect(self, image_path) -> pd.DataFrame:
+        result_cls = self.__cls_model.predict(image_path, stream=False)
+        # Prepare data for CSV
+        data1 = []
+        for result in result_cls:
+            cnt1 = 0
+            for i in result_cls[0].probs.top5:
+                data1.append({
+                    "Image/File Name": result_cls[0].path,
+                    "Detected class by cls": self.__cls_model.names[i],
+                    "Conf score": result_cls[0].probs.top5conf.tolist()[cnt1]
+                    })
+                cnt1 = cnt1 + 1
+        
+        # Convert to DataFrame and save as CSV
+        return pd.DataFrame(data1)
+        # df1.to_csv('classification_results.csv', index=False)

compare.py

@@ -0,0 +1,22 @@
+import pandas as pd
+
+class Compare:
+    def comparison(self, df, df1):
+        final = df.merge(df1, on='Image/File Name')
+        final.drop_duplicates(subset=['Confidence Score'], keep='first', inplace=True)
+        final = final.reset_index().drop(columns=['index'])
+        
+        fnl_cls = []
+        fnl_cnf = []
+        for i in range(0, len(final['Confidence Score'])):
+            if final['Confidence Score'][i] > final['Conf score'][i]:
+                fnl_cls.append(final['Detected Class'][i])
+                fnl_cnf.append(final['Confidence Score'][i])
+            else:
+                fnl_cls.append(final['Detected class by cls'][i])
+                fnl_cnf.append(final['Conf score'][i])
+        
+        final['fnl_cls'] = fnl_cls
+        final['fnl_pred'] = fnl_cnf
+        final = final.drop(columns=['Detected Class', 'Confidence Score', 'Detected class by cls', 'Conf score'])
+        return final

detect.py

@@ -0,0 +1,28 @@
+import pandas as pd
+from ultralytics import YOLO
+
+class Detection:
+    def __init__(self) -> None:
+        self.__obd_model = YOLO('/home/pratham/MDD/bin/model/obd_best.pt')
+
+    def detect_defect(self, image_path) -> pd.DataFrame:
+        result_obd = self.__obd_model.predict(image_path, stream=False)  # Adjust paths as needed
+        
+        # Prepare data for CSV
+        data = []
+        for result in result_obd:
+            cnt = 0
+            for i in result_obd[0].boxes.cls.tolist():
+                data.append({
+                    "Image/File Name": result.path,
+                    "Detected Class": self.__obd_model.names[int(i)],
+                    "Confidence Score": result.boxes.conf.tolist()[cnt],
+                    "x1": result.boxes.xyxy.tolist()[cnt][0],
+                    "y1": result.boxes.xyxy.tolist()[cnt][1],
+                    "x2": result.boxes.xyxy.tolist()[cnt][2],
+                    "y2": result.boxes.xyxy.tolist()[cnt][3]
+                    })
+                cnt = cnt + 1
+        
+        # Convert to DataFrame and save as CSV
+        return pd.DataFrame(data)

model/cls_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeae51ca684bae2c7bc8ca51a9b0646b72467f8354c526a0c7584f7d3e7b3be3
+size 2973377

model/obd_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f423bbe8c4876ff23e4585399acce014b51cf3db3a094f494514a881bed8a99
+size 6256729

requirements.txt

@@ -0,0 +1,77 @@
+altair==5.3.0
+attrs==23.2.0
+blinker==1.8.2
+cachetools==5.3.3
+certifi==2024.2.2
+charset-normalizer==3.3.2
+click==8.1.7
+contourpy==1.2.1
+cycler==0.12.1
+dill==0.3.8
+filelock==3.14.0
+fonttools==4.51.0
+fsspec==2024.5.0
+gitdb==4.0.11
+GitPython==3.1.43
+idna==3.7
+Jinja2==3.1.4
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.0
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.3
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.1.105
+opencv-python==4.9.0.80
+packaging==24.0
+pandas==2.2.2
+pillow==10.3.0
+protobuf==4.25.3
+psutil==5.9.8
+py-cpuinfo==9.0.0
+pyarrow==16.1.0
+pydeck==0.9.1
+Pygments==2.18.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.35.1
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.1
+scipy==1.13.0
+seaborn==0.13.2
+six==1.16.0
+smmap==5.0.1
+streamlit==1.34.0
+sympy==1.12
+tenacity==8.3.0
+thop==0.1.1.post2209072238
+toml==0.10.2
+toolz==0.12.1
+torch==2.3.0
+torchvision==0.18.0
+tornado==6.4
+tqdm==4.66.4
+triton==2.3.0
+typing_extensions==4.11.0
+tzdata==2024.1
+ultralytics==8.2.16
+urllib3==2.2.1
+watchdog==4.0.0