File size: 1,945 Bytes
7ebc30a |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import dataloader as dl
import torch
import argparse
import transformers
import PIL.Image as Image
from typing import Union, List
from transformers.image_processing_utils import BaseImageProcessor
from transformers.utils import PushToHubMixin
class CommForImageProcessor(BaseImageProcessor, PushToHubMixin):
"""
Image processor for Community Forensics VIT model. Processes PIL images and returns PyTorch tensors.
"""
image_processor_type = "commfor_image_processor"
model_input_names = ["pixel_values"]
def __init__(self, size=384, **kwargs):
super().__init__(**kwargs)
self.size = size
assert self.size in [224, 384], f"Unsupported size: {self.size}. Supported sizes are 224 and 384."
def preprocess(
self,
images: Union[Image.Image, List[Image.Image]],
mode: str = "test",
**kwargs
):
"""
Preprocess the input images to PyTorch tensors.
"""
assert mode in ["test", "train"], f"Unsupported mode: {mode}. Supported modes are 'test' and 'train'."
assert isinstance(images, (Image.Image, list)), "Input must be a PIL Image or a list of PIL Images."
if isinstance(images, Image.Image):
images = [images]
args = argparse.Namespace()
args.input_size = self.size
args.rsa_ops="JPEGinMemory,RandomResizeWithRandomIntpl,RandomCrop,RandomHorizontalFlip,RandomVerticalFlip,RRCWithRandomIntpl,RandomRotation,RandomTranslate,RandomShear,RandomPadding,RandomCutout"
args.rsa_min_num_ops='0'
args.rsa_max_num_ops='2'
transform = dl.get_transform(args, mode=mode)
processed_images = [transform(image) for image in images] # the output would be tensors
if len(processed_images) == 1:
return {"pixel_values": processed_images[0]}
else:
return {"pixel_values": torch.stack(processed_images)}
|