LTF-v3.3 / utils /tile_dataset_srm.py

Use non-overlapping full blocks with boundary-only completion

a2a2e29 verified about 14 hours ago

14.6 kB

from PIL import Image, ImageOps, ImageFilter
import random
import io, os
import math
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.transforms import InterpolationMode
BICUBIC = InterpolationMode.BICUBIC

def resize_with_random_scale(image: Image.Image, p: float = 0.5, scale_range=(0.25, 2.0)) -> Image.Image:
    if random.random() > p:
        return image 

    w, h = image.size
    scale = random.uniform(*scale_range)
    new_w = max(1, int(w * scale))
    new_h = max(1, int(h * scale))
    return image.resize((new_w, new_h), resample=Image.BICUBIC)

def pad_crop_image(image, target_width, target_height):
    src_width, src_height = image.size
    # 0
    pad_left = pad_right = pad_top = pad_bottom = 0
    crop_left = crop_top = 0

    if src_width < target_width:
        pad_left = random.randint(0, target_width - src_width)
        pad_right = target_width - src_width - pad_left
    else:
        crop_left = random.randint(0, src_width - target_width)

    if src_height < target_height:
        pad_top = random.randint(0, target_height - src_height)
        pad_bottom = target_height - src_height - pad_top
    else:
        crop_top = random.randint(0, src_height - target_height)

    if src_width > target_width or src_height > target_height:
        left = crop_left if src_width > target_width else 0
        top = crop_top if src_height > target_height else 0
        right = left + target_width
        bottom = top + target_height
        image = image.crop((left, top, right, bottom))
    # Padding
    if src_width < target_width or src_height < target_height:
        image = ImageOps.expand(
            image,
            border=(pad_left, pad_top, pad_right, pad_bottom),
            fill=(0, 0, 0)
        )
    return image

def get_start_positions(image_length, block_size, bias=0):
    """Non-overlapping full blocks, plus one boundary block when needed."""
    if image_length <= block_size:
        return [bias]
    positions = list(range(0, image_length - block_size + 1, block_size))
    if positions[-1] != image_length - block_size:
        positions.append(image_length - block_size)
    return [position + bias for position in positions]

def crop_image_to_tiles(image, block_size):
    """
    crop to block_size x block_size
    """
    width, height = image.size
    x_positions = get_start_positions(width, block_size)
    y_positions = get_start_positions(height, block_size)
    tiles = []
    for y in y_positions:
        for x in x_positions:
            tiles.append(image.crop((x, y, x + block_size, y + block_size)))
    return tiles

