import rclpy
from rclpy.node import Node

from std_msgs.msg import String
from sensor_msgs.msg import Image

from ultralytics import YOLO
import cv2


class YoloDetect(Node):

    def __init__(self):
        super().__init__('YoloDetect')
        self.publisher_ = self.create_publisher(String, 'yolodetect', 10)
        self.publisher_img = self.create_publisher(Image, 'yolodetect', 10)
        timer_period = 0.5  # seconds
        self.timer = self.create_timer(timer_period, self.timer_callback)
        self.cap = cv2.VideoCapture("/mnt/c/Users/Bemposta/Videos/madrid.mp4")
        if not self.cap.isOpened():
            print("Cannot open video:")
            exit()
        self.model = YOLO("yolo26n.pt")

    def timer_callback(self):
        ret, frame = self.cap.read()
        results = self.model.predict(source=frame, iou=0.5, verbose=False, imgsz=[736, 1280])
        results_json =  results[0].to_json(normalize=True, decimals=5)

        msg = String()
        msg.data = results_json
        self.publisher_.publish(msg)
        cv2.imshow("detcatdo", results[0].plot())
        cv2.waitKey(1)
        self.get_logger().info('Publishing: "%s"' % msg.data)


def main(args=None):
    rclpy.init(args=args)

    yolo_detect = YoloDetect()

    rclpy.spin(yolo_detect)

    yolo_detect.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()