UniTrain is an open-source, unified platform for effortless machine learning model training, evaluation, and deployment across diverse tasks. Experience seamless experimentation and model deployment with UniTrain.
import os
import glob
from PIL import Image
import torch
from torch.utils.data import Dataset

class DCGANdataset:
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.classes = os.listdir(data_dir)
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
        self.data = self.load_data()

    def load_data(self):
        data = []
        for cls in self.classes:
            class_path = os.path.join(self.data_dir, cls)
            class_idx = self.class_to_idx[cls]
            # Include .jpg, .jpeg, and .png extensions in the glob pattern
            for file_path in glob.glob(os.path.join(class_path, '*.jpg')) + \
                             glob.glob(os.path.join(class_path, '*.jpeg')) + \
                             glob.glob(os.path.join(class_path, '*.png')):
                data.append((file_path, class_idx))
        return data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path, target = self.data[idx]
        image = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            image = self.transform(image)

        return image, target

UniTrain/dataset/DCGAN.py
ishan121028/UniTrain
from .classification import ClassificationDataset
from .segmentation import SegmentationDataset
from .DCGAN import DCGANdataset
UniTrain/dataset/__init__.py
ishan121028/UniTrain
import os
import glob
from PIL import Image
import torch
from torch.utils.data import Dataset

class ClassificationDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.classes = os.listdir(data_dir)
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
        self.data = self.load_data()

    def load_data(self):
        data = []
        for cls in self.classes:
            class_path = os.path.join(self.data_dir, cls)
            class_idx = self.class_to_idx[cls]
            # Include .jpg, .jpeg, and .png extensions in the glob pattern
            for file_path in glob.glob(os.path.join(class_path, '*.jpg')) + \
                             glob.glob(os.path.join(class_path, '*.jpeg')) + \
                             glob.glob(os.path.join(class_path, '*.png')):
                data.append((file_path, class_idx))
        return data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path, target = self.data[idx]
        image = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            image = self.transform(image)

        return image, target

UniTrain/dataset/classification.py
ishan121028/UniTrain
import os
from PIL import Image
import torch
from torch.utils.data import Dataset
import cv2
import torchvision.transforms as transforms

class SegmentationDataset(Dataset):
    def __init__(self, image_paths: list, mask_paths: list, transform=None):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        # Load image and mask based on index
        image_path = self.image_paths[index]
        mask_path = self.mask_paths[index]

        print(image_path)

        image = Image.open(image_path).convert('RGB')
        mask = Image.open(mask_path).convert('L')

        img = cv2.imread(image_path)

        if self.transform is not None:
            image = self.transform(image)
            
        mask_transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
        mask = mask_transform(mask)
        mask = mask.to(torch.long)
        
        # You may need to further preprocess the mask if required
        # Example: Convert mask to tensor and perform class mapping
        return image, mask

UniTrain/dataset/segmentation.py
ishan121028/UniTrain
from torchsummary import summary
import torch.nn as nn
    

class disc(nn.Module):
    
    discriminator = nn.Sequential(
        # in: 3 x 64 x 64

        # Layer 01 : Converts a 3 Channel Image to a 64 Channel Image (FEATURE MAP)
        nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(64),
        nn.LeakyReLU(0.2, inplace=True),
        # out: 64 x 32 x 32 
        
        nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(128),
        nn.LeakyReLU(0.2, inplace=True),
        # out: 128 x 16 x 16

        nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(256),
        nn.LeakyReLU(0.2, inplace=True),
        # out: 256 x 8 x 8

        nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(512),
        nn.LeakyReLU(0.2, inplace=True),
        # out: 512 x 4 x 4

        nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),
        # out: 1 x 1 x 1
        

        nn.Flatten(),
        nn.Sigmoid()
        
     )

    
class gen(nn.Module):

    latent_size = 128

    generator = nn.Sequential(
        # in: latent_size x 1 x 1

        nn.ConvTranspose2d(latent_size, 512, kernel_size=4, stride=1, padding=0, bias=False),
        nn.BatchNorm2d(512),
        nn.ReLU(True),
        # out: 512 x 4 x 4

        nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(256),
        nn.ReLU(True),
        # out: 256 x 8 x 8

        nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(128),
        nn.ReLU(True),
        # out: 128 x 16 x 16

        nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(64),
        nn.ReLU(True),
        # out: 64 x 32 x 32

        nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False),
        nn.Tanh()
        # out: 3 x 64 x 64
    )
