update upload

serve/upload.py

@@ -5,6 +5,7 @@ from PIL import Image
 import requests
 import json
 import random
+import concurrent.futures
 from .constants import SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD, SSH_LOG, SSH_MSCOCO
 
 ssh_client = None
@@ -76,6 +77,20 @@ def create_remote_directory(remote_directory):
     else:
         print(f"Directory {remote_directory} created successfully.")
     return f'{SSH_LOG}/{remote_directory}'
+
+
+def upload_images(i, image_list, output_file_list, sftp_client):
+    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] = image_list[i].resize((512, 512), Image.ANTIALIAS)
+        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]}")
+    
+
     
 def upload_ssh_all(states, output_dir, data, data_path):
     global sftp_client
@@ -89,16 +104,19 @@ def upload_ssh_all(states, output_dir, data, data_path):
         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] = image_list[i].resize((512, 512), Image.ANTIALIAS)
-                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]}")
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = [executor.submit(upload_images, i, image_list, output_file_list, sftp_client) for i in range(len(output_file_list))]
+
+        # 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] = image_list[i].resize((512, 512), Image.ANTIALIAS)
+        #         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)