pad image before resize

app.py

@@ -66,6 +66,19 @@ def count_files(*inputs):
             Training_Steps = file_counter*200
     return(gr.update(visible=True, value=f"You are going to train {concept_counter} {type_of_thing}(s), with {file_counter} images for {Training_Steps} steps. This should take around {round(Training_Steps/1.5, 2)} seconds, or {round((Training_Steps/1.5)/3600, 2)} hours. As a reminder, the T4 GPU costs US$0.60 for 1h. Once training is over, don't forget to swap the hardware back to CPU."))
 
+def pad_image(image):
+    w, h = image.size
+    if w == h:
+        return image
+    elif w > h:
+        new_image = Image.new(image.mode, (w, w), (0, 0, 0))
+        new_image.paste(image, (0, (w - h) // 2))
+        return new_image
+    else:
+        new_image = Image.new(image.mode, (h, h), (0, 0, 0))
+        new_image.paste(image, ((h - w) // 2, 0))
+        return new_image
+
 def train(*inputs):
     if "IS_SHARED_UI" in os.environ:
         raise gr.Error("This Space only works in duplicated instances")
@@ -84,13 +97,7 @@ def train(*inputs):
                     raise gr.Error("You forgot to define your concept prompt")
                 for j, file_temp in enumerate(files):
                     file = Image.open(file_temp.name)
-                    width, height = file.size
-                    side_length = min(width, height)
-                    left = (width - side_length)/2
-                    top = (height - side_length)/2
-                    right = (width + side_length)/2
-                    bottom = (height + side_length)/2
-                    image = file.crop((left, top, right, bottom))
+                    image = pad_image(file)
                     image = image.resize((512, 512))
                     extension = file_temp.name.split(".")[1]
                     image = image.convert('RGB')