UniTrain/models/DCGAN.py
ishan121028/UniTrain
from .classification import ResNet9
from .segmentation import UNet
from .DCGAN import *
UniTrain/models/__init__.py
ishan121028/UniTrain
import torch
import torch.nn as nn

# Define the ResNet-9 model in a single class
class ResNet9(nn.Module):
    def __init__(self, num_classes):
        super(ResNet9, self).__init__()
        self.in_channels = 64

        # Initial convolution layer
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        # Residual blocks
        self.layer1 = self.make_layer(64, 2, stride=1)
        self.layer2 = self.make_layer(128, 2, stride=2)
        self.layer3 = self.make_layer(256, 2, stride=2)
        self.layer4 = self.make_layer(512, 2, stride=2)

        # Global average pooling and fully connected layer
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)

    def make_layer(self, out_channels, num_blocks, stride):
        layers = []
        layers.append(self.build_residual_block(self.in_channels, out_channels, stride))
        self.in_channels = out_channels
        for _ in range(1, num_blocks):
            layers.append(self.build_residual_block(self.in_channels, out_channels, stride=1))
        return nn.Sequential(*layers)

    def build_residual_block(self, in_channels, out_channels, stride):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

UniTrain/models/classification.py
ishan121028/UniTrain
import torch
from torch import nn
from torch.nn import functional as F    

class UNet(nn.Module):
    def __init__(self, n_class):
        super().__init__()
        
        # Encoder
        # In the encoder, convolutional layers with the Conv2d function are used to extract features from the input image. 
        # Each block in the encoder consists of two convolutional layers followed by a max-pooling layer, with the exception of the last block which does not include a max-pooling layer.
        # -------
        # input: 572x572x3
        self.e11 = nn.Conv2d(3, 64, kernel_size=3, padding=1) # output: 570x570x64
        self.e12 = nn.Conv2d(64, 64, kernel_size=3, padding=1) # output: 568x568x64
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) # output: 284x284x64

        # input: 284x284x64
        self.e21 = nn.Conv2d(64, 128, kernel_size=3, padding=1) # output: 282x282x128
        self.e22 = nn.Conv2d(128, 128, kernel_size=3, padding=1) # output: 280x280x128
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # output: 140x140x128

        # input: 140x140x128
        self.e31 = nn.Conv2d(128, 256, kernel_size=3, padding=1) # output: 138x138x256
        self.e32 = nn.Conv2d(256, 256, kernel_size=3, padding=1) # output: 136x136x256
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) # output: 68x68x256

        # input: 68x68x256
        self.e41 = nn.Conv2d(256, 512, kernel_size=3, padding=1) # output: 66x66x512
        self.e42 = nn.Conv2d(512, 512, kernel_size=3, padding=1) # output: 64x64x512
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) # output: 32x32x512

        # input: 32x32x512
        self.e51 = nn.Conv2d(512, 1024, kernel_size=3, padding=1) # output: 30x30x1024
        self.e52 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1) # output: 28x28x1024


        # Decoder
        self.upconv1 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.d11 = nn.Conv2d(1024, 512, kernel_size=3, padding=1)
        self.d12 = nn.Conv2d(512, 512, kernel_size=3, padding=1)

        self.upconv2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.d21 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.d22 = nn.Conv2d(256, 256, kernel_size=3, padding=1)

        self.upconv3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.d31 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
        self.d32 = nn.Conv2d(128, 128, kernel_size=3, padding=1)

        self.upconv4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.d41 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
        self.d42 = nn.Conv2d(64, 64, kernel_size=3, padding=1)

        # Output layer
        self.outconv = nn.Conv2d(64, n_class, kernel_size=1)
    
    def forward(self, x):
        # Encoder
        xe11 = F.relu(self.e11(x))
        xe12 = F.relu(self.e12(xe11))
        xp1 = self.pool1(xe12)

        xe21 = F.relu(self.e21(xp1))
        xe22 = F.relu(self.e22(xe21))
        xp2 = self.pool2(xe22)

        xe31 = F.relu(self.e31(xp2))
        xe32 = F.relu(self.e32(xe31))
        xp3 = self.pool3(xe32)

        xe41 = F.relu(self.e41(xp3))
        xe42 = F.relu(self.e42(xe41))
        xp4 = self.pool4(xe42)

        xe51 = F.relu(self.e51(xp4))
        xe52 = F.relu(self.e52(xe51))
        
        # Decoder
        xu1 = self.upconv1(xe52)
        xu11 = torch.cat([xu1, xe42], dim=1)
        xd11 = F.relu(self.d11(xu11))
        xd12 = F.relu(self.d12(xd11))

        xu2 = self.upconv2(xd12)
        xu22 = torch.cat([xu2, xe32], dim=1)
        xd21 = F.relu(self.d21(xu22))
        xd22 = F.relu(self.d22(xd21))

        xu3 = self.upconv3(xd22)
        xu33 = torch.cat([xu3, xe22], dim=1)
        xd31 = F.relu(self.d31(xu33))
        xd32 = F.relu(self.d32(xd31))

        xu4 = self.upconv4(xd32)
        xu44 = torch.cat([xu4, xe12], dim=1)
        xd41 = F.relu(self.d41(xu44))
        xd42 = F.relu(self.d42(xd41))

        # Output layer
        out = self.outconv(xd42)

        return out