def crop_image_to_tiles_new(image, block_size):
    width, height = image.size
    x_positions = get_start_positions(width, block_size)
    y_positions = get_start_positions(height, block_size)
    tiles = []
    for y in y_positions:
        for x in x_positions:
            tiles.append(image.crop((x, y, x + block_size, y + block_size)))
    if width > block_size and height > block_size:
        leave_space = min(block_size, width-block_size)
        x_positions = get_start_positions(width-leave_space, block_size, leave_space//2)
        leave_space = min(block_size, height-block_size)
        y_positions = get_start_positions(height-leave_space, block_size, leave_space//2)
        for y in y_positions:
            for x in x_positions:
                tiles.append(image.crop((x, y, x + block_size, y + block_size)))
    return tiles

def random_crop(image: Image.Image, size: int = 224, is_pad_black: bool = False) -> Image.Image:
    w, h = image.size
    crop_size = min(w, h, size)
    if w == crop_size:
        left = 0
    else:
        left = random.randint(0, w - crop_size)
    if h == crop_size:
        top = 0
    else:
        top = random.randint(0, h - crop_size)

    cropped = image.crop((left, top, left + crop_size, top + crop_size))

    if is_pad_black and crop_size < size:
        padded_img = Image.new("RGB", (size, size), (0, 0, 0))
        offset = ((size - crop_size) // 2, (size - crop_size) // 2)
        padded_img.paste(cropped, offset)
        return padded_img
    else:
        return cropped

def _convert_image_to_rgb(image):
    return image.convert("RGB")


def clip_transform_allresize(n_px, blur_radius=None, blur_prob=0.5):
    if blur_radius is None:
        return transforms.Compose([
            transforms.Resize((n_px, n_px), interpolation=BICUBIC),
            _convert_image_to_rgb,
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        ])
    else:
        return transforms.Compose([
            transforms.Resize((n_px, n_px), interpolation=BICUBIC),
            transforms.RandomApply(
                [transforms.GaussianBlur(kernel_size=5, sigma=blur_radius)], p=blur_prob  # (0.1, 3)
            ),
            _convert_image_to_rgb,
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        ])

def clip_transform(n_px, blur_radius=None, blur_prob=0.5):
    if blur_radius is None:
        return transforms.Compose([
            transforms.Resize(n_px, interpolation=BICUBIC),
            transforms.CenterCrop(n_px),
            _convert_image_to_rgb,
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        ])
    else:
        return transforms.Compose([
            transforms.Resize(n_px, interpolation=BICUBIC),
            transforms.CenterCrop(n_px),
            transforms.RandomApply(
                [transforms.GaussianBlur(kernel_size=5, sigma=blur_radius)], p=blur_prob  # (0.1, 3)
            ),
            _convert_image_to_rgb,
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        ])

class JpegAugment():
    def __init__(self, quality=None, p=1, seed=42):
        self.local_random = random.Random(seed)
        self.quality = quality
        self.p = p
    
    def __call__(self, img):
        if self.quality is None:
            return img
        if self.local_random.random() < self.p:
            buffer = io.BytesIO()
            if self.quality == -1:
                img.save(buffer, format='JPEG', quality=self.local_random.randint(30, 100))
            else:
                if isinstance(self.quality, tuple):
                    img.save(buffer, format='JPEG', quality=self.local_random.randint(*self.quality))
                else:
                    img.save(buffer, format='JPEG', quality=self.quality)
            buffer.seek(0)
            return Image.open(buffer).convert('RGB')
        return img

def webp_compress(img, quality=80, method=4):
    buffer = io.BytesIO() 
    img.save(buffer, format='WEBP', quality=quality, method=method)
    buffer.seek(0)
    return Image.open(buffer).convert('RGB')

# Image extensions supported when scanning real/fake dirs (case-insensitive).
# Missing .JPG / .webp etc. can cause real folder to show 0 images if only those exist.
IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tif', '.tiff', '.ppm')


def find_files_with_extensions(directory, extensions=None, max_num=None):
    """Find image paths under directory. If extensions is None, use IMG_EXTENSIONS (case-insensitive)."""
    if extensions is None:
        extensions = IMG_EXTENSIONS
    ext_lower = tuple(e.lower() for e in extensions)
    matching_files = []
    check_num = 0
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(ext_lower):
                matching_files.append(os.path.join(root, file))
                check_num += 1
                if max_num is not None and check_num >= max_num:
                    return matching_files
    return matching_files

class TiledRFDataset(Dataset):
    """ Dataset -> (img, label)"""
    def __init__(self, real_dirs:list, fake_dirs:list, tile_size=224, set_size=None, shuffle_seed=42, balance=True, 
                 jpeg_augment=None, jpeg_p=0.1,
                 webp_augment=None, webp_p=0.1,
                 resize_augment=None, resize_p=0.5,
                 blur_augment=None, blur_p=0.5,
                 default_train_tiles=4, is_train=False):
        self.set_size = set_size
        self.shuffle_seed = shuffle_seed
        self.samples = []
        self.real_samples = []
        self.fake_samples = []
        self.tile_size=tile_size
        self.transform = clip_transform(tile_size, blur_augment, blur_p)
        self.default_train_tiles = default_train_tiles
        self.is_train = is_train
        
        self.jpeg_augment=jpeg_augment
        self.jpeg_p = jpeg_p
        self.JpegEnhance = JpegAugment(quality=jpeg_augment, p=jpeg_p, seed=shuffle_seed)
        self.webp_augment = webp_augment
        self.webp_p = webp_p
        self.resize_augment = resize_augment
        self.resize_p = resize_p

        for p in real_dirs:
            find_paths_real = find_files_with_extensions(p)
            self.real_samples += [(f, 0) for f in find_paths_real]
        for p in fake_dirs:
            find_paths_fake = find_files_with_extensions(p)
            self.fake_samples += [(f, 1) for f in find_paths_fake]
        if balance:
            target_len = min(len(self.real_samples), len(self.fake_samples))
            self.real_samples = self.real_samples[:target_len]
            self.fake_samples = self.fake_samples[:target_len]
        self.samples = self.real_samples + self.fake_samples
        # if self.set_size is not None and len(self.samples) < self.set_size:
        #     n_repeats = math.ceil(self.set_size / len(self.samples))
        #     self.samples = (self.samples * n_repeats)[:self.set_size]
        if self.set_size is not None and len(self.samples) > self.set_size:
            local_random = random.Random(self.shuffle_seed)
            self.samples = local_random.sample(self.samples, self.set_size)
    def __len__(self):
        return len(self.samples)
    
    def _get_train_item(self, idx):
        path, label = self.samples[idx]
        try:
            img = Image.open(path).convert('RGB')
        except Exception as e:
            print(f"Skipping corrupted image: {path}")
            return self.__getitem__((idx + 1) % len(self.samples))
        if self.resize_augment is not None:
            if self.resize_augment == -1:
                img = resize_with_random_scale(img, p=self.resize_p)
            else:
                if isinstance(self.resize_augment, tuple):
                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=self.resize_augment)
                else:
                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=(self.resize_augment, self.resize_augment))
        if self.jpeg_augment is not None:
            img = self.JpegEnhance(img)
        
        # crops_one_side = 2
        # img = pad_crop_image(img, crops_one_side*self.tile_size, crops_one_side*self.tile_size)
        # tiles = crop_image_to_tiles(img, self.tile_size)[:crops_one_side**2]
        tiles = []
        for i in range(self.default_train_tiles):  # 10
            tiles.append(random_crop(img, self.tile_size, is_pad_black=False))
        img_resize = self.transform(img)
        img_list = [img_resize]
        for tile in tiles:
            img_list.append(self.transform(tile))
        return img_list, img, label

    def _get_test_item(self, idx):
        path, label = self.samples[idx]
        try:
            img = Image.open(path).convert('RGB')
        except Exception as e:
            print(f"Skipping corrupted image: {path}")
            return self.__getitem__((idx + 1) % len(self.samples))
        if self.resize_augment is not None:
            if self.resize_augment == -1:
                img = resize_with_random_scale(img, p=self.resize_p)
            else:
                if isinstance(self.resize_augment, tuple):
                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=self.resize_augment)
                else:
                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=(self.resize_augment, self.resize_augment))
        if self.jpeg_augment is not None:
            img = self.JpegEnhance(img)
        if self.webp_augment is not None:
            if random.random() < self.webp_p:
                if isinstance(self.webp_augment, tuple):
                    img = webp_compress(img, quality=random.randint(*self.webp_augment))
                else:
                    img = webp_compress(img, quality=self.webp_augment)
        # testtime batchsize=1
        tiles = crop_image_to_tiles(img, self.tile_size)
        img_resize = self.transform(img)
        img_list = [img_resize]
        for tile in tiles:
            img_list.append(self.transform(tile))
        return img_list, img, label, path

    def __getitem__(self, idx):
        if self.is_train:
            return self._get_train_item(idx)
        else:
            return self._get_test_item(idx)
        

