'''
Write data loader pipeline in Python for YOLOv7
'''
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from yolov7.utils.datasets import LoadImagesAndLabels

def create_dataloader(img_size, batch_size, dataset, data_path, augment=False, multiscale=False):
    # Define image transformations
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
    ])
    
    # Create dataset
    if dataset == 'custom':
        dataset = LoadImagesAndLabels(data_path, img_size=img_size, augment=augment, multiscale=multiscale)
    else:
        dataset = datasets.ImageFolder(data_path, transform=transform)
    
    # Create data loader
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=dataset.collate_fn)
    
    return dataloader