UniTrain/models/segmentation.py
ishan121028/UniTrain
from models.segmentation import UNet
from utils.segmentation import parse_folder, get_data_loader, train_unet
from torchvision import transforms
import glob


def main():
    if parse_folder('data'):


        # Make Your Custom Data Transformations 
        # transform = transforms.Compose([
        #     transforms.Resize((224, 224)),  # Resize images to a fixed size
        #     transforms.ToTensor(),  # Convert images to PyTorch tensors
        #     transforms.Normalize((0.485, 0.456, 0.406),
        #                          (0.229, 0.224, 0.225))  # Normalize with ImageNet stats
        # ])
        train_image_paths = glob.glob("data/train/images/*.jpg")
        train_mask_paths = glob.glob("data/train/masks/*.png")

        test_image_paths = glob.glob("data/test/images/*.jpg")
        test_mask_paths = glob.glob("data/test/masks/*.png")

        print(train_image_paths, train_mask_paths, test_image_paths, test_mask_paths)

        train_dataloader = get_data_loader(train_image_paths,train_mask_paths, 1, True)
        test_dataloader = get_data_loader(test_image_paths, test_mask_paths, 1, True)

        model = UNet(n_class=20)

        train_unet(model, train_dataloader, test_dataloader, num_epochs=10, learning_rate=1e-3, checkpoint_dir='checkpoints')

    else:
        print("Invalid dataset folder.")
        return None

if __name__ == '__main__':
    main()
UniTrain/train.py
ishan121028/UniTrain
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as T
import torch
import os
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
from torchvision.utils import save_image
from UniTrain.dataset.DCGAN import DCGANdataset
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as T
import torch
import os
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
from torchvision.utils import save_image
from tqdm.notebook import tqdm
import torch.nn.functional as F


latent_size = 128
stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)

def denorm(img_tensors):
    return img_tensors * stats[1][0] + stats[0][0]

def get_data_loader(data_dir, batch_size, shuffle=True, transform = None, split='real'):
    """
    Create and return a data loader for a custom dataset.

    Args:
        data_dir (str): Path to the dataset directory.
        batch_size (int): Batch size for the data loader.
        shuffle (bool): Whether to shuffle the data (default is True).

    Returns:
        DataLoader: PyTorch data loader.
    """
    
    # Define data transformations (adjust as needed)
    if split == 'real':
        data_dir = os.path.join(data_dir, 'real_images')
    else:
        raise ValueError(f"Invalid split choice: {split}")

    image_size = 64
    batch_size = 128
    stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)

    if transform is None:
        transform=T.Compose([T.Resize(image_size), T.CenterCrop(image_size),T.ToTensor(),T.Normalize(*stats)])

    # Create a custom dataset
    dataset = DCGANdataset(data_dir, transform=transform)

    # Create a data loader
    print(batch_size)
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle
    )

    return data_loader


def parse_folder(dataset_path):
    print(dataset_path)
    print(os.getcwd())
    print(os.path.exists(dataset_path))
    try:
        if os.path.exists(dataset_path):
            # Store paths to train, test, and eval folders if they exist
            real_path = os.path.join(dataset_path, 'real_images')

            if os.path.exists(real_path) :
                print("Real Data folder path:", real_path)

                real_classes = set(os.listdir(real_path))

                if real_classes:
                    return real_classes
                else:
                    print("Real Data Set is empty")
                    return None
            else:
                print("One or more of the train, test, or eval folders does not exist.")
                return None
        else:
            print(f"The '{dataset_path}' folder does not exist in the current directory.")
            return None
    except Exception as e:
        print("An error occurred:", str(e))
        return None
        
