Spaces:
Running
on
T4
Running
on
T4
File size: 737 Bytes
7b127f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import torch
import numpy as np
import torchvision.transforms as T
from pathlib import Path
from PIL import Image
from pathlib import Path
from app.DataProcessor.DataProcessor import DataProcessor
class ImageProcessor(DataProcessor):
def _get_img_tensor(self, image_file: Path) -> torch.Tensor:
"""
Return a (3, 224, 224) shape tensor
"""
transform = T.Compose([
T.ToPILImage(),
T.Resize((224, 224)),
T.ToTensor(),
T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])
img = np.array(Image.open(Path(image_file)).convert("RGB"))
img = transform(img).to(self._device)
return img |