ArtiAgent - DefectDiffu/engine/DefectDiffu/train.py

import os
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.transforms import Lambda
from diffusers.models import AutoencoderKL
import argparse

import clip.clip as clip
from models_add_cross_concate import DiT
from diffusion import create_diffusion

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


# =========================================================================
# 1. TOP-LEVEL HELPER FUNCTION (Prevents Pickle Errors on Windows)
# =========================================================================
def scale_to_neg_one_to_one(t):
    return (t * 2) - 1


# =========================================================================
# 2. DATASET CLASS FOR 4-LEVEL STRUCTURE
# =========================================================================
class Dataset_self(Dataset):
    def __init__(self, img_root, preprocess):
        self.img_root = img_root
        self.img_process = preprocess
        self.img = []
        self.label_word = []
        self.label_mask = []

        # Parse 4-level structure: root / object_class / img / defect_type / image.png
        for name_class in os.listdir(self.img_root):  # e.g., 'vcsel'
            class_path = os.path.join(self.img_root, name_class)
            img_base = os.path.join(class_path, 'img')
            gt_base = os.path.join(class_path, 'ground_truth')

            if os.path.exists(img_base):
                for defect in os.listdir(img_base):  # e.g., 'good', 'scratch', 'bubble', 'crack'
                    defect_img_dir = os.path.join(img_base, defect)
                    defect_gt_dir = os.path.join(gt_base, defect)

                    if os.path.isdir(defect_img_dir):
                        for name_img in os.listdir(defect_img_dir):
                            if name_img.lower().endswith(('.png', '.jpg', '.jpeg')):
                                img_path = os.path.join(defect_img_dir, name_img)
                                
                                # Flexible check: supports '000_mask.png' or '000.png'
                                base_name = os.path.splitext(name_img)[0]
                                mask_candidate_1 = os.path.join(defect_gt_dir, f"{base_name}_mask.png")
                                mask_candidate_2 = os.path.join(defect_gt_dir, name_img)

                                if os.path.exists(mask_candidate_1):
                                    mask_path = mask_candidate_1
                                elif os.path.exists(mask_candidate_2):
                                    mask_path = mask_candidate_2
                                else:
                                    print(f"⚠️ Warning: Mask missing for image {img_path}")
                                    continue

                                self.img.append(img_path)
                                self.label_word.append(f"{defect} {name_class}")
                                self.label_mask.append(mask_path)

        print(f" Successfully loaded {len(self.img)} samples across all defect categories.")

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

    def __getitem__(self, idx):
        img_path = self.img[idx]
        label_mask_path = self.label_mask[idx]
        
        image = Image.open(img_path).convert('RGB')
        label_mask_img = Image.open(label_mask_path).convert('RGB')

        label_mask = self.img_process[1](label_mask_img)
        mask_resize = self.img_process[2](label_mask_img)
        mask_loss = self.img_process[3](label_mask_img)

        mask_loss = mask_loss[0, :, :]
        mask_loss[mask_loss != 0] = 1
        mask_resize_res = torch.cat([mask_resize, mask_resize[0, :, :].unsqueeze(0)], dim=0)

        label = self.label_word[idx]
        image = self.img_process[0](image)

        return image, label, label_mask, mask_resize_res, mask_loss


# =========================================================================
# 3. MAIN TRAINING LOGIC
# =========================================================================
def main(args):
    device = "cuda"
    model_clip, _ = clip.load('RN50', device)

    data_path = args.data
    image_size = args.imagesize
    batch_size = args.batchsize
    latent_size = image_size // 8

    model = DiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, input_size=latent_size, num_classes=1000).to(device)
    state_dict = torch.load(args.ckpt)
    model.load_state_dict(state_dict, strict=False)

    diffusion = create_diffusion(timestep_respacing="")
    vae = AutoencoderKL.from_pretrained(args.vae).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-5, weight_decay=1e-8)

    transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        Lambda(scale_to_neg_one_to_one),
    ])

    transform_mask = transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        Lambda(scale_to_neg_one_to_one),
    ])
    
    transform_resize_mask = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(latent_size),
        transforms.CenterCrop(latent_size),
    ])
    
    transform_mask_loss = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(latent_size // 2),
        transforms.CenterCrop(latent_size // 2),
    ])

    dataset = Dataset_self(img_root=data_path, preprocess=[transform, transform_mask, transform_resize_mask, transform_mask_loss])

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=0,  # 0 workers avoids Windows multiprocessing crashes
        pin_memory=True,
        drop_last=True
    )

    model.train()
    EPOCH = args.epochs

    for epoch in range(EPOCH):
        for x, y, mask, mask_resize, mask_loss in loader:
            x = x.to(device)
            mask = mask.to(device)
            mask_resize = mask_resize.to(device)
            mask_loss = mask_loss.to(device)

            drop_rat = 0.2
            if args.free == 2:
                for i in range(len(y)):
                    c = y[i]
                    if c.split()[0] == 'good':
                        rat_1 = torch.rand(1)
                        if rat_1 < drop_rat:
                            y[i] = 'good industry'
                    else:
                        rat = torch.rand(1)
                        if rat < drop_rat:
                            y[i] = ('good ' + c.split()[1])
            else:
                for i in range(len(y)):
                    c = y[i]
                    if c.split()[0] != 'good':
                        rat_1 = torch.rand(1)
                        if rat_1 < drop_rat:
                            y[i] = ('good ' + c.split()[1])

            defect = torch.cat([clip.tokenize(f"a photo of {c.split()[0]}") for c in y]).to(device)
            classes = torch.cat([clip.tokenize(f"a photo of {c.split()[1]}") for c in y]).to(device)
            y_all = torch.cat([clip.tokenize(f"a photo of {c}") for c in y]).to(device)

            with torch.no_grad():
                defect = model_clip.encode_text(defect)
                classes = model_clip.encode_text(classes)
                y_all = model_clip.encode_text(y_all)

            defect /= defect.norm(dim=-1, keepdim=True)
            defect = defect.float().to(device)

            classes /= classes.norm(dim=-1, keepdim=True)
            classes = classes.float().to(device)

            y_all /= y_all.norm(dim=-1, keepdim=True)
            y_all = y_all.float().to(device)

            with torch.no_grad():
                x = vae.encode(x).latent_dist.sample().mul_(0.18215)
                mask_gt = vae.encode(mask).latent_dist.sample().mul_(0.18215)
                
            t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
            model_kwargs = dict(y=[defect, classes, y_all])
            loss_dict = diffusion.training_losses(model, x, t, model_kwargs, mask_resize=mask_resize, mask_att=mask_loss, label_mask=mask_gt)
            loss = loss_dict["loss"].mean()
            
            opt.zero_grad()
            loss.backward()
            opt.step()
            print(f"Epoch {epoch} | Loss: {loss.item():.4f}")

        if epoch % 100 == 0 and 2000 >= epoch >= 100:
            os.makedirs('checkpoint', exist_ok=True)
            torch.save({
                'model_state_dict': model.state_dict(),
            }, f'checkpoint/model_{epoch}.pth')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batchsize", type=int, default=2)
    parser.add_argument("--free", type=int, default=1)
    parser.add_argument("--data", type=str, required=True)
    parser.add_argument("--imagesize", type=int, choices=[256, 512], default=256)
    parser.add_argument("--ckpt", type=str, required=True, help="Optional path to a DiT checkpoint.")
    parser.add_argument("--vae", type=str, required=True, help="Optional path to a vae checkpoint.")
    parser.add_argument("--epochs", type=int, default=501, help="Number of training epochs.")
    args = parser.parse_args()
    main(args)