pad non square image before resize

#1
by radames - opened
Files changed (1) hide show
  1. app.py +14 -0
app.py CHANGED
@@ -28,8 +28,22 @@ pipe = DiffusionPipeline.from_pretrained(
28
  device = "cuda" if torch.cuda.is_available() else "cpu"
29
  pipe = pipe.to(device)
30
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
  def process_image(image, text, prompt):
 
33
  image = image.resize((512, 512))
34
  with autocast("cuda"):
35
  inpainted_image = pipe(image=image, text=text, prompt=prompt).images[0]
 
28
  device = "cuda" if torch.cuda.is_available() else "cpu"
29
  pipe = pipe.to(device)
30
 
31
+ def pad_image(image):
32
+ w, h = image.size
33
+ if w == h:
34
+ return image
35
+ elif w > h:
36
+ new_image = Image.new(image.mode, (w, w), (0, 0, 0))
37
+ new_image.paste(image, (0, (w - h) // 2))
38
+ return new_image
39
+ else:
40
+ new_image = Image.new(image.mode, (h, h), (0, 0, 0))
41
+ new_image.paste(image, ((h - w) // 2, 0))
42
+ return new_image
43
+
44
 
45
  def process_image(image, text, prompt):
46
+ image = pad_image(image)
47
  image = image.resize((512, 512))
48
  with autocast("cuda"):
49
  inpainted_image = pipe(image=image, text=text, prompt=prompt).images[0]