def train_discriminator(discriminator, generator, real_images, opt_d, batch_size, latent_size, device):
    opt_d.zero_grad()

    
    print(real_images.shape)
    real_preds = discriminator(real_images)
    print(real_preds.shape)
    real_targets = torch.ones(real_images.size(0), 1, device=device)
    real_loss = F.binary_cross_entropy(real_preds, real_targets)
    real_score = torch.mean(real_preds).item()

    # Generate fake images
    latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
    fake_images = generator(latent)
    # Pass fake images through discriminator
    fake_targets = torch.zeros(fake_images.size(0), 1, device=device)
    fake_preds = discriminator(fake_images)
    fake_loss = F.binary_cross_entropy(fake_preds, fake_targets)
    fake_score = torch.mean(fake_preds).item()

    # Update discriminator weights
    loss = real_loss + fake_loss
    loss.backward()
    opt_d.step()
    return loss.item(), real_score, fake_score


def train_generator(opt_g, discriminator, generator, batch_size, device):
    # Clear generator gradients
    opt_g.zero_grad()

    # Generate fake images
    latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
    fake_images = generator(latent)

    # Try to fool the discriminator
    preds = discriminator(fake_images)
    targets = torch.ones(batch_size, 1, device=device)
    loss = F.binary_cross_entropy(preds, targets)

    # Update generator weights
    loss.backward()
    opt_g.step()

    return loss.item()

generated_dir = 'generated'
os.makedirs(generated_dir, exist_ok=True)

def save_samples(index, generator_model, latent_tensors, show=True):
    fake_images = generator_model(latent_tensors)
    fake_fname = 'generated-images-{0:0=4d}.png'.format(index)
    save_image(denorm(fake_images), os.path.join(generated_dir, fake_fname), nrow=8)
    print('Saving', fake_fname)
    if show:
        fig, ax = plt.subplots(figsize=(8, 8))
        ax.set_xticks([]); ax.set_yticks([])
        ax.imshow(make_grid(fake_images.cpu().detach(), nrow=8).permute(1, 2, 0))


def train_model(discriminator_model, generator_model, train_data_loader , batch_size, epochs, learning_rate, checkpoint_dir , device= torch.device('cpu') ,logger=None, iou=False, ):

    os.makedirs(checkpoint_dir + "/discriminator_checkpoint" , exist_ok=True)
    os.makedirs(checkpoint_dir + "/generator_checkpoint" , exist_ok=True)

    fixed_latent = torch.randn(128, latent_size, 1, 1, device=device)

    # Losses & scores
    losses_g = []
    losses_d = []
    real_scores = []
    fake_scores = []

    # Create optimizers
    opt_d = torch.optim.Adam(discriminator_model.parameters(), lr=learning_rate, betas=(0.5, 0.999))
    opt_g = torch.optim.Adam(generator_model.parameters(), lr=learning_rate, betas=(0.5, 0.999))

    # e = 0
    for epoch in range(epochs):
        i=0
        for real_images, _ in tqdm(train_data_loader):
            # Train discriminator
            print('epochs',epoch,'image',i)
            i=i+1
            loss_d, real_score, fake_score = train_discriminator(discriminator_model, generator_model, real_images, opt_d,128,128, device='cpu' )
            # Train generator
            loss_g = train_generator(opt_g, discriminator_model, generator_model, batch_size, device='cpu')

        # Record losses & scores
        losses_g.append(loss_g)
        losses_d.append(loss_d)
        real_scores.append(real_score)
        fake_scores.append(fake_score)

        # Save generated images
        save_samples(epoch+epoch,generator_model , fixed_latent, show=False)

    print('Finished Training')

def evaluate_model(discriminator_model, dataloader):
    discriminator_model.eval()  # Set the model to evaluation mode
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in dataloader:
            outputs = discriminator_model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total

    return accuracy



from tqdm.notebook import tqdm
# from tqdm.notebook import tqdm
import torch.nn.functional as F
# %matplotlib inline

latent_size = 128
stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)

def denorm(img_tensors):
    return img_tensors * stats[1][0] + stats[0][0]