def collate_img_list_to_tensor_train(batch):
    """
    batch = list of (img_list, label)
    - img_list: List of N Tensor[B, 3, H, W]
    return：
    - images: Tensor[B, N, 3, H, W]
    - labels: Tensor[B]
    """
    img_lists, raw_imgs, labels = zip(*batch)  # len = B
    # img_lists: list of length B, each is List[N x (3,H,W)]
    stacked = torch.stack([torch.stack(imgs) for imgs in img_lists])  # (B, N, 3, H, W)
    raw_imgs = list(raw_imgs)
    labels = torch.tensor(labels)
    return stacked, raw_imgs, labels

def collate_img_list_to_tensor_test(batch):
    """
    batch = list of (img_list, label)
    - img_list: List of N Tensor[B, 3, H, W]
    return：
    - images: Tensor[B, N, 3, H, W]
    - labels: Tensor[B]
    """
    img_lists, raw_imgs, labels, paths = zip(*batch)  # len = B
    # img_lists: list of length B, each is List[N x (3,H,W)]
    stacked = torch.stack([torch.stack(imgs) for imgs in img_lists])  # (B, N, 3, H, W)
    raw_imgs = list(raw_imgs)
    labels = torch.tensor(labels)
    return stacked, raw_imgs, labels, paths


