import cv2
from sahi.predict import get_sliced_prediction
from sahi.utils.cv import read_image
from sahi import AutoDetectionModel
import torch

# Initialize Sahi and YOLOv8 models
yolov8_model_path = "C:/Users/moham/OneDrive/Desktop/Sliced_video/best_ATM.pt"

# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# Load YOLOv8 model
detection_model = AutoDetectionModel.from_pretrained(
    model_type='yolov8',
    model_path=yolov8_model_path,
    confidence_threshold=0.5,
    device=device
)

# Function to process video with Sahi
def process_video_with_sahi(video_path, output_dir):
    cap = cv2.VideoCapture(video_path)
    frame_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Convert frame to RGB
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Get sliced prediction
        result = get_sliced_prediction(
            frame_rgb,
            detection_model,
            slice_height=640,
            slice_width=512,
            overlap_height_ratio=0.2,
            overlap_width_ratio=0.2
        )

        # Extract bounding boxes, classes, names, and confidences
        boxes = result[0].boxes.xyxy.tolist()
        classes = result[0].boxes.cls.tolist()
        names = result[0].names
        confidences = result[0].boxes.conf.tolist()

        # Iterate through the results
        for box, cls, conf in zip(boxes, classes, confidences):
            x1, y1, x2, y2 = box
            confidence = conf
            detected_class = cls
            name = names[int(cls)]

            # Example: print bounding box coordinates, class, confidence, and name
            print(f"Bounding box: [{x1}, {y1}, {x2}, {y2}], Class: {detected_class}, Confidence: {confidence}, Name: {name}")

        # Export visuals
        output_path = f"{output_dir}/frame_{frame_count}.jpg"
        result.export_visuals(export_dir=output_path)

        frame_count += 1

    # Release resources
    cap.release()


# Example usage
if __name__ == "__main__":
    input_video_path = 'C:/Users/moham/OneDrive/Desktop/VSDRONE 19K/flag.mp4'
    output_directory = 'C:/Users/moham/OneDrive/Desktop/Sliced_video/sliced_frames'
    process_video_with_sahi(input_video_path, output_directory)