import pandas as pd
from ultralytics import YOLO

class Detection:
    def __init__(self) -> None:
        self.__obd_model = YOLO('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)