Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,579 Bytes
3427608 b177a48 3427608 02f8ed6 3427608 02f8ed6 3427608 b177a48 02f8ed6 6dc341c b177a48 3427608 02f8ed6 3427608 243f5a6 02f8ed6 243f5a6 02f8ed6 b39062c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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
ssh_client = None
sftp_client = None
def create_ssh_client(server, port, user, password):
global ssh_client, sftp_client
ssh_client = paramiko.SSHClient()
ssh_client.load_system_host_keys()
ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh_client.connect(server, port, user, password)
sftp_client = ssh_client.open_sftp()
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 get_random_mscoco_prompt():
global sftp_client
num = random.randint(0, 2999)
file = "{}.txt".format(num)
remote_file_path = os.path.join(SSH_MSCOCO, file)
with sftp_client.file(remote_file_path, 'r') as f:
content = f.read().decode('utf-8')
print(f"Content of {file}:")
print("\n")
return content
def create_remote_directory(remote_directory):
global ssh_client
stdin, stdout, stderr = ssh_client.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.")
return f'{SSH_LOG}/{remote_directory}'
def upload_ssh_all(states, output_dir, data, data_path):
global sftp_client
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 sftp_client 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_client.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_client.putfo(json_byte_stream, data_path)
print(f"Successfully uploaded JSON data to {data_path}")
create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD) |