update upload

serve/upload.py

@@ -80,15 +80,16 @@ def create_remote_directory(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]}")
+    with sftp_client as sftp:
+        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]}")
     
 
     
@@ -99,24 +100,25 @@ def upload_ssh_all(states, output_dir, data, data_path):
     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)
-
+    output_file = os.path.join(output_dir, f"{i}.jpg")
+    output_file_list.append(output_file)
+    image_list.append(states[i].output)
+
+
+    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]}")
     with sftp_client as sftp:
-        # 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)