def get_data_loader(data_dir, batch_size, shuffle=True, transform = None, split='real'):
    """
    Create and return a data loader for a custom dataset.

    Args:
        data_dir (str): Path to the dataset directory.
        batch_size (int): Batch size for the data loader.
        shuffle (bool): Whether to shuffle the data (default is True).

    Returns:
        DataLoader: PyTorch data loader.
    """
    # Define data transformations (adjust as needed)

    if split == 'real':
        data_dir = os.path.join(data_dir, 'real_images')
    else:
        raise ValueError(f"Invalid split choice: {split}")

    image_size = 64
    batch_size = 128
    stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)

    if transform is None:
        transform=T.Compose([T.Resize(image_size), T.CenterCrop(image_size),T.ToTensor(),T.Normalize(*stats)])

    # Create a custom dataset
    dataset = DCGANdataset(data_dir, transform=transform)

    # Create a data loader
    print(batch_size)
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle
    )

    return data_loader



def parse_folder(dataset_path):
    print(dataset_path)
    print(os.getcwd())
    print(os.path.exists(dataset_path))
    try:
        if os.path.exists(dataset_path):
            # Store paths to train, test, and eval folders if they exist
            real_path = os.path.join(dataset_path, 'real_images')

            # if os.path.exists(real_path) and os.path.exists(test_path) and os.path.exists(eval_path):
            if os.path.exists(real_path) :
                print("Real Data folder path:", real_path)

                real_classes = set(os.listdir(real_path))

                if real_classes:
                    return real_classes
                else:
                    print("Real Data Set is empty")
                    return None
            else:
                print("One or more of the train, test, or eval folders does not exist.")
                return None
        else:
            print(f"The '{dataset_path}' folder does not exist in the current directory.")
            return None
    except Exception as e:
        print("An error occurred:", str(e))
        return None
    
def train_discriminator(discriminator, generator, real_images, opt_d, batch_size, latent_size, device):
    # Clear discriminator gradients
    opt_d.zero_grad()

    # Pass real images through discriminator
    print(real_images.shape)
    real_preds = discriminator(real_images)
    print(real_preds.shape)
    real_targets = torch.ones(real_images.size(0), 1, device=device)
    real_loss = F.binary_cross_entropy(real_preds, real_targets)
    real_score = torch.mean(real_preds).item()

    # Generate fake images
    latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
    fake_images = generator(latent)
    # Pass fake images through discriminator
    fake_targets = torch.zeros(fake_images.size(0), 1, device=device)
    fake_preds = discriminator(fake_images)
    fake_loss = F.binary_cross_entropy(fake_preds, fake_targets)
    fake_score = torch.mean(fake_preds).item()

    # Update discriminator weights
    loss = real_loss + fake_loss
    loss.backward()
    opt_d.step()
    return loss.item(), real_score, fake_score


def train_generator(opt_g, discriminator, generator, batch_size, device):
    # Clear generator gradients
    opt_g.zero_grad()

    # Generate fake images
    latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
    fake_images = generator(latent)

    # Try to fool the discriminator
    preds = discriminator(fake_images)
    targets = torch.ones(batch_size, 1, device=device)
    loss = F.binary_cross_entropy(preds, targets)

    # Update generator weights
    loss.backward()
    opt_g.step()

    return loss.item()

generated_dir = 'generated'
os.makedirs(generated_dir, exist_ok=True)

def save_samples(index, generator_model, latent_tensors, show=True):
    fake_images = generator_model(latent_tensors)
    fake_fname = 'generated-images-{0:0=4d}.png'.format(index)
    save_image(denorm(fake_images), os.path.join(generated_dir, fake_fname), nrow=8)
    print('Saving', fake_fname)
    if show:
        fig, ax = plt.subplots(figsize=(8, 8))
        ax.set_xticks([]); ax.set_yticks([])
        ax.imshow(make_grid(fake_images.cpu().detach(), nrow=8).permute(1, 2, 0))


def train_model(discriminator_model, generator_model, train_data_loader ,batch_size, epochs, learning_rate, checkpoint_dir,device="cpu" , logger=None, iou=False, ):
    if device == 'cpu':
        device = torch.device('cpu')
    elif device == 'cuda':
        device = torch.device('cuda')
    else:
        print(f"{device} is not a valid device.")
        return None

    os.makedirs(checkpoint_dir + "/discriminator_checkpoint")
    os.makedirs(checkpoint_dir + "/generator_checkpoint")

    fixed_latent = torch.randn(128, latent_size, 1, 1, device=device)

    # Losses & scores
    losses_g = []
    losses_d = []
    real_scores = []
    fake_scores = []

    # Create optimizers
    opt_d = torch.optim.Adam(discriminator_model.parameters(), lr=learning_rate, betas=(0.5, 0.999))
    opt_g = torch.optim.Adam(generator_model.parameters(), lr=learning_rate, betas=(0.5, 0.999))

    # e = 0
    for epoch in range(epochs):
        i=0
        for real_images, _ in tqdm(train_data_loader):
            # Train discriminator
            print('epochs',epoch,'image',i)
            i=i+1
            loss_d, real_score, fake_score = train_discriminator(discriminator_model, generator_model, real_images, opt_d,128,128, device='cpu' )
            # Train generator
            loss_g = train_generator(opt_g, discriminator_model, generator_model, batch_size, device='cpu')

        # Record losses & scores
        losses_g.append(loss_g)
        losses_d.append(loss_d)
        real_scores.append(real_score)
        fake_scores.append(fake_score)

        # Save generated images
        save_samples(epoch+epoch,generator_model , fixed_latent, show=False)

    print('Finished Training')

