100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
import urllib.request
|
|
import os
|
|
from pathlib import Path
|
|
from bingart import BingArt
|
|
|
|
from custom_apps.utils import cookie_manager
|
|
from main_rest.app.utils.parsing_utils import prompt_to_filenames
|
|
from main_rest.app.utils.date_utils import D
|
|
from const import TEMP_FOLDER
|
|
from utils.custom_sftp import sftp_client
|
|
from config import rest_config
|
|
|
|
|
|
class BingArtGenerator:
|
|
|
|
model = 'dalle3'
|
|
detail = 'art'
|
|
output_folder = os.path.join(TEMP_FOLDER,"dalle","art")
|
|
remote_folder = os.path.join(rest_config.remote_folder,"dalle","art")
|
|
|
|
def __init__(self):
|
|
self.bing_art = BingArt(auth_cookie_U=cookie_manager.get_cookie())
|
|
self.prompt = ''
|
|
self.datetime = D.date_file_name()
|
|
Path(self.output_folder).mkdir(parents=True, exist_ok=True)
|
|
|
|
def get_image_links(self, prompt, image_len:int):
|
|
"""
|
|
이미지 url값 생성
|
|
"""
|
|
results = None
|
|
try:
|
|
self.prompt = prompt_to_filenames(prompt)
|
|
results = self.bing_art.generate_images(prompt)
|
|
|
|
finally:
|
|
self.bing_art.close_session()
|
|
|
|
if results == None:
|
|
raise Exception("result is null")
|
|
|
|
images = results.get("images")
|
|
|
|
if images != None:
|
|
_temp = []
|
|
for i in images:
|
|
_temp.append(i.get('url'))
|
|
images = _temp
|
|
|
|
#중복제거
|
|
images = list(set(_temp))
|
|
|
|
if len(images) < image_len:
|
|
return images[0:len(images)]
|
|
else:
|
|
return images[0:image_len]
|
|
else:
|
|
return None
|
|
|
|
def link_to_img_remote(self, img_links:list):
|
|
"""
|
|
이미지 링크로 이미지 파일 원격지에 저장
|
|
"""
|
|
jpeg_index = 1
|
|
|
|
if img_links == None:
|
|
raise Exception("result is null")
|
|
|
|
for i in img_links:
|
|
path= f"{self.model}_{self.detail}_{self.prompt}_{jpeg_index}_{self.datetime}.png"
|
|
|
|
urllib.request.urlretrieve(i, os.path.join(self.output_folder,path))
|
|
sftp_client.remote_copy_data(os.path.join(self.output_folder,path), os.path.join(self.remote_folder,path))
|
|
jpeg_index += 1
|
|
|
|
def get_images(self, prompt, image_len):
|
|
image_links = self.get_image_links(prompt,image_len)
|
|
|
|
if image_links == None:
|
|
return 0
|
|
|
|
else:
|
|
self.link_to_img_remote(img_links=image_links)
|
|
return len(image_links)
|
|
|
|
def get_query_image(self,prompt):
|
|
|
|
create_time = D.date_file_name()
|
|
|
|
image_links = self.get_image_links(prompt,1)
|
|
|
|
if image_links == None:
|
|
return 0
|
|
else:
|
|
image = image_links[0]
|
|
query_path = os.path.join(TEMP_FOLDER,f"query_bingarg_{create_time}.png")
|
|
urllib.request.urlretrieve(image, query_path)
|
|
|
|
return query_path
|