edit : 원격 저장으로 변경

.gitignore

@@ -142,7 +142,7 @@ cython_debug/
 .idea/
 
 log
-sheet_counter.txt
 
 *output
+*temp
 datas/eyewear_all/*

const.py

@@ -1,4 +1,4 @@
-# OUTPUT_FOLDER = "/home/fermat/STORAGE/01.Projects/A2TEC/K_EYEWEAR/02.ML_DATA/Image_generator_result"
+REMOTE_FOLDER = "/home/fermat/STORAGE/01.Projects/A2TEC/K_EYEWEAR/02.ML_DATA/Image_generator_result"
-OUTPUT_FOLDER = "./output"
+TEMP_FOLDER = "./temp"
 
 ILLEGAL_FILE_NAME = ['<', '>', ':', '"', '/', '\ ', '|', '?', '*']

custom_apps/bingart/bingart.py

@@ -6,14 +6,16 @@ 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 OUTPUT_FOLDER
+from const import REMOTE_FOLDER, TEMP_FOLDER
+from utils.custom_sftp import sftp_client
 
 
 class BingArtGenerator:
     
     model = 'dalle3'
     detail = 'art'
-    output_folder = os.path.join(OUTPUT_FOLDER,"dalle","art")
+    output_folder = os.path.join(TEMP_FOLDER,"dalle","art")
+    remote_folder = os.path.join(REMOTE_FOLDER,"dalle","art")
 
     def __init__(self):
         self.bing_art = BingArt(auth_cookie_U=cookie_manager.get_cookie())
@@ -67,6 +69,7 @@ class BingArtGenerator:
             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):

custom_apps/bingimagecreator/utils.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from pathlib import Path
 
 from custom_apps.bingimagecreator.BingImageCreator import *
-from const import OUTPUT_FOLDER
+from const import REMOTE_FOLDER
 from custom_apps.utils import cookie_manager
 
 
@@ -24,7 +24,7 @@ class DallEArgument:
     U = cookie_manager.get_cookie()
     prompt: str
     cookie_file: str|None = None
-    output_dir: str = os.path.join(OUTPUT_FOLDER,"dalle","img")
+    output_dir: str = os.path.join(REMOTE_FOLDER,"dalle","img")
     download_count: int = 1
     debug_file: str|None = None
     quiet: bool = False

custom_apps/imagen/custom_imagen.py

@@ -4,10 +4,10 @@ import vertexai
 import os
 from vertexai.preview.vision_models import ImageGenerationModel
 
-from const import OUTPUT_FOLDER
+from const import REMOTE_FOLDER, TEMP_FOLDER
 from main_rest.app.utils.parsing_utils import prompt_to_filenames
 from main_rest.app.utils.date_utils import D
-
+from utils.custom_sftp import sftp_client
 
 class ImagenConst:
     project_id = "glasses-imagen"
@@ -21,11 +21,12 @@ def imagen_generate_image(prompt,download_count=1):
     model = ImageGenerationModel.from_pretrained(ImagenConst.model)
 
     _file_name = prompt_to_filenames(prompt)
-    _folder = os.path.join(OUTPUT_FOLDER,"imagen")
+    _folder = os.path.join(TEMP_FOLDER)
+    _remote_folder = os.path.join(REMOTE_FOLDER,"imagen")
     _datetime = D.date_file_name()
     
-    if not os.path.isdir(_folder):
+    # if not os.path.isdir(_folder):
-        raise FileExistsError(f"FileExistsError: {_folder}")
+    #     raise FileExistsError(f"FileExistsError: {_folder}")
     
     images = model.generate_images(
         prompt=prompt,
@@ -42,12 +43,13 @@ def imagen_generate_image(prompt,download_count=1):
     
     if len(images.images) <=1 :
         images[0].save(location=os.path.join(_folder,f"imagen_{_file_name}_{_datetime}.png"), include_generation_parameters=False)
+        sftp_client.remote_copy_data(os.path.join(_folder,f"imagen_{_file_name}_{_datetime}.png"), os.path.join(_remote_folder,f"imagen_{_file_name}_{_datetime}.png"))
         return 1
     
     else:
         for i in range(len(images.images)):
             images[i].save(location=os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), include_generation_parameters=False)
-        
+            sftp_client.remote_copy_data(os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), os.path.join(_remote_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"))
         return len(images.images)
     
 def imagen_generate_image_data(prompt,download_count=1):
@@ -76,7 +78,7 @@ def imagen_generate_image_path(image_prompt):
     QUERY = "query"
     create_time = D.date_file_name()
     
-    folder_name = os.path.join(OUTPUT_FOLDER,f"{MODEL}_{QUERY}_{image_prompt}_{create_time}")
+    folder_name = os.path.join(TEMP_FOLDER,f"{MODEL}_{QUERY}_{image_prompt}_{create_time}")
     if not os.path.exists(folder_name):
         os.makedirs(folder_name)
     

main_rest/app/routes/services.py

@@ -9,7 +9,7 @@
 @brief: services api
 """
 
-import requests, json, traceback
+import requests, json, traceback, os
 from fastapi import APIRouter, Depends, Body
 from starlette.requests import Request
 from typing import Annotated, List
@@ -24,6 +24,8 @@ from custom_apps.bingart.bingart import BingArtGenerator
 from custom_apps.imagen.custom_imagen import imagen_generate_image, imagen_generate_image_path
 from main_rest.app.utils.parsing_utils import download_range
 from custom_apps.utils import cookie_manager
+from utils.custom_sftp import sftp_client
+from const import REMOTE_FOLDER, TEMP_FOLDER
 
 router = APIRouter(prefix="/services")
 
@@ -176,6 +178,17 @@ async def vactor_image(request: Request, request_body_info: M.VactorImageSearchR
         if json.loads(vactor_response.text)["error"] != None:
             raise Exception(f"vactor error: {json.loads(vactor_response.text)['error']}")
         
+        # remote
+        _directory_path, _file = os.path.split(img_path)
+        _base_bame = os.path.basename(_directory_path)
+        
+        # remote 폴더 생성
+        sftp_client.remote_mkdir(os.path.join(REMOTE_FOLDER, _base_bame))
+        
+        # remote 폴더에 이미지 저장
+        for i in os.listdir(_directory_path):
+            sftp_client.remote_copy_data(local_path=os.path.join(_directory_path, i), remote_path=os.path.join(REMOTE_FOLDER, _base_bame, i))        
+        
         return response.set_message()
     
     except Exception as e:

requirements_main.txt

@@ -31,7 +31,5 @@ bingart==1.1.0
 #DALL-E 3
 # openai
 
-# faiss
+# SFTP
-# scikit-learn
+paramiko
-# flax
-# tensorflow

utils/custom_sftp.py

@@ -0,0 +1,29 @@
+import paramiko
+
+class CustomSFTPClient():
+    def __init__(self):
+        host = "192.168.200.230"
+        port = 22
+        id = "fermat"
+        pw = "1234"
+        
+        self.ssh_client = paramiko.SSHClient()
+        self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        self.ssh_client.connect(hostname=host, port=port, username=id, password=pw)
+        
+        self.sftp_client = self.ssh_client.open_sftp()
+
+    def remote_mkdir(self, remote_path):
+        self.sftp_client.mkdir(remote_path)
+    
+    def remote_copy_data(self,local_path, remote_path):
+        self.sftp_client.put(local_path, remote_path)
+    
+    def client_close(self):
+        self.sftp_client.close()
+        self.ssh_client.close()
+        
+sftp_client = None
+
+if sftp_client is None:
+    sftp_client = CustomSFTPClient()