def evaluate_model(discriminator_model, dataloader):
    discriminator_model.eval()  # Set the model to evaluation mode
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in dataloader:
            outputs = discriminator_model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total

    return accuracy


UniTrain/utils/DCGAN.py
ishan121028/UniTrain
from .classification import get_data_loader, parse_folder, train_model
from .segmentation import get_data_loader, parse_folder, train_unet, generate_model_summary, iou_score
from .DCGAN import get_data_loader, parse_folder , train_model
UniTrain/utils/__init__.py
ishan121028/UniTrain
import os
from torch.utils.data import DataLoader
from torchvision import transforms
from ..dataset.classification import ClassificationDataset
import torch.optim as optim
import torch.nn as nn
import torch
import logging

def get_data_loader(data_dir, batch_size, shuffle=True, transform = None, split='train'):
    """
    Create and return a data loader for a custom dataset.

    Args:
        data_dir (str): Path to the dataset directory.
        batch_size (int): Batch size for the data loader.
        shuffle (bool): Whether to shuffle the data (default is True).

    Returns:
        DataLoader: PyTorch data loader.
    """
    # Define data transformations (adjust as needed)

    if split == 'train':
        data_dir = os.path.join(data_dir, 'train')
    elif split == 'test':
        data_dir = os.path.join(data_dir, 'test')
    elif split == 'eval':
        data_dir = os.path.join(data_dir, 'eval')
    else:
        raise ValueError(f"Invalid split choice: {split}")



    if transform is None:
        transform = transforms.Compose([
            transforms.Resize((224, 224)),  # Resize images to a fixed size
            transforms.ToTensor(),  # Convert images to PyTorch tensors
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # Normalize with ImageNet stats
        ])

    # Create a custom dataset
    dataset = ClassificationDataset(data_dir, transform=transform)

    # Create a data loader
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle
    )

    return data_loader


def parse_folder(dataset_path):
    try:
        if os.path.exists(dataset_path):
            # Store paths to train, test, and eval folders if they exist
            train_path = os.path.join(dataset_path, 'train')
            test_path = os.path.join(dataset_path, 'test')
            eval_path = os.path.join(dataset_path, 'eval')

            if os.path.exists(train_path) and os.path.exists(test_path) and os.path.exists(eval_path):
                print("Train folder path:", train_path)
                print("Test folder path:", test_path)
                print("Eval folder path:", eval_path)

                train_classes = set(os.listdir(train_path))
                test_classes = set(os.listdir(test_path))
                eval_classes = set(os.listdir(eval_path))

                if train_classes == test_classes == eval_classes:
                    return train_classes, train_path, test_path, eval_path
                else:
                    print("Classes are not the same in train, test, and eval folders.")
                    return None
            else:
                print("One or more of the train, test, or eval folders does not exist.")
                return None
        else:
            print(f"The '{dataset_path}' folder does not exist in the current directory.")
            return None
    except Exception as e:
        print("An error occurred:", str(e))
        return None

