import torch from ultralytics import YOLO from faster_whisper import WhisperModel from pyannote.audio import Pipeline from google.cloud import translate_v2 as translate from speechbrain.inference.speaker import SpeakerRecognition from tools.tools import UniverseToolsClass ModelsUBL_process = 'cpu' # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device ModelsBAT_process = 'cpu' ModelFace_process = 'cpu' ModelBATHotspot_process = 'cpu' ModelBevarageIndrasty_process = 'cpu' ModelMaterial_process = 'cpu' ModelMemoOcr_process = 'cpu' ModelIntImgAna_process = None ModelSpeechAnalyzer_process = None class ModelsUBL: def __init__(self,device=ModelsUBL_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.daModel = YOLO("AI_Models/UBL_Models/ublDA_v8.1.pt").to(self.device) self.qpdsModel = YOLO("AI_Models/UBL_Models/ublQPDS_2_v1.2.pt").to(self.device) self.sosModel = YOLO("AI_Models/UBL_Models/ublSOS_v4.5.pt").to(self.device) self.mtSOSModel = YOLO("AI_Models/UBL_Models/ublMTSOS_v1.5.pt").to(self.device) self.sachetModel = YOLO("AI_Models/UBL_Models/sachetModel_v4.2.pt").to(self.device) self.megaHangerModel = YOLO("AI_Models/UBL_Models/megaHanger_v1.2.pt").to(self.device) self.sovmModel = YOLO("AI_Models/UBL_Models/sovmPOSM_v1.3.pt").to(self.device) self.mPOSM = YOLO("AI_Models/UBL_Models/POSM1_POSM2_FAT_v15.1.pt").to(self.device) self.st_orientation = YOLO("AI_Models/UBL_Models/st_orientation_v5.1.pt").to(self.device) self.exclusivity_model = YOLO("AI_Models/UBL_Models/exclusivity_v6.3.pt").to(self.device) self.sosModel_square = YOLO("AI_Models/UBL_Models/sosSquare_v1.2.pt").to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelsBAT: def __init__(self,device=ModelsBAT_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.blanksModel = YOLO('AI_Models/BAT_Models/blanks_ghw_v8.1.pt').to(self.device) self.posmModel = YOLO('AI_Models/BAT_Models/batPOSMall_v3.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelFace: def __init__(self,device=ModelFace_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.faceModel = YOLO('AI_Models/Face_Models/faceModel.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelBATHotspot: def __init__(self,device=ModelBATHotspot_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.BATHotspotModel = YOLO('AI_Models/BAT_Models/batPOSMall_v3.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelBevarageIndrasty: def __init__(self,device=ModelBevarageIndrasty_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.FridgeModel_fdz = YOLO('AI_Models/Fridge_Models/fdzModel.pt').to(self.device) self.FridgeModel_count = YOLO('AI_Models/Fridge_Models/countModel.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelMaterial: def __init__(self,device=ModelMaterial_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.MaterialGpModel_gp = YOLO('AI_Models/Material_Models/gp_v1.pt').to(self.device) self.MaterialGpModel_marchent = YOLO('AI_Models/Material_Models/marchent_v13.1.pt').to(self.device) self.MaterialGpModel_uddokta = YOLO('AI_Models/Material_Models/uddokta_v13.1.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelMemoOcr: def __init__(self,device=ModelMemoOcr_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.MemoOcrModel_crop = YOLO('AI_Models/MemoOCR_Models/cropModel.pt').to(self.device) self.MemoOcrModel_n_item = YOLO('AI_Models/MemoOCR_Models/n_item.pt').to(self.device) self.MemoOcrModel_n_price = YOLO('AI_Models/MemoOCR_Models/n_price.pt').to(self.device) self.MemoOcrModel_n_qty = YOLO('AI_Models/MemoOCR_Models/n_qty.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelIntImgAna: def __init__(self,device=ModelIntImgAna_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.ClasfiModel = YOLO('AI_Models/IntImgAna_Models/classification.pt').to(self.device) self.IntImgAnaModel = YOLO('AI_Models/IntImgAna_Models/prego_v4.1.pt').to(self.device) self.IntImgAnaFrameModel = YOLO('AI_Models/IntImgAna_Models/prego_frame_v3.1.pt').to(self.device) self.HeinekenModel = YOLO('AI_Models/IntImgAna_Models/heineken_sku.pt').to(self.device) self.HeinekenFrameModel = YOLO('AI_Models/IntImgAna_Models/heineken_frame.pt').to(self.device) self.HeinekenPlanogramModel = YOLO('AI_Models/IntImgAna_Models/heineken_planogram.pt').to(self.device) if 'cuda' in self.device: torch.cuda.empty_cache() class ModelMthouseOcr: def __init__(self,device=ModelIntImgAna_process): self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') self.PageCropModel = YOLO('AI_Models/MtHouseOCR_Model/page.pt').to(self.device) class ModelSpeechAnalyzer: def __init__(self,device=ModelSpeechAnalyzer_process): # Determine target device self.device = ( "cpu" if device == "cpu" else device if device else ("cuda" if torch.cuda.is_available() else "cpu") ) self.torch_device = torch.device(self.device) # Hugging Face token: environment override if set self.tools_universe = UniverseToolsClass() self.hf_token = self.tools_universe.decrypt_token() if not self.hf_token: print("⚠️ No Hugging Face token provided — gated models may fail to load.") # ----------------------- # 1️⃣ Whisper model (Speech-to-Text) # ----------------------- compute_type = "int8" if self.device == "cpu" else "int8_float32" self.whisper_model = WhisperModel( "large-v3", device=self.device, compute_type=compute_type ) # ----------------------- # 2️⃣ Speaker Verification (Speaker Embeddings) # ----------------------- self.speaker_verification_model = SpeakerRecognition.from_hparams( source="speechbrain/spkrec-ecapa-voxceleb", savedir="AI_Models/SpeechAnalyzer_Models/pretrained_models/spkrec", run_opts={"device": self.device} ) # ----------------------- # 3️⃣ Speaker Diarization (Who Spoke When) # ----------------------- try: self.diarization_pipeline = Pipeline.from_pretrained( "pyannote/speaker-diarization-community-1", token=self.hf_token ).to(self.torch_device) print(f"✅ Diarization pipeline loaded on {self.device}") except Exception as e: print(f"❌ Diarization pipeline failed: {e}") self.diarization_pipeline = None # ----------------------- # 4️⃣ translate # ----------------------- self.translate_client = translate.Client.from_service_account_json("tools/service_account.json") # ----------------------- # 5️⃣ Optional GPU memory cleanup # ----------------------- if self.device == "cuda": torch.cuda.empty_cache() print("🧠 GPU memory cache cleared.") print("🧠 GPU memory cache cleared.") # class ModelsUBL: # def __init__(self,device=ModelsUBL_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.daModel = None # self.qpdsModel = None # self.sosModel = None # self.mtSOSModel = None # self.sachetModel = None # self.megaHangerModel = None # self.sovmModel = None # self.mPOSM = None # self.st_orientation = None # self.exclusivity_model = None # self.sosModel_square = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelsBAT: # def __init__(self,device=ModelsBAT_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.blanksModel = None # self.posmModel = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelFace: # def __init__(self,device=ModelFace_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.faceModel = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelBATHotspot: # def __init__(self,device=ModelBATHotspot_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.BATHotspotModel = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelBevarageIndrasty: # def __init__(self,device=ModelBevarageIndrasty_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.FridgeModel_fdz = None # self.FridgeModel_count = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelMaterial: # def __init__(self,device=ModelMaterial_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.MaterialGpModel_gp = None # self.MaterialGpModel_marchent = None # self.MaterialGpModel_uddokta = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelMemoOcr: # def __init__(self,device=ModelMemoOcr_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.MemoOcrModel_crop = None # self.MemoOcrModel_n_item = None # self.MemoOcrModel_n_price = None # self.MemoOcrModel_n_qty = None # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelIntImgAna: # def __init__(self,device=ModelIntImgAna_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.ClasfiModel = YOLO('AI_Models/IntImgAna_Models/classification.pt').to(self.device) # self.IntImgAnaModel = YOLO('AI_Models/IntImgAna_Models/prego_v4.1.pt').to(self.device) # self.IntImgAnaFrameModel = YOLO('AI_Models/IntImgAna_Models/prego_frame_v3.0.pt').to(self.device) # self.HeinekenModel = YOLO('AI_Models/IntImgAna_Models/heineken_sku.pt').to(self.device) # self.HeinekenFrameModel = YOLO('AI_Models/IntImgAna_Models/heineken_frame.pt').to(self.device) # self.HeinekenPlanogramModel = YOLO('AI_Models/IntImgAna_Models/heineken_planogram.pt').to(self.device) # if 'cuda' in self.device: # torch.cuda.empty_cache() # class ModelMthouseOcr: # def __init__(self,device=ModelIntImgAna_process): # self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu') # self.PageCropModel = None # class ModelSpeechAnalyzer: # def __init__(self,device=ModelSpeechAnalyzer_process): # self.whisper_model = None # self.speaker_verification_model = None # self.diarization_pipeline = None # self.translate_client = None