import torch
import cv2
import os
import numpy as np
from ultralytics import YOLO
from threading import Thread
import logging

from nodes.Node import Node
from nodes.YOLONode import YOLONode
import time

from multiprocessing import Process, Queue

from nodes.SORTTrackNode import SORTTrackNode

class FishYOLONode(YOLONode):

    def __init__(self,name="FishYOLONode",parent=None):
        super().__init__(name, parent)

    def run(self):
        device = torch.device("cuda:0")
        torch.cuda.set_device(device)
        self.open()

        self.is_stopped = False
        while True:
            if self.is_stopped:
                if self.is_opened:
                    self.close()
                break
            else:
                if( not self._framePool.empty()):

                    #try:
                    xtime = time.time()

                    frame,videoIndex,frameIndex = self._framePool.get()

                    #print("video Index:", videoIndex)

                    results = self.model(frame,imgsz=640,device=0)

                    endTime = time.time()
                    elapsedTime = endTime - xtime
                    #print(f"YoloTime:{ elapsedTime }")

                    if frame is not None:
                        self.passFrame([frame, videoIndex, frameIndex, results[0].boxes.cpu().numpy()])