def train_model(model, train_data_loader, test_data_loader, num_epochs, learning_rate=0.001, checkpoint_dir='checkpoints', logger=None, device=torch.device('cpu')):
    '''Train a PyTorch model for a classification task.
    Args:
    model (nn.Module): Torch model to train.
    train_data_loader (DataLoader): Training data loader.
    test_data_loader (DataLoader): Testing data loader.
    num_epochs (int): Number of epochs to train the model for.
    learning_rate (float): Learning rate for the optimizer.
    checkpoint_dir (str): Directory to save model checkpoints.
    logger (Logger): Logger to log training details.
    device (torch.device): Device to run training on (GPU or CPU).

    Returns:
    None
    '''

    if logger:
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - Epoch %(epoch)d - Train Acc: %(train_acc).4f - Val Acc: %(val_acc).4f - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S', filename=logger, filemode='w')
        logger = logging.getLogger(__name__)

    # Define loss function and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    best_accuracy = 0.0

    # Training loop
    for epoch in range(num_epochs):
        model.train()  # Set the model to training mode
        running_loss = 0.0

        for batch_idx, (inputs, labels) in enumerate(train_data_loader):
            optimizer.zero_grad()  # Zero the parameter gradients

            inputs = inputs.to(device)
            labels = labels.to(device)

            # Forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            if batch_idx % 100 == 99:  # Print and log every 100 batches
                avg_loss = running_loss / 100
                if logger:
                    logger.info(f'Epoch {epoch + 1}, Batch {batch_idx + 1}, Loss: {avg_loss:.4f}')
                print(f'Epoch {epoch + 1}, Batch {batch_idx + 1}, Loss: {avg_loss:.4f}')
                running_loss = 0.0

        # Save model checkpoint if accuracy improves
        accuracy = evaluate_model(model, test_data_loader)
        if logger:
            logger.info(f'Epoch {epoch + 1}, Validation Accuracy: {accuracy:.2f}%')


        if accuracy > best_accuracy:
            best_accuracy = accuracy
            checkpoint_path = os.path.join(checkpoint_dir, f'model_epoch_{epoch + 1}.pth')
            torch.save(model.state_dict(), checkpoint_path)
            if logger:
                logger.info(f'Saved checkpoint to {checkpoint_path}')

    print('Finished Training')


def evaluate_model(model, dataloader):
    model.eval()  # Set the model to evaluation mode
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in dataloader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total

    return accuracy

    

UniTrain/utils/classification.py
ishan121028/UniTrain
import os
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from torch.utils.data import DataLoader
from torchvision import transforms
from ..dataset.segmentation import SegmentationDataset
import torchsummary
import logging
import glob

def get_data_loader(data_dir: str, batch_size:int, shuffle:bool=True, transform=None, split='train') -> DataLoader:
    """,
    Create and return a data loader for a custom dataset.

    Args:
        data_dir (str): Path to the dataset directory.
        batch_size (int): Batch size for the data loader.
        shuffle (bool): Whether to shuffle the data (default is True).

    Returns:
        DataLoader: PyTorch data loader.
    """
    # Define data transformations (adjust as needed)
    if transform is None:
        transform = transforms.Compose([
            transforms.Resize((224, 224)),  # Resize images to a fixed size
            transforms.ToTensor(),  # Convert images to PyTorch tensors
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # Normalize with ImageNet stats
        ])

    image_paths = glob.glob(os.path.join(data_dir, split, 'images', '*'))
    mask_paths = glob.glob(os.path.join(data_dir, split, 'masks', '*'))

    # Create a custom dataset
    dataset = SegmentationDataset(image_paths=image_paths, mask_paths=mask_paths, transform=transform)

    # Create a data loader
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle
    )

    return data_loader

def parse_folder(dataset_path):
    '''Parse the dataset folder and return True if the folder structure is valid, False otherwise.

    Args:
    dataset_path (str): Path to the dataset folder.

    Returns:
    bool: Whether the dataset folder is valid.
    '''
    try:
        if os.path.exists(dataset_path):
            # Store paths to train, test, and eval folders if they exist
            train_path = os.path.join(dataset_path, "train")
            test_path = os.path.join(dataset_path, "test")
            eval_path = os.path.join(dataset_path, "eval")

            if os.path.exists(train_path) and os.path.exists(test_path) and os.path.exists(eval_path):

                print(f"Train path: {train_path}")
                print(f"Test path: {test_path}")
                print(f"Eval path: {eval_path}")

                root_dir_list = os.listdir(dataset_path)

                for dir in root_dir_list:
                    masks_path = os.path.join(dataset_path, dir, "masks")
                    images_path = os.path.join(dataset_path, dir, "images")

                    if os.path.exists(masks_path) and os.path.exists(images_path):
                        pass
                    else:
                        return False
                
                return True

        else:
            print(f"The '{dataset_path}' folder does not exist in the current directory.")
            return False
    except Exception as e:
        print("An error occurred:", str(e))
        return False


