Update vlog4chat.py

vlog4chat.py

@@ -1,21 +1,21 @@
 import os
-import cv2
-import pdb
+#import cv2
+#import pdb
 import sys
 import time
 import numpy as np
-from PIL import Image
-from transformers import logging
-logging.set_verbosity_error()
-
-from models.kts_model import VideoSegmentor
-from models.clip_model import FeatureExtractor
-from models.blip2_model import ImageCaptioner
-from models.grit_model import DenseCaptioner
-from models.whisper_model import AudioTranslator
-from models.gpt_model import LlmReasoner
+#from PIL import Image
+#from transformers import logging
+#logging.set_verbosity_error()
+
+#from models.kts_model import VideoSegmentor
+#from models.clip_model import FeatureExtractor
+#from models.blip2_model import ImageCaptioner
+#from models.grit_model import DenseCaptioner
+#from models.whisper_model import AudioTranslator
+#from models.gpt_model import LlmReasoner
 from utils.utils import logger_creator, format_time
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, LlamaTokenizer, AutoModelForSeq2SeqLM
+#from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, LlamaTokenizer, AutoModelForSeq2SeqLM
 
 import together
 import warnings
@@ -38,13 +38,13 @@ from langchain.memory import ConversationBufferMemory
 from langchain import LLMChain, PromptTemplate
 from paddleocr import PaddleOCR, draw_ocr
 
-sys.path.append('/root/autodl-tmp/recognize-anything')
+#sys.path.append('/root/autodl-tmp/recognize-anything')
 
-from ram.models import ram
-from ram.models import tag2text
-from ram import inference_ram as inference
+#from ram.models import ram
+#from ram.models import tag2text
+#from ram import inference_ram as inference
 #from ram import inference_tag2text as inference
-from ram import get_transform
+#from ram import get_transform
 
 warnings.filterwarnings("ignore", category=UserWarning)
 B_INST, E_INST = "[INST]", "[/INST]"
