bacancydataprophets/Metal_Defect_Detection

1import pandas as pd2from ultralytics import YOLO3 4class Classification:5    def __init__(self) -> None:6        self.__cls_model = YOLO('model/cls_best.pt')7 8    def classify_defect(self, image_path) -> pd.DataFrame:9        result_cls = self.__cls_model.predict(image_path, stream=False)10        # Prepare data for CSV11        data1 = []12        for result in result_cls:13            cnt1 = 014            for i in result_cls[0].probs.top5:15                data1.append({16                    "Image/File Name": result_cls[0].path,17                    "Detected class by cls": self.__cls_model.names[i],18                    "Conf score": result_cls[0].probs.top5conf.tolist()[cnt1]19                    })20                cnt1 = cnt1 + 121        22        # Convert to DataFrame and save as CSV23        return pd.DataFrame(data1)24        # df1.to_csv('classification_results.csv', index=False)