Spaces:
Running on Zero
Running on Zero
File size: 1,188 Bytes
f7b58ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | from PIL import Image
import numpy as np
class ImageProcess:
def __init__(self, size=256):
self.size = size
def process(self, image_path):
with Image.open(image_path) as img:
image = img.convert('RGB')
w, h = image.size
# 1. Match PyTorch's resizing dimensions (cast to int, don't use round)
if w < h:
new_w = self.size
new_h = int(h * self.size / w)
else:
new_h = self.size
new_w = int(w * self.size / h)
image = image.resize((new_w, new_h), Image.Resampling.BILINEAR)
# 2. Center crop math
left = (new_w - self.size) // 2
top = (new_h - self.size) // 2
image = image.crop((left, top, left + self.size, top + self.size))
# 3. Convert to float array AND scale to [0.0, 1.0]
# CRITICAL: This is what transforms.ToTensor() does under the hood
image = np.asarray(image, dtype=np.float32) / 255.0
# 4. Reorder dimensions to CHW and expand to BCHW
image = np.transpose(image, (2, 0, 1))
image = np.expand_dims(image, 0)
return image
|