From b9e9878f418fd0d9826445c0eab11104520b6373 Mon Sep 17 00:00:00 2001 From: nfilin480 Date: Sun, 15 Oct 2023 01:41:11 +0300 Subject: [PATCH] add multiply videos --- model/generator/main.py | 17 ++++++++-------- model/preprocessor/libs/config.py | 12 ++++++++++++ model/preprocessor/libs/models.py | 32 ++++++++++++++++++++----------- model/preprocessor/main.py | 31 ++++++++++++++++++++---------- 4 files changed, 63 insertions(+), 29 deletions(-) create mode 100644 model/preprocessor/libs/config.py diff --git a/model/generator/main.py b/model/generator/main.py index 72d7625..e32a63f 100644 --- a/model/generator/main.py +++ b/model/generator/main.py @@ -6,13 +6,14 @@ from libs import logger from libs.models import ImageGenerationModel import urllib.request +from libs.config import settings class Pipeline: def __init__(self) -> None: - self.pb: PocketBase = PocketBase("https://pb.apps.npod.space/") + self.pb: PocketBase = PocketBase(settings.PB_LINK) self.admin_data = self.pb.admins.auth_with_password( - "dev@email.local", - "6c6297287af76472") + settings.PB_LOGIN, + settings.PB_PWD) self.model = ImageGenerationModel() self.model.load_kandinsky() @@ -21,7 +22,7 @@ def run(self): while True: # logger.info('listening to the collection') try: - records = self.pb.collection('text_generation_mvp').get_list( + records = self.pb.collection(settings.COL_NAME).get_list( page=1, per_page=1, query_params={ @@ -47,7 +48,7 @@ def run(self): images = self.process_generation(records.items[0], width=width, height=height) if len(images) == 0: - self.pb.collection('text_generation_mvp').update( + self.pb.collection(settings.COL_NAME).update( id=records.items[0].id, body_params={ "status": "error" @@ -56,7 +57,7 @@ def run(self): logger.error('Не было создано ни одного изображения') continue - self.pb.collection('text_generation_mvp').update( + self.pb.collection(settings.COL_NAME).update( id=records.items[0].id, body_params={ "status": "generated", @@ -74,7 +75,7 @@ def run(self): image.save(path) uploads.append((path, open(path, mode='rb'))) - self.pb.collection('text_generation_mvp').update( + self.pb.collection(settings.COL_NAME).update( id=records.items[0].id, body_params={ "status": "generated", @@ -92,7 +93,7 @@ def process_generation(self, data, width, height) -> list[Image.Image]: if data.input_image != '' and data.type == 'avatar': try: - image_url = self.pb.collection('text_generation_mvp').get_file_url(data, filename=data.input_image) + image_url = self.pb.collection(settings.COL_NAME).get_file_url(data, filename=data.input_image) urllib.request.urlretrieve(image_url, data.input_image) self.model.remove_models() diff --git a/model/preprocessor/libs/config.py b/model/preprocessor/libs/config.py new file mode 100644 index 0000000..cea2574 --- /dev/null +++ b/model/preprocessor/libs/config.py @@ -0,0 +1,12 @@ +from pydantic_settings import BaseSettings, SettingsConfigDict + +class Settings(BaseSettings): + model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8', extra='ignore') + +class ModuleSettings(Settings): + PB_LINK: str = "https://pb.apps.npod.space/" + PB_LOGIN: str = "dev@email.local" + PB_PWD: str = "6c6297287af76472" + COL_NAME: str = "text_generation_mvp" + +settings = ModuleSettings() \ No newline at end of file diff --git a/model/preprocessor/libs/models.py b/model/preprocessor/libs/models.py index 43781b1..dd2829f 100644 --- a/model/preprocessor/libs/models.py +++ b/model/preprocessor/libs/models.py @@ -99,20 +99,30 @@ def video_parser(self, video_file): video_clip.save_frame(frame_filename, current_duration) return folder_dir - def get_desc_from_video(self, video_file: str) -> str: - if not video_file.endswith('.mp4'): - return '' + def get_desc_from_video(self, video_files) -> str: + + result_summary = [] + for video_file in video_files: + if not video_file.endswith('.mp4'): + return '' - folder_dir = self.video_parser(video_file) - desc = [] - for images in os.listdir(folder_dir): - desc.append(self.generate_desc(folder_dir + '/' + images)) + folder_dir = self.video_parser(video_file) + desc = [] + for images in os.listdir(folder_dir): + desc.append(self.generate_desc(folder_dir + '/' + images)) - if os.path.exists(folder_dir): - shutil.rmtree(folder_dir) + if os.path.exists(folder_dir): + shutil.rmtree(folder_dir) - summary_desc = self.get_summary(desc) - return summary_desc + summary_desc = self.get_summary(desc) + result_summary.append(summary_desc) + + if len(result_summary) > 1: + return self.get_summary(result_summary) + elif len(result_summary) == 1: + return result_summary[0] + else: + return '' def get_summary(self, desc) -> str: diff --git a/model/preprocessor/main.py b/model/preprocessor/main.py index 5aa8cff..14f8cbd 100644 --- a/model/preprocessor/main.py +++ b/model/preprocessor/main.py @@ -9,13 +9,15 @@ from libs import logger from libs.models import Translator, VideoParser import urllib.request +from libs.config import settings +import re class Pipeline: def __init__(self) -> None: - self.pb: PocketBase = PocketBase("https://pb.apps.npod.space/") + self.pb: PocketBase = PocketBase(settings.PB_LINK) self.admin_data = self.pb.admins.auth_with_password( - "dev@email.local", - "6c6297287af76472") + settings.PB_LOGIN, + settings.PB_PWD) self.translator = Translator() self.vparser = VideoParser() @@ -23,7 +25,7 @@ def __init__(self) -> None: def run(self): while True: try: - records = self.pb.collection('text_generation_mvp').get_list( + records = self.pb.collection(settings.COL_NAME).get_list( page=1, per_page=1, query_params={ @@ -53,12 +55,21 @@ def run(self): style = records.items[0].style try: - video_url = self.pb.collection('text_generation_mvp').get_file_url(records.items[0], filename=records.items[0].video) + video_url = self.pb.collection(settings.COL_NAME).get_file_url(records.items[0], filename=records.items[0].video) - video_name = 'video.mp4' - urllib.request.urlretrieve(video_url, video_name) - - summary = self.vparser.get_desc_from_video(video_name)[0]['summary_text'] + s = re.findall('\[(.*)\]', url) + ss = re.sub('[\' ]', '', s[0]) + video_names = re.split(',', ss) + + url = re.sub('\[.*\]', '', url) + + #video_name = 'video.mp4' + + for video_name in video_names: + urllib.request.urlretrieve(video_url + video_name, video_name) + + + summary = self.vparser.get_desc_from_video(video_names)[0]['summary_text'] print(summary) if prompt != '': @@ -68,7 +79,7 @@ def run(self): except: logger.info('Нет приложенного видео') - self.pb.collection('text_generation_mvp').update( + self.pb.collection(settings.COL_NAME).update( id=records.items[0].id, body_params={ "status": "preprocessed",