def train_unet(model, train_data_loader, test_data_loader, num_epochs, learning_rate, checkpoint_dir, logger=None, iou=False, device=torch.device('cpu')) -> None:
    '''Train the model using the given train and test data loaders.
    
    Args: 
    model (nn.Module): PyTorch model to train.
    train_data_loader (DataLoader): Data loader of the training dataset.
    test_data_loader (DataLoader): Data loader of the test dataset.
    num_epochs (int): Number of epochs to train the model.
    learning_rate (float): Learning rate for the optimizer.
    checkpoint_dir (str): Directory to save model checkpoints.
    logger (Logger): Logger to log training information.
    iou (bool): Whether to calculate the IOU score.
    device (torch.device): Device to run training on (GPU or CPU).

    Returns:
    None
    '''

    if logger:
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - Epoch %(epoch)d - Train Acc: %(train_acc).4f - Val Acc: %(val_acc).4f - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S', filename=logger, filemode='w')
        logger = logging.getLogger(__name__)


    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    best_loss = float('inf')

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0.0
        iou_score_mean = 0.0

        for inputs, targets in tqdm(train_data_loader, desc=f'Epoch {epoch + 1}/{num_epochs}', leave=False):
            optimizer.zero_grad()
            outputs = model(inputs)
            targets = targets.squeeze(1)
            outputs.to(device)
            targets.to(device)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            iou_score_mean += iou_score(outputs, targets)

        iou_score_mean = iou_score_mean / len(train_data_loader)
        average_train_loss = train_loss / len(train_data_loader)

        if logger and iou:
            logger.info(f'Epoch {epoch + 1}/{num_epochs}, Train Loss: {average_train_loss:.4f}, IOU Score: {iou_score_mean:.4f}')
        elif logger is not None:
            logger.info(f"Epoch {epoch + 1}/{num_epochs}, Train Loss: {average_train_loss:.4f}")

        # Validation
        model.eval()
        val_loss = 0.0
        iou_score_mean = 0.0
        with torch.no_grad():
            for inputs, targets in tqdm(test_data_loader, desc=f'Validation', leave=False):
                outputs = model(inputs)
                targets = targets.squeeze(1)
                outputs.to(device)
                targets.to(device)
                loss = criterion(outputs, targets)
                val_loss += loss.item()
                iou_score_mean += iou_score(outputs, targets)

        iou_score_mean = iou_score_mean / len(test_data_loader)
        average_val_loss = val_loss / len(test_data_loader)

        if logger and iou:
            logger.info(f'Epoch {epoch + 1}/{num_epochs}, Validation Loss: {average_val_loss:.4f}. IOU Score: {iou_score_mean:.4f}')
        elif logger is not None:
            logger.info(f"Epoch {epoch + 1}/{num_epochs}, Validation Loss: {average_val_loss:.4f}")

        # Save model checkpoint if validation loss improves
        if average_val_loss < best_loss:
            best_loss = average_val_loss
            checkpoint_path = f'{checkpoint_dir}/unet_model_epoch_{epoch + 1}.pth'
            torch.save(model.state_dict(), checkpoint_path)
            if logger:
                logger.info(f'Saved checkpoint to {checkpoint_path}')

    print('Finished Training')

def generate_model_summary(model, input_size):
    '''Generate a summary of the model.'''
    torchsummary.summary(model, input_size=input_size)

def iou_score(output, target):
    '''Computes the intersection over union (IoU) for the given prediction and target masks.
    Args:
    output (torch.Tensor): Predicted masks of shape (N x H x W).
    target (torch.Tensor): Target masks of shape (N x H x W).

    Returns:
    float: The average IoU score.
    '''
    smooth = 1e-6
    output = output.argmax(1)
    intersection = (output & target).float().sum((1, 2))
    union = (output | target).float().sum((1, 2))
    iou = (intersection + smooth) / (union + smooth)
    return iou.mean().item()
UniTrain/utils/segmentation.py
ishan121028/UniTrain
from setuptools import setup, find_packages

setup(
    name="UniTrain",
    version="0.2.4",
    author="Ishan Upadhyay",
    author_email="ishan.upadhyay.iitb@gmail.com",
    description="A generalized training framework for Deep Learning Tasks",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    url="https://github.com/ishan121028/UniTrain",
    download_url="https://github.com/ishan121028/UniTrain/archive/v0.1.tar.gz",
    packages=find_packages(),
    install_requires=[
        "torch",
        "torchvision",
        "numpy",
        "pandas",
        "torchsummary",
        "opencv-python",
    ],
    keywords=["Deep Learning", "Machine Learning", "Training Framework"],
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Developers",
        "Intended Audience :: Education",
        "License :: OSI Approved :: MIT License",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
)

setup.py
ishan121028/UniTrain