if __name__ == '__main__':
    block_size = 224
    width = 512
    leave_space = min(block_size, width-block_size)
    x_positions = get_start_positions(width-leave_space, block_size, leave_space//2)
    print(leave_space, x_positions)

	from PIL import Image, ImageOps, ImageFilter
	import random
	import io, os
	import math
	import torch
	from torch.utils.data import Dataset, DataLoader
	from torchvision import transforms
	from torchvision.transforms import InterpolationMode
	BICUBIC = InterpolationMode.BICUBIC

	def resize_with_random_scale(image: Image.Image, p: float = 0.5, scale_range=(0.25, 2.0)) -> Image.Image:
	if random.random() > p:
	return image

	w, h = image.size
	scale = random.uniform(*scale_range)
	new_w = max(1, int(w * scale))
	new_h = max(1, int(h * scale))
	return image.resize((new_w, new_h), resample=Image.BICUBIC)

	def pad_crop_image(image, target_width, target_height):
	src_width, src_height = image.size
	# 0
	pad_left = pad_right = pad_top = pad_bottom = 0
	crop_left = crop_top = 0

	if src_width < target_width:
	pad_left = random.randint(0, target_width - src_width)
	pad_right = target_width - src_width - pad_left
	else:
	crop_left = random.randint(0, src_width - target_width)

	if src_height < target_height:
	pad_top = random.randint(0, target_height - src_height)
	pad_bottom = target_height - src_height - pad_top
	else:
	crop_top = random.randint(0, src_height - target_height)

	if src_width > target_width or src_height > target_height:
	left = crop_left if src_width > target_width else 0
	top = crop_top if src_height > target_height else 0
	right = left + target_width
	bottom = top + target_height
	image = image.crop((left, top, right, bottom))
	# Padding
	if src_width < target_width or src_height < target_height:
	image = ImageOps.expand(
	image,
	border=(pad_left, pad_top, pad_right, pad_bottom),
	fill=(0, 0, 0)
	)
	return image

	def get_start_positions(image_length, block_size, bias=0):
	"""Non-overlapping full blocks, plus one boundary block when needed."""
	if image_length <= block_size:
	return [bias]
	positions = list(range(0, image_length - block_size + 1, block_size))
	if positions[-1] != image_length - block_size:
	positions.append(image_length - block_size)
	return [position + bias for position in positions]

	def crop_image_to_tiles(image, block_size):
	"""
	crop to block_size x block_size
	"""
	width, height = image.size
	x_positions = get_start_positions(width, block_size)
	y_positions = get_start_positions(height, block_size)
	tiles = []
	for y in y_positions:
	for x in x_positions:
	tiles.append(image.crop((x, y, x + block_size, y + block_size)))
	return tiles

	def crop_image_to_tiles_new(image, block_size):
	width, height = image.size
	x_positions = get_start_positions(width, block_size)
	y_positions = get_start_positions(height, block_size)
	tiles = []
	for y in y_positions:
	for x in x_positions:
	tiles.append(image.crop((x, y, x + block_size, y + block_size)))
	if width > block_size and height > block_size:
	leave_space = min(block_size, width-block_size)
	x_positions = get_start_positions(width-leave_space, block_size, leave_space//2)
	leave_space = min(block_size, height-block_size)
	y_positions = get_start_positions(height-leave_space, block_size, leave_space//2)
	for y in y_positions:
	for x in x_positions:
	tiles.append(image.crop((x, y, x + block_size, y + block_size)))
	return tiles

	def random_crop(image: Image.Image, size: int = 224, is_pad_black: bool = False) -> Image.Image:
	w, h = image.size
	crop_size = min(w, h, size)
	if w == crop_size:
	left = 0
	else:
	left = random.randint(0, w - crop_size)
	if h == crop_size:
	top = 0
	else:
	top = random.randint(0, h - crop_size)

	cropped = image.crop((left, top, left + crop_size, top + crop_size))

	if is_pad_black and crop_size < size:
	padded_img = Image.new("RGB", (size, size), (0, 0, 0))
	offset = ((size - crop_size) // 2, (size - crop_size) // 2)
	padded_img.paste(cropped, offset)
	return padded_img
	else:
	return cropped

	def _convert_image_to_rgb(image):
	return image.convert("RGB")


	def clip_transform_allresize(n_px, blur_radius=None, blur_prob=0.5):
	if blur_radius is None:
	return transforms.Compose([
	transforms.Resize((n_px, n_px), interpolation=BICUBIC),
	_convert_image_to_rgb,
	transforms.ToTensor(),
	transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
	])
	else:
	return transforms.Compose([
	transforms.Resize((n_px, n_px), interpolation=BICUBIC),
	transforms.RandomApply(
	[transforms.GaussianBlur(kernel_size=5, sigma=blur_radius)], p=blur_prob # (0.1, 3)
	),
	_convert_image_to_rgb,
	transforms.ToTensor(),
	transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
	])

	def clip_transform(n_px, blur_radius=None, blur_prob=0.5):
	if blur_radius is None:
	return transforms.Compose([
	transforms.Resize(n_px, interpolation=BICUBIC),
	transforms.CenterCrop(n_px),
	_convert_image_to_rgb,
	transforms.ToTensor(),
	transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
	])
	else:
	return transforms.Compose([
	transforms.Resize(n_px, interpolation=BICUBIC),
	transforms.CenterCrop(n_px),
	transforms.RandomApply(
	[transforms.GaussianBlur(kernel_size=5, sigma=blur_radius)], p=blur_prob # (0.1, 3)
	),
	_convert_image_to_rgb,
	transforms.ToTensor(),
	transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
	])

	class JpegAugment():
	def __init__(self, quality=None, p=1, seed=42):
	self.local_random = random.Random(seed)
	self.quality = quality
	self.p = p

	def __call__(self, img):
	if self.quality is None:
	return img
	if self.local_random.random() < self.p:
	buffer = io.BytesIO()
	if self.quality == -1:
	img.save(buffer, format='JPEG', quality=self.local_random.randint(30, 100))
	else:
	if isinstance(self.quality, tuple):
	img.save(buffer, format='JPEG', quality=self.local_random.randint(*self.quality))
	else:
	img.save(buffer, format='JPEG', quality=self.quality)
	buffer.seek(0)
	return Image.open(buffer).convert('RGB')
	return img

	def webp_compress(img, quality=80, method=4):
	buffer = io.BytesIO()
	img.save(buffer, format='WEBP', quality=quality, method=method)
	buffer.seek(0)
	return Image.open(buffer).convert('RGB')

	# Image extensions supported when scanning real/fake dirs (case-insensitive).
	# Missing .JPG / .webp etc. can cause real folder to show 0 images if only those exist.
	IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tif', '.tiff', '.ppm')


	def find_files_with_extensions(directory, extensions=None, max_num=None):
	"""Find image paths under directory. If extensions is None, use IMG_EXTENSIONS (case-insensitive)."""
	if extensions is None:
	extensions = IMG_EXTENSIONS
	ext_lower = tuple(e.lower() for e in extensions)
	matching_files = []
	check_num = 0
	for root, dirs, files in os.walk(directory):
	for file in files:
	if file.lower().endswith(ext_lower):
	matching_files.append(os.path.join(root, file))
	check_num += 1
	if max_num is not None and check_num >= max_num:
	return matching_files
	return matching_files

	class TiledRFDataset(Dataset):
	""" Dataset -> (img, label)"""
	def __init__(self, real_dirs:list, fake_dirs:list, tile_size=224, set_size=None, shuffle_seed=42, balance=True,
	jpeg_augment=None, jpeg_p=0.1,
	webp_augment=None, webp_p=0.1,
	resize_augment=None, resize_p=0.5,
	blur_augment=None, blur_p=0.5,
	default_train_tiles=4, is_train=False):
	self.set_size = set_size
	self.shuffle_seed = shuffle_seed
	self.samples = []
	self.real_samples = []
	self.fake_samples = []
	self.tile_size=tile_size
	self.transform = clip_transform(tile_size, blur_augment, blur_p)
	self.default_train_tiles = default_train_tiles
	self.is_train = is_train

	self.jpeg_augment=jpeg_augment
	self.jpeg_p = jpeg_p
	self.JpegEnhance = JpegAugment(quality=jpeg_augment, p=jpeg_p, seed=shuffle_seed)
	self.webp_augment = webp_augment
	self.webp_p = webp_p
	self.resize_augment = resize_augment
	self.resize_p = resize_p

	for p in real_dirs:
	find_paths_real = find_files_with_extensions(p)
	self.real_samples += [(f, 0) for f in find_paths_real]
	for p in fake_dirs:
	find_paths_fake = find_files_with_extensions(p)
	self.fake_samples += [(f, 1) for f in find_paths_fake]
	if balance:
	target_len = min(len(self.real_samples), len(self.fake_samples))
	self.real_samples = self.real_samples[:target_len]
	self.fake_samples = self.fake_samples[:target_len]
	self.samples = self.real_samples + self.fake_samples
	# if self.set_size is not None and len(self.samples) < self.set_size:
	# n_repeats = math.ceil(self.set_size / len(self.samples))
	# self.samples = (self.samples * n_repeats)[:self.set_size]
	if self.set_size is not None and len(self.samples) > self.set_size:
	local_random = random.Random(self.shuffle_seed)
	self.samples = local_random.sample(self.samples, self.set_size)
	def __len__(self):
	return len(self.samples)

	def _get_train_item(self, idx):
	path, label = self.samples[idx]
	try:
	img = Image.open(path).convert('RGB')
	except Exception as e:
	print(f"Skipping corrupted image: {path}")
	return self.__getitem__((idx + 1) % len(self.samples))
	if self.resize_augment is not None:
	if self.resize_augment == -1:
	img = resize_with_random_scale(img, p=self.resize_p)
	else:
	if isinstance(self.resize_augment, tuple):
	img = resize_with_random_scale(img, p=self.resize_p, scale_range=self.resize_augment)
	else:
	img = resize_with_random_scale(img, p=self.resize_p, scale_range=(self.resize_augment, self.resize_augment))
	if self.jpeg_augment is not None:
	img = self.JpegEnhance(img)

	# crops_one_side = 2
	# img = pad_crop_image(img, crops_one_sideself.tile_size, crops_one_sideself.tile_size)
	# tiles = crop_image_to_tiles(img, self.tile_size)[:crops_one_side**2]
	tiles = []
	for i in range(self.default_train_tiles): # 10
	tiles.append(random_crop(img, self.tile_size, is_pad_black=False))
	img_resize = self.transform(img)
	img_list = [img_resize]
	for tile in tiles:
	img_list.append(self.transform(tile))
	return img_list, img, label

	def _get_test_item(self, idx):
	path, label = self.samples[idx]
	try:
	img = Image.open(path).convert('RGB')
	except Exception as e:
	print(f"Skipping corrupted image: {path}")
	return self.__getitem__((idx + 1) % len(self.samples))
	if self.resize_augment is not None:
	if self.resize_augment == -1:
	img = resize_with_random_scale(img, p=self.resize_p)
	else:
	if isinstance(self.resize_augment, tuple):
	img = resize_with_random_scale(img, p=self.resize_p, scale_range=self.resize_augment)
	else:
	img = resize_with_random_scale(img, p=self.resize_p, scale_range=(self.resize_augment, self.resize_augment))
	if self.jpeg_augment is not None:
	img = self.JpegEnhance(img)
	if self.webp_augment is not None:
	if random.random() < self.webp_p:
	if isinstance(self.webp_augment, tuple):
	img = webp_compress(img, quality=random.randint(*self.webp_augment))
	else:
	img = webp_compress(img, quality=self.webp_augment)
	# testtime batchsize=1
	tiles = crop_image_to_tiles(img, self.tile_size)
	img_resize = self.transform(img)
	img_list = [img_resize]
	for tile in tiles:
	img_list.append(self.transform(tile))
	return img_list, img, label, path

	def __getitem__(self, idx):
	if self.is_train:
	return self._get_train_item(idx)
	else:
	return self._get_test_item(idx)


	def collate_img_list_to_tensor_train(batch):
	"""
	batch = list of (img_list, label)
	- img_list: List of N Tensor[B, 3, H, W]
	return：
	- images: Tensor[B, N, 3, H, W]
	- labels: Tensor[B]
	"""
	img_lists, raw_imgs, labels = zip(*batch) # len = B
	# img_lists: list of length B, each is List[N x (3,H,W)]
	stacked = torch.stack([torch.stack(imgs) for imgs in img_lists]) # (B, N, 3, H, W)
	raw_imgs = list(raw_imgs)
	labels = torch.tensor(labels)
	return stacked, raw_imgs, labels

	def collate_img_list_to_tensor_test(batch):
	"""
	batch = list of (img_list, label)
	- img_list: List of N Tensor[B, 3, H, W]
	return：
	- images: Tensor[B, N, 3, H, W]
	- labels: Tensor[B]
	"""
	img_lists, raw_imgs, labels, paths = zip(*batch) # len = B
	# img_lists: list of length B, each is List[N x (3,H,W)]
	stacked = torch.stack([torch.stack(imgs) for imgs in img_lists]) # (B, N, 3, H, W)
	raw_imgs = list(raw_imgs)
	labels = torch.tensor(labels)
	return stacked, raw_imgs, labels, paths


	if __name__ == '__main__':
	block_size = 224
	width = 512
	leave_space = min(block_size, width-block_size)
	x_positions = get_start_positions(width-leave_space, block_size, leave_space//2)
	print(leave_space, x_positions)