Spaces:
Running
on
Zero
Running
on
Zero
| import paramiko | |
| import numpy as np | |
| import io, os | |
| from PIL import Image | |
| import requests | |
| import json | |
| import random | |
| from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_LOG, SSH_MSCOCO | |
| def create_ssh_client(server, port, user, password): | |
| ssh = paramiko.SSHClient() | |
| ssh.load_system_host_keys() | |
| ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) | |
| ssh.connect(server, port, user, password) | |
| return ssh | |
| def get_image_from_url(image_url): | |
| response = requests.get(image_url) | |
| response.raise_for_status() # success | |
| return Image.open(io.BytesIO(response.content)) | |
| # def upload_image_via_sftp(ssh, image, remote_image_path): | |
| # if isinstance(image, str): | |
| # print("get url image") | |
| # image = get_image_from_url(image) | |
| # with ssh.open_sftp() as sftp: | |
| # with io.BytesIO() as image_byte_stream: | |
| # image.save(image_byte_stream, format='JPEG') | |
| # image_byte_stream.seek(0) | |
| # sftp.putfo(image_byte_stream, remote_image_path) | |
| # print(f"Successfully uploaded image to {remote_image_path}") | |
| # def upload_json_via_sftp(ssh, data, remote_json_path): | |
| # json_data = json.dumps(data, indent=4) | |
| # with ssh.open_sftp() as sftp: | |
| # with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream: | |
| # sftp.putfo(json_byte_stream, remote_json_path) | |
| # print(f"Successfully uploaded JSON data to {remote_json_path}") | |
| def upload_image(states, output_dir): | |
| pass | |
| # ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| # for i in range(len(states)): | |
| # output_file = os.path.join(output_dir, f"{i}.jpg") | |
| # upload_image_via_sftp(ssh, states[i].output, output_file) | |
| def upload_informance(data, data_path): | |
| pass | |
| # ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| # upload_json_via_sftp(ssh, data, data_path) | |
| def get_random_mscoco_prompt(): | |
| ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| sftp = ssh.open_sftp() | |
| files = sftp.listdir(SSH_MSCOCO) | |
| txt_files = [file for file in files if file.endswith('.txt')] | |
| selected_files = random.sample(txt_files, 1) # get one prompt | |
| for file in selected_files: | |
| remote_file_path = os.path.join(SSH_MSCOCO, file) | |
| with sftp.file(remote_file_path, 'r') as f: | |
| content = f.read().decode('utf-8') | |
| print(f"Content of {file}:") | |
| print("\n") | |
| sftp.close() | |
| ssh.close() | |
| return content | |
| def create_remote_directory(remote_directory): | |
| ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| stdin, stdout, stderr = ssh.exec_command(f'mkdir -p {SSH_LOG}/{remote_directory}') | |
| error = stderr.read().decode('utf-8') | |
| if error: | |
| print(f"Error: {error}") | |
| else: | |
| print(f"Directory {remote_directory} created successfully.") | |
| # ssh.close() | |
| return f'{SSH_LOG}/{remote_directory}' | |
| def upload_ssh_all(states, output_dir, data, data_path): | |
| ssh = create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) | |
| output_file_list = [] | |
| image_list = [] | |
| for i in range(len(states)): | |
| output_file = os.path.join(output_dir, f"{i}.jpg") | |
| output_file_list.append(output_file) | |
| image_list.append(states[i].output) | |
| with ssh.open_sftp() as sftp: | |
| for i in range(len(output_file_list)): | |
| if isinstance(image_list[i], str): | |
| print("get url image") | |
| image_list[i] = get_image_from_url(image_list[i]) | |
| with io.BytesIO() as image_byte_stream: | |
| image_list[i].save(image_byte_stream, format='JPEG') | |
| image_byte_stream.seek(0) | |
| sftp.putfo(image_byte_stream, output_file_list[i]) | |
| print(f"Successfully uploaded image to {output_file_list[i]}") | |
| json_data = json.dumps(data, indent=4) | |
| with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream: | |
| sftp.putfo(json_byte_stream, data_path) | |
| print(f"Successfully uploaded JSON data to {data_path}") |