@@ -68,7 +68,7 @@ Chat History:
 Follow Up Input: {question}
 Standalone question:"""
 
-os.environ['HF_HOME'] = '/root/autodl-tmp/cache/'
+#os.environ['HF_HOME'] = '/root/autodl-tmp/cache/'
 os.environ["TOGETHER_API_KEY"] = "48bf2536f85b599c7d5d7f9921cc9ee7056f40ed535fd2174d061e1b9abcf8af"
 
 def get_prompt(instruction, new_system_prompt=DEFAULT_SYSTEM_PROMPT ):
@@ -142,29 +142,10 @@ class Vlogger4chat :
         self.tmp_dir = args.tmp_dir
         self.models_flag = False
         self.init_llm()
-        self.init_tag2txt()
         self.history = []
         if not os.path.exists(self.tmp_dir):
             os.makedirs(self.tmp_dir)
-        
-    def init_models(self): 
-        print('\033[1;34m' + "Welcome to the our Vlog toolbox...".center(50, '-') + '\033[0m')
-        print('\033[1;33m' + "Initializing models...".center(50, '-') + '\033[0m')
-        print('\033[1;31m' + "This may time-consuming, please wait...".center(50, '-') + '\033[0m')
-        self.ocr = PaddleOCR(lang='ch')  # need to run only once to download and load model into memory
-        self.feature_extractor = FeatureExtractor(self.args)
-        self.video_segmenter = VideoSegmentor(alpha=self.alpha, beta=self.beta)
-        self.image_captioner = ImageCaptioner(model_name=self.args.captioner_base_model, device=self.args.image_captioner_device)
-        self.dense_captioner = DenseCaptioner(device=self.args.dense_captioner_device)
-        self.audio_translator = AudioTranslator(model=self.args.audio_translator, device=self.args.audio_translator_device)
-        print('\033[1;32m' + "Model initialization finished!".center(50, '-') + '\033[0m')
-        # 翻译文档
-        # 初始化 tokenizer 和 model
-        model_name = 'Helsinki-NLP/opus-mt-en-zh' #'Helsinki-NLP/opus-mt-zh-en'
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        self.my_embedding = HuggingFaceEmbeddings(model_name='BAAI/bge-m3', model_kwargs={'device': 'cuda'} ,encode_kwargs={'normalize_embeddings': True}) 
-    
+            
     def init_llm(self):
         print('\033[1;33m' + "Initializing LLM Reasoner...".center(50, '-') + '\033[0m')
         self.llm = TogetherLLM(
@@ -173,31 +154,7 @@ class Vlogger4chat :
             max_tokens=768
         )
         print('\033[1;32m' + "LLM initialization finished!".center(50, '-') + '\033[0m')
-        
-        if not self.models_flag:
-            self.init_models()
-            self.models_flag = True
             
-    def init_tag2txt(self):
-        self.transform = get_transform(image_size=384)
-
-        # delete some tags that may disturb captioning
-        # 127: "quarter"; 2961: "back", 3351: "two"; 3265: "three"; 3338: "four"; 3355: "five"; 3359: "one"
-        delete_tag_index = [127,2961, 3351, 3265, 3338, 3355, 3359]
-
-        #######load model
-        #self.tag2txt_model = tag2text(pretrained='/root/autodl-tmp/recognize-anything/pretrained/tag2text_swin_14m.pth',
-        #                              image_size=384, vit='swin_b', delete_tag_index=delete_tag_index)
-        self.ram_model = ram(pretrained='/root/autodl-tmp/recognize-anything/pretrained/ram_swin_large_14m.pth',
-                             image_size=384,
-                             vit='swin_l')
-        #self.tag2txt_model.threshold = 0.68  # threshold for tagging
-        #self.tag2txt_model.eval()
-        self.ram_model.eval()
-
-        #self.tag2txt_model = self.tag2txt_model.to(device=self.args.dense_captioner_device)
-        self.ram_model = self.ram_model.to(device=self.args.dense_captioner_device)
-        
     def exist_videolog(self, video_id):
         if isinstance(self.data_dir, tuple):
             self.data_dir = self.data_dir[0]  # 或者根据实际情况选择合适的索引
@@ -239,61 +196,6 @@ class Vlogger4chat :
         if self.exist_videolog(video_id):
             return self.printlog(video_id)
 
-        if not self.models_flag:
-            self.init_models()
-            self.models_flag = True
-            
-        logger = logger_creator(video_id)
-        clip_features, video_length = self.feature_extractor(video_path, video_id)
-        seg_windows = self.video_segmenter(clip_features, video_length)
-        
-        cap = cv2.VideoCapture(video_path)
-        fps = cap.get(cv2.CAP_PROP_FPS)
-        audio_results = self.audio_translator(video_path)
-
-        for start_sec, end_sec in seg_windows:
-            middle_sec = (start_sec + end_sec) // 2
-            middle_frame_idx = int(middle_sec * fps)
-            cap.set(cv2.CAP_PROP_POS_FRAMES, middle_frame_idx)
-            ret, frame = cap.read()
-            
-            if ret:
-                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                image_caption = self.image_captioner.image_caption(frame)
-                dense_caption = self.dense_captioner.image_dense_caption(frame)
-                image = self.transform(Image.fromarray(frame)).unsqueeze(0).to(device=self.args.dense_captioner_device)
-                #tag2txt = inference(image, self.tag2txt_model, 'None')
-                ram = inference(image, self.ram_model)
-                audio_transcript = self.audio_translator.match(audio_results, start_sec, end_sec)
-                OCR_result = self.ocr.ocr(frame)
-                # 提取所有文本块中的所有行的文字
-                texts = []
-                for block in OCR_result:
-                    if block is not None:  # 检查 block 是否为 None
-                        for line in block:
-                            if line is not None:  # 检查 line 是否为 None
-                                text = line[1][0]  # 提取文字部分
-                                texts.append(text)
-                # 将列表中的所有文字合并成一个字符串
-                OCR_result_str = ' '.join(texts)
-                    
-                logger.info(f"When {format_time(start_sec)} - {format_time(end_sec)}")
-                chinese_image_caption = self.translate_text(image_caption, self.tokenizer, self.model)
-                #chinese_tag2txt = self.translate_text(tag2txt[2], self.tokenizer, self.model)
-                chinese_dense_caption = self.translate_text(dense_caption, self.tokenizer, self.model)
-                logger.info(f"我看到这些画面：\"{chinese_image_caption}\"")
-                #logger.info(f"我看见 {chinese_tag2txt}.")
-                logger.info(f"我发现这些内容：\"{chinese_dense_caption}\"")
-                logger.info(f"我检测到这些标签：\"{ram[1]}.\"")
-                logger.info(f"我识别到这些文字：\"{OCR_result_str}\"")
-                    
-                if len(audio_transcript) > 0:
-                    #english_audio_text = self.translate_text(audio_transcript, self.tokenizer, self.model)
-                    logger.info(f"我听到有人说：\"{audio_transcript}\"")
-                logger.info("\n")
-        
-        cap.release()
-        self.create_videolog(video_id)
         return self.printlog(video_id)
         
     def printlog(self, video_id):
@@ -303,17 +205,6 @@ class Vlogger4chat :
             for line in f:
                 log_list.append(line.strip())
         return log_list
-    
-    def translate_text(self, text, tokenizer, model):
-        # 编码文本
-        encoded_text = tokenizer.prepare_seq2seq_batch([text], return_tensors='pt')
-
-        # 生成翻译
-        translated = model.generate(**encoded_text)
-
-        # 解码翻译后的文本
-        translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
-        return translated_text
 
     def chat2video(self, question):
         print(f"Question: {question}")