File size: 579 Bytes
a68b822 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import torch
from torchvision import transforms
from PIL import Image
import io
# Same transformations used during training
def get_transform():
return transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def process_image(image_bytes):
"""Process an image from bytes to tensor"""
image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
transform = get_transform()
return transform(image).unsqueeze(0) # Add batch dimension
|