from typing import NamedTuple
import warnings
warnings.filterwarnings("ignore")
from towhee.utils.ndarray_utils import to_ndarray
from deepface import DeepFace
import numpy as np
from towhee.types.image import Image
from towhee.operator import NNOperator
from towhee import register
from towhee.types.image_utils import to_pil
from towhee._types import Image
from towhee.types import arg, to_image_color


@register(output_schema=['vec'])
class DeepfaceFaceEmbedding(NNOperator):
  def __init__(self, model_name: str) -> None:
      super().__init__()
      self.model_name=model_name
  @arg(1, to_image_color('RGB'))
  def __call__(self,img: Image) -> np.ndarray:
      open_cv_image = np.array(img) 
      # Convert RGB to BGR 
      open_cv_image = open_cv_image[:, :, ::-1].copy() 
      # img.to_ndarray()
      embedding=DeepFace.represent(open_cv_image, model_name = self.model_name)
      embedding=np.array(embedding)
      return embedding
  def train(self):
          """
          For training model
          """
          pass