Update README.md

README.md

@@ -56,27 +56,33 @@ def preprocess(image_path, ref_size=512):
     image = Image.open(image_path).convert("RGB")
     w, h = image.size
 
+    # Resize while maintaining aspect ratio
     scale = ref_size / max(h, w)
     new_w, new_h = int(w * scale), int(h * scale)
     image_resized = image.resize((new_w, new_h), Image.BILINEAR)
 
+    # Create a new blank image of ref_size x ref_size and paste the resized image
+    # This will pad the image to ref_size x ref_size
+    padded_image = Image.new("RGB", (ref_size, ref_size), (0, 0, 0)) # Black padding
+    padded_image.paste(image_resized, ((ref_size - new_w) // 2, (ref_size - new_h) // 2))
+
     transform = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5, 0.5, 0.5],
                              std=[0.5, 0.5, 0.5])
     ])
 
-    tensor = transform(image_resized).unsqueeze(0)
+    tensor = transform(padded_image).unsqueeze(0)
     return tensor, (w, h)
 
 # Run inference
-inp, original_size = preprocess("input.jpg")
+inp, original_size = preprocess("/content/portrait.jpg")
 
 import torch
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model.to(device)
-inp.to(device)
+inp = inp.to(device) # Assign the tensor to the device
 
 with torch.no_grad():
     semantic, detail, matte = model(inp, True)