diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/.gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/autoencoder.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..8772bc03371ffcf5e4b5d646395bf62cfac2229e --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/autoencoder.py @@ -0,0 +1,584 @@ +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange +import os +import torchvision.transforms as transforms +from torchvision.utils import save_image +import os +from PIL import Image + + +os.environ["CUDA_VISIBLE_DEVICES"] = "1" +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3) + k = k.softmax(dim=-1) + context = torch.einsum('bhdn,bhen->bhde', k, v) + out = torch.einsum('bhde,bhdn->bhen', context, q) + out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) + return self.to_out(out) + + +def nonlinearity(x): + # swish + return x*torch.sigmoid(x) + + +def Normalize(in_channels, num_groups=32): + return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + + +class Upsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=3, + stride=2, + padding=0) + + def forward(self, x): + if self.with_conv: + pad = (0,1,0,1) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + return x + + +class ResnetBlock(nn.Module): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, + dropout, temb_channels=512): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + if temb_channels > 0: + self.temb_proj = torch.nn.Linear(temb_channels, + out_channels) + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = torch.nn.Conv2d(out_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + else: + self.nin_shortcut = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0) + + def forward(self, x, temb): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None] + + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x+h + + +class LinAttnBlock(LinearAttention): + """to match AttnBlock usage""" + def __init__(self, in_channels): + super().__init__(dim=in_channels, heads=1, dim_head=in_channels) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + self.k = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + self.v = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + self.proj_out = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b,c,h,w = q.shape + q = q.reshape(b,c,h*w) + q = q.permute(0,2,1) # b,hw,c + k = k.reshape(b,c,h*w) # b,c,hw + w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c)**(-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b,c,h*w) + w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v,w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b,c,h,w) + + h_ = self.proj_out(h_) + + return x+h_ + + +def make_attn(in_channels, attn_type="vanilla"): + assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown' + print(f"making attention of type '{attn_type}' with {in_channels} in_channels") + if attn_type == "vanilla": + return AttnBlock(in_channels) + elif attn_type == "none": + return nn.Identity(in_channels) + else: + return LinAttnBlock(in_channels) + + +class Encoder(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla", + **ignore_kwargs): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + # downsampling + self.conv_in = torch.nn.Conv2d(in_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + 2*z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + # timestep embedding + temb = None + + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False, + attn_type="vanilla", **ignorekwargs): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.give_pre_end = give_pre_end + self.tanh_out = tanh_out + + # compute in_ch_mult, block_in and curr_res at lowest res + in_ch_mult = (1,)+tuple(ch_mult) + block_in = ch*ch_mult[self.num_resolutions-1] + curr_res = resolution // 2**(self.num_resolutions-1) + self.z_shape = (1,z_channels,curr_res,curr_res) + print("Working with z of shape {} = {} dimensions.".format( + self.z_shape, np.prod(self.z_shape))) + + # z to block_in + self.conv_in = torch.nn.Conv2d(z_channels, + block_in, + kernel_size=3, + stride=1, + padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, z): + #assert z.shape[1:] == self.z_shape[1:] + self.last_z_shape = z.shape + + # timestep embedding + temb = None + + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks+1): + h = self.up[i_level].block[i_block](h, temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + if self.tanh_out: + h = torch.tanh(h) + return h + + +class FrozenAutoencoderKL(nn.Module): + def __init__(self, ddconfig, embed_dim, pretrained_path, scale_factor=0.18215): + super().__init__() + print(f'Create autoencoder with scale_factor={scale_factor}') + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + assert ddconfig["double_z"] + self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + self.embed_dim = embed_dim + self.scale_factor = scale_factor + m, u = self.load_state_dict(torch.load(pretrained_path, map_location='cpu')) + assert len(m) == 0 and len(u) == 0 + self.eval() + self.requires_grad_(False) + + def encode_moments(self, x): + h = self.encoder(x) + moments = self.quant_conv(h) + return moments + + def sample(self, moments): + mean, logvar = torch.chunk(moments, 2, dim=1) + logvar = torch.clamp(logvar, -30.0, 20.0) + std = torch.exp(0.5 * logvar) + z = mean + std * torch.randn_like(mean) + z = self.scale_factor * z + return z + + def encode(self, x): + moments = self.encode_moments(x) + z = self.sample(moments) + return z + + def decode(self, z): + z = (1. / self.scale_factor) * z + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward(self, inputs, fn): + if fn == 'encode_moments': + return self.encode_moments(inputs) + elif fn == 'encode': + return self.encode(inputs) + elif fn == 'decode': + return self.decode(inputs) + else: + raise NotImplementedError + + +def get_model(pretrained_path, scale_factor=0.18215): + ddconfig = dict( + double_z=True, + z_channels=4, + resolution=256, + in_channels=3, + out_ch=3, + ch=128, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_resolutions=[], + dropout=0.0 + ) + return FrozenAutoencoderKL(ddconfig, 4, pretrained_path, scale_factor) + + +def main(): + import torchvision.transforms as transforms + from torchvision.utils import save_image + import os + from PIL import Image + + model = get_model('/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/save_auto/checkpoint/autoencoder_kl.pth') + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + model = model.to(device) + + scale_factor = 0.18215 + T = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(256), transforms.ToTensor()]) + path = '/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/auto_data/test' + fnames = os.listdir(path) + for fname in fnames: + p = os.path.join(path, fname) + img = Image.open(p) + img = T(img) + img = img * 2. - 1 + img = img[None, ...] + img = img.to(device) + + # with torch.cuda.amp.autocast(): + # moments = model.encode_moments(img) + # mean, logvar = torch.chunk(moments, 2, dim=1) + # logvar = torch.clamp(logvar, -30.0, 20.0) + # std = torch.exp(0.5 * logvar) + # zs = [(mean + std * torch.randn_like(mean)) * scale_factor for _ in range(4)] + # recons = [model.decode(z) for z in zs] + + with torch.cuda.amp.autocast(): + print('test encode & decode') + recons = [model.decode(model.encode(img)) for _ in range(4)] + + out = torch.cat([img, *recons], dim=0) + out = (out + 1) * 0.5 + save_image(out, '/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/save_auto/img/' + f'recons_{fname}') + + +from torch.utils.data import Dataset, DataLoader +from torchvision.transforms import ToTensor, Compose, CenterCrop, Resize, RandomCrop, Lambda +from torch.optim import Adam +class Dataset_self(Dataset): + def __init__(self,img_root, preprocess): + self.img_root = img_root + self.img_process = preprocess + self.img = [] + for name_img in os.listdir(self.img_root): + self.img.append(self.img_root + '/' + name_img) + + def __len__(self): + return len(self.img) + + def __getitem__(self, idx): + img_path = self.img[idx] + image = Image.open(img_path).convert('RGB') + image = self.img_process(image) + return image + +def train(): + model = get_model('assets/stable-diffusion/autoencoder_kl.pth') + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + model = model.to(device) + scale_factor = 0.18215 + path_test = ['/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/auto_data/test/389.png', '/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/auto_data/test/390.png'] + T = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(256), transforms.ToTensor(), Lambda(lambda t: (t * 2) - 1)]) + dataset = Dataset_self(img_root= '/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/auto_data/train/good', preprocess=T) + dataloader = DataLoader(dataset, batch_size=4, shuffle=True) + criterion = nn.MSELoss() + optimizer = Adam(model.parameters(), lr=1e-5) + epochs = 100000 + for epoch in range(epochs): + step = 0 + for batch in dataloader: + step += 1 + optimizer.zero_grad() + batch = batch.to(device) + output_model = model.decode(model.encode(batch)) + loss = criterion(output_model, batch) + loss.backward() + optimizer.step() + print(epoch, loss.item()) + torch.save({ + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + }, '/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/save_auto/checkpoint/model_hazelnut_last.pth') + for i in range(2): + img_path = path_test[i] + img = Image.open(img_path) + img = T(img) + img = img[None, ...] + img = img.to(device) + out = model.decode(model.encode(img)) + out = (out + 1) * 0.5 + save_image(out, '/data/CASIA/Projects/sqf/Diffsion_Model/self_operate/save_auto/img/' + f'recons_{epoch}_{i}.png') + + +if __name__ == "__main__": + # train() + main() \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__init__.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dcc5619538c0f7c782508bdbd9587259d805e0d9 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__init__.py @@ -0,0 +1 @@ +from .clip import * diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2439b07e3b80e450232a40b4da28dd55d7157f71 Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/clip.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/clip.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f10c055bd05ded21b73e181f110cfe65c1d4d6d0 Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/clip.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/model.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51882398d667e0032b33fa3692e9906d3781e535 Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/model.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/simple_tokenizer.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/simple_tokenizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f37f1522556632135ce2180f2320541f9db3001c Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/__pycache__/simple_tokenizer.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/bpe_simple_vocab_16e6.txt.gz b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/bpe_simple_vocab_16e6.txt.gz new file mode 100644 index 0000000000000000000000000000000000000000..36a15856e00a06a9fbed8cdd34d2393fea4a3113 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/bpe_simple_vocab_16e6.txt.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a +size 1356917 diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/clip.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..257511e1d40c120e0d64a0f1562d44b2b8a40a17 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/clip.py @@ -0,0 +1,237 @@ +import hashlib +import os +import urllib +import warnings +from typing import Any, Union, List +from pkg_resources import packaging + +import torch +from PIL import Image +from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize +from tqdm import tqdm + +from .model import build_model +from .simple_tokenizer import SimpleTokenizer as _Tokenizer + +try: + from torchvision.transforms import InterpolationMode + BICUBIC = InterpolationMode.BICUBIC +except ImportError: + BICUBIC = Image.BICUBIC + + +if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"): + warnings.warn("PyTorch version 1.7.1 or higher is recommended") + + +__all__ = ["available_models", "load", "tokenize"] +_tokenizer = _Tokenizer() + +_MODELS = { + "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt", + "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt", + "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt", + "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt", + "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt", + "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt", + "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt", + "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt", + "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt", +} + + +def _download(url: str, root: str): + os.makedirs(root, exist_ok=True) + filename = os.path.basename(url) + + expected_sha256 = url.split("/")[-2] + download_target = os.path.join(root, filename) + + if os.path.exists(download_target) and not os.path.isfile(download_target): + raise RuntimeError(f"{download_target} exists and is not a regular file") + + if os.path.isfile(download_target): + if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256: + return download_target + else: + warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file") + + with urllib.request.urlopen(url) as source, open(download_target, "wb") as output: + with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop: + while True: + buffer = source.read(8192) + if not buffer: + break + + output.write(buffer) + loop.update(len(buffer)) + + if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256: + raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match") + + return download_target + + +def _convert_image_to_rgb(image): + return image.convert("RGB") + + +def _transform(n_px): + return Compose([ + Resize(n_px, interpolation=BICUBIC), + CenterCrop(n_px), + _convert_image_to_rgb, + ToTensor(), + Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), + ]) + + +def available_models() -> List[str]: + """Returns the names of available CLIP models""" + return list(_MODELS.keys()) + + +def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None): + """Load a CLIP model + + Parameters + ---------- + name : str + A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict + + device : Union[str, torch.device] + The device to put the loaded model + + jit : bool + Whether to load the optimized JIT model or more hackable non-JIT model (default). + + download_root: str + path to download the model files; by default, it uses "~/.cache/clip" + + Returns + ------- + model : torch.nn.Module + The CLIP model + + preprocess : Callable[[PIL.Image], torch.Tensor] + A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input + """ + if name in _MODELS: + model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip")) + elif os.path.isfile(name): + model_path = name + else: + raise RuntimeError(f"Model {name} not found; available models = {available_models()}") + + with open(model_path, 'rb') as opened_file: + try: + # loading JIT archive + model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval() + state_dict = None + except RuntimeError: + # loading saved state dict + if jit: + warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead") + jit = False + state_dict = torch.load(opened_file, map_location="cpu") + + if not jit: + model = build_model(state_dict or model.state_dict()).to(device) + if str(device) == "cpu": + model.float() + return model, _transform(model.visual.input_resolution) + + # patch the device names + device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[]) + device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1] + + def patch_device(module): + try: + graphs = [module.graph] if hasattr(module, "graph") else [] + except RuntimeError: + graphs = [] + + if hasattr(module, "forward1"): + graphs.append(module.forward1.graph) + + for graph in graphs: + for node in graph.findAllNodes("prim::Constant"): + if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"): + node.copyAttributes(device_node) + + model.apply(patch_device) + patch_device(model.encode_image) + patch_device(model.encode_text) + + # patch dtype to float32 on CPU + if str(device) == "cpu": + float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[]) + float_input = list(float_holder.graph.findNode("aten::to").inputs())[1] + float_node = float_input.node() + + def patch_float(module): + try: + graphs = [module.graph] if hasattr(module, "graph") else [] + except RuntimeError: + graphs = [] + + if hasattr(module, "forward1"): + graphs.append(module.forward1.graph) + + for graph in graphs: + for node in graph.findAllNodes("aten::to"): + inputs = list(node.inputs()) + for i in [1, 2]: # dtype can be the second or third argument to aten::to() + if inputs[i].node()["value"] == 5: + inputs[i].node().copyAttributes(float_node) + + model.apply(patch_float) + patch_float(model.encode_image) + patch_float(model.encode_text) + + model.float() + + return model, _transform(model.input_resolution.item()) + + +def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]: + """ + Returns the tokenized representation of given input string(s) + + Parameters + ---------- + texts : Union[str, List[str]] + An input string or a list of input strings to tokenize + + context_length : int + The context length to use; all CLIP models use 77 as the context length + + truncate: bool + Whether to truncate the text in case its encoding is longer than the context length + + Returns + ------- + A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]. + We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long. + """ + if isinstance(texts, str): + texts = [texts] + + sot_token = _tokenizer.encoder["<|startoftext|>"] + eot_token = _tokenizer.encoder["<|endoftext|>"] + all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] + if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"): + result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + else: + result = torch.zeros(len(all_tokens), context_length, dtype=torch.int) + + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + if truncate: + tokens = tokens[:context_length] + tokens[-1] = eot_token + else: + raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}") + result[i, :len(tokens)] = torch.tensor(tokens) + + return result diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/model.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/model.py new file mode 100644 index 0000000000000000000000000000000000000000..48d293d100254f885eb2deeae4ce5ce2c952a7e1 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/model.py @@ -0,0 +1,434 @@ +from collections import OrderedDict +from typing import Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1): + super().__init__() + + # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1 + self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.relu1 = nn.ReLU(inplace=True) + + self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.relu2 = nn.ReLU(inplace=True) + + self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity() + + self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.relu3 = nn.ReLU(inplace=True) + + self.downsample = None + self.stride = stride + + if stride > 1 or inplanes != planes * Bottleneck.expansion: + # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1 + self.downsample = nn.Sequential(OrderedDict([ + ("-1", nn.AvgPool2d(stride)), + ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)), + ("1", nn.BatchNorm2d(planes * self.expansion)) + ])) + + def forward(self, x: torch.Tensor): + identity = x + + out = self.relu1(self.bn1(self.conv1(x))) + out = self.relu2(self.bn2(self.conv2(out))) + out = self.avgpool(out) + out = self.bn3(self.conv3(out)) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu3(out) + return out + + +class AttentionPool2d(nn.Module): + def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None): + super().__init__() + self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5) + self.k_proj = nn.Linear(embed_dim, embed_dim) + self.q_proj = nn.Linear(embed_dim, embed_dim) + self.v_proj = nn.Linear(embed_dim, embed_dim) + self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) + self.num_heads = num_heads + + def forward(self, x): + x = x.flatten(start_dim=2).permute(2, 0, 1) # NCHW -> (HW)NC + x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC + x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC + x, _ = F.multi_head_attention_forward( + query=x[:1], key=x, value=x, + embed_dim_to_check=x.shape[-1], + num_heads=self.num_heads, + q_proj_weight=self.q_proj.weight, + k_proj_weight=self.k_proj.weight, + v_proj_weight=self.v_proj.weight, + in_proj_weight=None, + in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]), + bias_k=None, + bias_v=None, + add_zero_attn=False, + dropout_p=0, + out_proj_weight=self.c_proj.weight, + out_proj_bias=self.c_proj.bias, + use_separate_proj_weight=True, + training=self.training, + need_weights=False + ) + return x.squeeze(0) + + +class ModifiedResNet(nn.Module): + """ + A ResNet class that is similar to torchvision's but contains the following changes: + - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool. + - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1 + - The final pooling layer is a QKV attention instead of an average pool + """ + + def __init__(self, layers, output_dim, heads, input_resolution=224, width=64): + super().__init__() + self.output_dim = output_dim + self.input_resolution = input_resolution + + # the 3-layer stem + self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(width // 2) + self.relu1 = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(width // 2) + self.relu2 = nn.ReLU(inplace=True) + self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False) + self.bn3 = nn.BatchNorm2d(width) + self.relu3 = nn.ReLU(inplace=True) + self.avgpool = nn.AvgPool2d(2) + + # residual layers + self._inplanes = width # this is a *mutable* variable used during construction + self.layer1 = self._make_layer(width, layers[0]) + self.layer2 = self._make_layer(width * 2, layers[1], stride=2) + self.layer3 = self._make_layer(width * 4, layers[2], stride=2) + self.layer4 = self._make_layer(width * 8, layers[3], stride=2) + + embed_dim = width * 32 # the ResNet feature dimension + self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim) + + def _make_layer(self, planes, blocks, stride=1): + layers = [Bottleneck(self._inplanes, planes, stride)] + + self._inplanes = planes * Bottleneck.expansion + for _ in range(1, blocks): + layers.append(Bottleneck(self._inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + def stem(x): + x = self.relu1(self.bn1(self.conv1(x))) + x = self.relu2(self.bn2(self.conv2(x))) + x = self.relu3(self.bn3(self.conv3(x))) + x = self.avgpool(x) + return x + + x = x.type(self.conv1.weight.dtype) + x = stem(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + x = self.attnpool(x) + + return x + + +class LayerNorm(nn.LayerNorm): + """Subclass torch's LayerNorm to handle fp16.""" + + def forward(self, x: torch.Tensor): + orig_type = x.dtype + ret = super().forward(x.type(torch.float32)) + return ret.type(orig_type) + + +class QuickGELU(nn.Module): + def forward(self, x: torch.Tensor): + return x * torch.sigmoid(1.702 * x) + + +class ResidualAttentionBlock(nn.Module): + def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): + super().__init__() + + self.attn = nn.MultiheadAttention(d_model, n_head) + self.ln_1 = LayerNorm(d_model) + self.mlp = nn.Sequential(OrderedDict([ + ("c_fc", nn.Linear(d_model, d_model * 4)), + ("gelu", QuickGELU()), + ("c_proj", nn.Linear(d_model * 4, d_model)) + ])) + self.ln_2 = LayerNorm(d_model) + self.attn_mask = attn_mask + + def attention(self, x: torch.Tensor): + self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None + return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] + + def forward(self, x: torch.Tensor): + x = x + self.attention(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +class Transformer(nn.Module): + def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None): + super().__init__() + self.width = width + self.layers = layers + self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]) + + def forward(self, x: torch.Tensor): + return self.resblocks(x) + + +class VisionTransformer(nn.Module): + def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int): + super().__init__() + self.input_resolution = input_resolution + self.output_dim = output_dim + self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False) + + scale = width ** -0.5 + self.class_embedding = nn.Parameter(scale * torch.randn(width)) + self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width)) + self.ln_pre = LayerNorm(width) + + self.transformer = Transformer(width, layers, heads) + + self.ln_post = LayerNorm(width) + self.proj = nn.Parameter(scale * torch.randn(width, output_dim)) + + def forward(self, x: torch.Tensor): + x = self.conv1(x) # shape = [*, width, grid, grid] + x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] + x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] + x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width] + x = x + self.positional_embedding.to(x.dtype) + x = self.ln_pre(x) + + x = x.permute(1, 0, 2) # NLD -> LND + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + + x = self.ln_post(x[:, 0, :]) + + if self.proj is not None: + x = x @ self.proj + + return x + + +class CLIP(nn.Module): + def __init__(self, + embed_dim: int, + # vision + image_resolution: int, + vision_layers: Union[Tuple[int, int, int, int], int], + vision_width: int, + vision_patch_size: int, + # text + context_length: int, + vocab_size: int, + transformer_width: int, + transformer_heads: int, + transformer_layers: int + ): + super().__init__() + + self.context_length = context_length + + if isinstance(vision_layers, (tuple, list)): + vision_heads = vision_width * 32 // 64 + self.visual = ModifiedResNet( + layers=vision_layers, + output_dim=embed_dim, + heads=vision_heads, + input_resolution=image_resolution, + width=vision_width + ) + else: + vision_heads = vision_width // 64 + self.visual = VisionTransformer( + input_resolution=image_resolution, + patch_size=vision_patch_size, + width=vision_width, + layers=vision_layers, + heads=vision_heads, + output_dim=embed_dim + ) + + self.transformer = Transformer( + width=transformer_width, + layers=transformer_layers, + heads=transformer_heads, + attn_mask=self.build_attention_mask() + ) + + self.vocab_size = vocab_size + self.token_embedding = nn.Embedding(vocab_size, transformer_width) + self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width)) + self.ln_final = LayerNorm(transformer_width) + + self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim)) + self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) + + self.initialize_parameters() + + def initialize_parameters(self): + nn.init.normal_(self.token_embedding.weight, std=0.02) + nn.init.normal_(self.positional_embedding, std=0.01) + + if isinstance(self.visual, ModifiedResNet): + if self.visual.attnpool is not None: + std = self.visual.attnpool.c_proj.in_features ** -0.5 + nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std) + nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std) + + for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]: + for name, param in resnet_block.named_parameters(): + if name.endswith("bn3.weight"): + nn.init.zeros_(param) + + proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5) + attn_std = self.transformer.width ** -0.5 + fc_std = (2 * self.transformer.width) ** -0.5 + for block in self.transformer.resblocks: + nn.init.normal_(block.attn.in_proj_weight, std=attn_std) + nn.init.normal_(block.attn.out_proj.weight, std=proj_std) + nn.init.normal_(block.mlp.c_fc.weight, std=fc_std) + nn.init.normal_(block.mlp.c_proj.weight, std=proj_std) + + if self.text_projection is not None: + nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5) + + def build_attention_mask(self): + # lazily create causal attention mask, with full attention between the vision tokens + # pytorch uses additive attention mask; fill with -inf + mask = torch.empty(self.context_length, self.context_length) + mask.fill_(float("-inf")) + mask.triu_(1) # zero out the lower diagonal + return mask + + @property + def dtype(self): + return self.visual.conv1.weight.dtype + + def encode_image(self, image): + return self.visual(image.type(self.dtype)) + + def encode_text(self, text): + x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model] + + x = x + self.positional_embedding.type(self.dtype) + x = x.permute(1, 0, 2) # NLD -> LND + x = self.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + x = self.ln_final(x).type(self.dtype) + + x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection + + return x + + def forward(self, image, text): + image_features = self.encode_image(image) + text_features = self.encode_text(text) + + # normalized features + image_features = image_features / image_features.norm(dim=1, keepdim=True) + text_features = text_features / text_features.norm(dim=1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_image = logit_scale * image_features @ text_features.t() + logits_per_text = logits_per_image.t() + + # shape = [global_batch_size, global_batch_size] + return logits_per_image, logits_per_text + + +def convert_weights(model: nn.Module): + """Convert applicable model parameters to fp16""" + + def _convert_weights_to_fp16(l): + if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)): + l.weight.data = l.weight.data.half() + if l.bias is not None: + l.bias.data = l.bias.data.half() + + if isinstance(l, nn.MultiheadAttention): + for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]: + tensor = getattr(l, attr) + if tensor is not None: + tensor.data = tensor.data.half() + + for name in ["text_projection", "proj"]: + if hasattr(l, name): + attr = getattr(l, name) + if attr is not None: + attr.data = attr.data.half() + + model.apply(_convert_weights_to_fp16) + + +def build_model(state_dict: dict): + vit = "visual.proj" in state_dict + + if vit: + vision_width = state_dict["visual.conv1.weight"].shape[0] + vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")]) + vision_patch_size = state_dict["visual.conv1.weight"].shape[-1] + grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5) + image_resolution = vision_patch_size * grid_size + else: + counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]] + vision_layers = tuple(counts) + vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0] + output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5) + vision_patch_size = None + assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0] + image_resolution = output_width * 32 + + embed_dim = state_dict["text_projection"].shape[1] + context_length = state_dict["positional_embedding"].shape[0] + vocab_size = state_dict["token_embedding.weight"].shape[0] + transformer_width = state_dict["ln_final.weight"].shape[0] + transformer_heads = transformer_width // 64 + transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks"))) + + model = CLIP( + embed_dim, + image_resolution, vision_layers, vision_width, vision_patch_size, + context_length, vocab_size, transformer_width, transformer_heads, transformer_layers + ) + + for key in ["input_resolution", "context_length", "vocab_size"]: + if key in state_dict: + del state_dict[key] + + convert_weights(model) + model.load_state_dict(state_dict) + return model.eval() diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/simple_tokenizer.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/simple_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..0a66286b7d5019c6e221932a813768038f839c91 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/clip/simple_tokenizer.py @@ -0,0 +1,132 @@ +import gzip +import html +import os +from functools import lru_cache + +import ftfy +import regex as re + + +@lru_cache() +def default_bpe(): + return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz") + + +@lru_cache() +def bytes_to_unicode(): + """ + Returns list of utf-8 byte and a corresponding list of unicode strings. + The reversible bpe codes work on unicode strings. + This means you need a large # of unicode characters in your vocab if you want to avoid UNKs. + When you're at something like a 10B token dataset you end up needing around 5K for decent coverage. + This is a signficant percentage of your normal, say, 32K bpe vocab. + To avoid that, we want lookup tables between utf-8 bytes and unicode strings. + And avoids mapping to whitespace/control characters the bpe code barfs on. + """ + bs = list(range(ord("!"), ord("~")+1))+list(range(ord("ยก"), ord("ยฌ")+1))+list(range(ord("ยฎ"), ord("รฟ")+1)) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8+n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + """Return set of symbol pairs in a word. + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +class SimpleTokenizer(object): + def __init__(self, bpe_path: str = default_bpe()): + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + merges = gzip.open(bpe_path).read().decode("utf-8").split('\n') + merges = merges[1:49152-256-2+1] + merges = [tuple(merge.split()) for merge in merges] + vocab = list(bytes_to_unicode().values()) + vocab = vocab + [v+'' for v in vocab] + for merge in merges: + vocab.append(''.join(merge)) + vocab.extend(['<|startoftext|>', '<|endoftext|>']) + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.decoder = {v: k for k, v in self.encoder.items()} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'} + self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE) + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + ( token[-1] + '',) + pairs = get_pairs(word) + + if not pairs: + return token+'' + + while True: + bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf'))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + new_word.extend(word[i:j]) + i = j + except: + new_word.extend(word[i:]) + break + + if word[i] == first and i < len(word)-1 and word[i+1] == second: + new_word.append(first+second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = ' '.join(word) + self.cache[token] = word + return word + + def encode(self, text): + bpe_tokens = [] + text = whitespace_clean(basic_clean(text)).lower() + for token in re.findall(self.pat, text): + token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) + bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' ')) + return bpe_tokens + + def decode(self, tokens): + text = ''.join([self.decoder[token] for token in tokens]) + text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('', ' ') + return text diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__init__.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c536a98da92c4d051458803737661e5ecf974c2 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__init__.py @@ -0,0 +1,46 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +from . import gaussian_diffusion as gd +from .respace import SpacedDiffusion, space_timesteps + + +def create_diffusion( + timestep_respacing, + noise_schedule="linear", + use_kl=False, + sigma_small=False, + predict_xstart=False, + learn_sigma=True, + rescale_learned_sigmas=False, + diffusion_steps=1000 +): + betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) + if use_kl: + loss_type = gd.LossType.RESCALED_KL + elif rescale_learned_sigmas: + loss_type = gd.LossType.RESCALED_MSE + else: + loss_type = gd.LossType.MSE + if timestep_respacing is None or timestep_respacing == "": + timestep_respacing = [diffusion_steps] + return SpacedDiffusion( + use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), + betas=betas, + model_mean_type=( + gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X + ), + model_var_type=( + ( + gd.ModelVarType.FIXED_LARGE + if not sigma_small + else gd.ModelVarType.FIXED_SMALL + ) + if not learn_sigma + else gd.ModelVarType.LEARNED_RANGE + ), + loss_type=loss_type + # rescale_timesteps=rescale_timesteps, + ) diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..839c0d82cab864b22d57f8f8c0d6237fd8c74bc7 Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/diffusion_utils.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/diffusion_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aabd41a57cfd5d5495b7bf154c6cac9896870765 Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/diffusion_utils.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/gaussian_diffusion.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/gaussian_diffusion.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65cf483a6652f7da8455cd23cb61b527937fe94b Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/gaussian_diffusion.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/respace.cpython-310.pyc b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/respace.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb2a5513541e95caa2b87cfc329d363c72e59fe3 Binary files /dev/null and b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/__pycache__/respace.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/diffusion_utils.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/diffusion_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e493a6a3ecb91e553a53cc7eadee5cc0d1753060 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/diffusion_utils.py @@ -0,0 +1,88 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +import torch as th +import numpy as np + + +def normal_kl(mean1, logvar1, mean2, logvar2): + """ + Compute the KL divergence between two gaussians. + Shapes are automatically broadcasted, so batches can be compared to + scalars, among other use cases. + """ + tensor = None + for obj in (mean1, logvar1, mean2, logvar2): + if isinstance(obj, th.Tensor): + tensor = obj + break + assert tensor is not None, "at least one argument must be a Tensor" + + # Force variances to be Tensors. Broadcasting helps convert scalars to + # Tensors, but it does not work for th.exp(). + logvar1, logvar2 = [ + x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) + for x in (logvar1, logvar2) + ] + + return 0.5 * ( + -1.0 + + logvar2 + - logvar1 + + th.exp(logvar1 - logvar2) + + ((mean1 - mean2) ** 2) * th.exp(-logvar2) + ) + + +def approx_standard_normal_cdf(x): + """ + A fast approximation of the cumulative distribution function of the + standard normal. + """ + return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) + + +def continuous_gaussian_log_likelihood(x, *, means, log_scales): + """ + Compute the log-likelihood of a continuous Gaussian distribution. + :param x: the targets + :param means: the Gaussian mean Tensor. + :param log_scales: the Gaussian log stddev Tensor. + :return: a tensor like x of log probabilities (in nats). + """ + centered_x = x - means + inv_stdv = th.exp(-log_scales) + normalized_x = centered_x * inv_stdv + log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x) + return log_probs + + +def discretized_gaussian_log_likelihood(x, *, means, log_scales): + """ + Compute the log-likelihood of a Gaussian distribution discretizing to a + given image. + :param x: the target images. It is assumed that this was uint8 values, + rescaled to the range [-1, 1]. + :param means: the Gaussian mean Tensor. + :param log_scales: the Gaussian log stddev Tensor. + :return: a tensor like x of log probabilities (in nats). + """ + assert x.shape == means.shape == log_scales.shape + centered_x = x - means + inv_stdv = th.exp(-log_scales) + plus_in = inv_stdv * (centered_x + 1.0 / 255.0) + cdf_plus = approx_standard_normal_cdf(plus_in) + min_in = inv_stdv * (centered_x - 1.0 / 255.0) + cdf_min = approx_standard_normal_cdf(min_in) + log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) + log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) + cdf_delta = cdf_plus - cdf_min + log_probs = th.where( + x < -0.999, + log_cdf_plus, + th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), + ) + assert log_probs.shape == x.shape + return log_probs diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/gaussian_diffusion.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/gaussian_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..1279992c64d6ad7b007bcd5959496a8d3b842300 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/gaussian_diffusion.py @@ -0,0 +1,903 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + + +import math + +import numpy as np +import torch as th +import enum + +from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl + + +def mean_flat(tensor): + """ + Take the mean over all non-batch dimensions. + """ + return tensor.mean(dim=list(range(1, len(tensor.shape)))) + + +class ModelMeanType(enum.Enum): + """ + Which type of output the model predicts. + """ + + PREVIOUS_X = enum.auto() # the model predicts x_{t-1} + START_X = enum.auto() # the model predicts x_0 + EPSILON = enum.auto() # the model predicts epsilon + + +class ModelVarType(enum.Enum): + """ + What is used as the model's output variance. + The LEARNED_RANGE option has been added to allow the model to predict + values between FIXED_SMALL and FIXED_LARGE, making its job easier. + """ + + LEARNED = enum.auto() + FIXED_SMALL = enum.auto() + FIXED_LARGE = enum.auto() + LEARNED_RANGE = enum.auto() + + +class LossType(enum.Enum): + MSE = enum.auto() # use raw MSE loss (and KL when learning variances) + RESCALED_MSE = ( + enum.auto() + ) # use raw MSE loss (with RESCALED_KL when learning variances) + KL = enum.auto() # use the variational lower-bound + RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB + + def is_vb(self): + return self == LossType.KL or self == LossType.RESCALED_KL + + +def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + warmup_time = int(num_diffusion_timesteps * warmup_frac) + betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) + return betas + + +def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): + """ + This is the deprecated API for creating beta schedules. + See get_named_beta_schedule() for the new library of schedules. + """ + if beta_schedule == "quad": + betas = ( + np.linspace( + beta_start ** 0.5, + beta_end ** 0.5, + num_diffusion_timesteps, + dtype=np.float64, + ) + ** 2 + ) + elif beta_schedule == "linear": + betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "warmup10": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) + elif beta_schedule == "warmup50": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) + elif beta_schedule == "const": + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1 + betas = 1.0 / np.linspace( + num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 + ) + else: + raise NotImplementedError(beta_schedule) + assert betas.shape == (num_diffusion_timesteps,) + return betas + + +def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): + """ + Get a pre-defined beta schedule for the given name. + The beta schedule library consists of beta schedules which remain similar + in the limit of num_diffusion_timesteps. + Beta schedules may be added, but should not be removed or changed once + they are committed to maintain backwards compatibility. + """ + if schedule_name == "linear": + # Linear schedule from Ho et al, extended to work for any number of + # diffusion steps. + scale = 1000 / num_diffusion_timesteps + return get_beta_schedule( + "linear", + beta_start=scale * 0.0001, + beta_end=scale * 0.02, + num_diffusion_timesteps=num_diffusion_timesteps, + ) + elif schedule_name == "squaredcos_cap_v2": + return betas_for_alpha_bar( + num_diffusion_timesteps, + lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, + ) + else: + raise NotImplementedError(f"unknown beta schedule: {schedule_name}") + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, + which defines the cumulative product of (1-beta) over time from t = [0,1]. + :param num_diffusion_timesteps: the number of betas to produce. + :param alpha_bar: a lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + :param max_beta: the maximum beta to use; use values lower than 1 to + prevent singularities. + """ + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +class GaussianDiffusion: + """ + Utilities for training and sampling diffusion models. + Original ported from this codebase: + https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 + :param betas: a 1-D numpy array of betas for each diffusion timestep, + starting at T and going to 1. + """ + + def __init__( + self, + *, + betas, + model_mean_type, + model_var_type, + loss_type + ): + + self.model_mean_type = model_mean_type + self.model_var_type = model_var_type + self.loss_type = loss_type + + # Use float64 for accuracy. + betas = np.array(betas, dtype=np.float64) + self.betas = betas + assert len(betas.shape) == 1, "betas must be 1-D" + assert (betas > 0).all() and (betas <= 1).all() + + self.num_timesteps = int(betas.shape[0]) + + alphas = 1.0 - betas + self.alphas_cumprod = np.cumprod(alphas, axis=0) + self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) + self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) + assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) + self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) + self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) + self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) + self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + self.posterior_variance = ( + betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain + self.posterior_log_variance_clipped = np.log( + np.append(self.posterior_variance[1], self.posterior_variance[1:]) + ) if len(self.posterior_variance) > 1 else np.array([]) + + self.posterior_mean_coef1 = ( + betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + self.posterior_mean_coef2 = ( + (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod) + ) + + # self.defect_w = th.nn.Parameter(torch.FloatTensor(1), requires_grad=True) + + + def q_mean_variance(self, x_start, t): + """ + Get the distribution q(x_t | x_0). + :param x_start: the [N x C x ...] tensor of noiseless inputs. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ + mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) + log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) + return mean, variance, log_variance + + def q_sample(self, x_start, t, noise=None): + """ + Diffuse the data for a given number of diffusion steps. + In other words, sample from q(x_t | x_0). + :param x_start: the initial data batch. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :param noise: if specified, the split-out normal noise. + :return: A noisy version of x_start. + """ + if noise is None: + noise = th.randn_like(x_start) + assert noise.shape == x_start.shape + return ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + def q_posterior_mean_variance(self, x_start, x_t, t): + """ + Compute the mean and variance of the diffusion posterior: + q(x_{t-1} | x_t, x_0) + """ + assert x_start.shape == x_t.shape + posterior_mean = ( + _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x_t.shape + ) + assert ( + posterior_mean.shape[0] + == posterior_variance.shape[0] + == posterior_log_variance_clipped.shape[0] + == x_start.shape[0] + ) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): + """ + Apply the model to get p(x_{t-1} | x_t), as well as a prediction of + the initial x, x_0. + :param model: the model, which takes a signal and a batch of timesteps + as input. + :param x: the [N x C x ...] tensor at time t. + :param t: a 1-D Tensor of timesteps. + :param clip_denoised: if True, clip the denoised signal into [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. Applies before + clip_denoised. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict with the following keys: + - 'mean': the model mean output. + - 'variance': the model variance output. + - 'log_variance': the log of 'variance'. + - 'pred_xstart': the prediction for x_0. + """ + if model_kwargs is None: + model_kwargs = {} + new_mask=None + B, C = x.shape[:2] + assert t.shape == (B,) + if model_kwargs == {}: + model_output = model(x, t, **model_kwargs) + else: + model_output, new_mask, _ = model(x, t, **model_kwargs) + if isinstance(model_output, tuple): + model_output, extra = model_output + else: + extra = None + + if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: + assert model_output.shape == (B, C * 2, *x.shape[2:]) + model_output, model_var_values = th.split(model_output, C, dim=1) + min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape) + max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) + # The model_var_values is [-1, 1] for [min_var, max_var]. + frac = (model_var_values + 1) / 2 + model_log_variance = frac * max_log + (1 - frac) * min_log + model_variance = th.exp(model_log_variance) + else: + model_variance, model_log_variance = { + # for fixedlarge, we set the initial (log-)variance like so + # to get a better decoder log likelihood. + ModelVarType.FIXED_LARGE: ( + np.append(self.posterior_variance[1], self.betas[1:]), + np.log(np.append(self.posterior_variance[1], self.betas[1:])), + ), + ModelVarType.FIXED_SMALL: ( + self.posterior_variance, + self.posterior_log_variance_clipped, + ), + }[self.model_var_type] + model_variance = _extract_into_tensor(model_variance, t, x.shape) + model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) + + def process_xstart(x): + if denoised_fn is not None: + x = denoised_fn(x) + if clip_denoised: + return x.clamp(-1, 1) + return x + + if self.model_mean_type == ModelMeanType.START_X: + pred_xstart = process_xstart(model_output) + else: + pred_xstart = process_xstart( + self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) + ) + model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t) + + assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape + if new_mask is None: + new_mask = pred_xstart + return { + "mean": model_mean, + "variance": model_variance, + "log_variance": model_log_variance, + "pred_xstart": pred_xstart, + "extra": extra, + "mask":new_mask, + } + + def _predict_xstart_from_eps(self, x_t, t, eps): + assert x_t.shape == eps.shape + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t + - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps + ) + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute the mean for the previous step, given a function cond_fn that + computes the gradient of a conditional log probability with respect to + x. In particular, cond_fn computes grad(log(p(y|x))), and we want to + condition on y. + This uses the conditioning strategy from Sohl-Dickstein et al. (2015). + """ + gradient = cond_fn(x, t, **model_kwargs) + new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() + return new_mean + + + def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute what the p_mean_variance output would have been, should the + model's score function be conditioned by cond_fn. + See condition_mean() for details on cond_fn. + Unlike condition_mean(), this instead uses the conditioning strategy + from Song et al (2020). + """ + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + + eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) + eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs) + + out = p_mean_var.copy() + out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) + out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t) + return out + + def p_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + ): + """ + Sample x_{t-1} from the model at the given timestep. + :param model: the model to sample from. + :param x: the current tensor at x_{t-1}. + :param t: the value of t, starting at 0 for the first diffusion step. + :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - 'sample': a random sample from the model. + - 'pred_xstart': a prediction of x_0. + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + noise = th.randn_like(x) + nonzero_mask = ( + (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + if cond_fn is not None: + out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs) + sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"], "mask":out["mask"]} + + def p_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + """ + Generate samples from the model. + :param model: the model module. + :param shape: the shape of the samples, (N, C, H, W). + :param noise: if specified, the noise from the encoder to sample. + Should be of the same shape as `shape`. + :param clip_denoised: if True, clip x_start predictions to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param device: if specified, the device to create the samples on. + If not specified, use a model parameter's device. + :param progress: if True, show a tqdm progress bar. + :return: a non-differentiable batch of samples. + """ + final = None + mask = 0 + num = 0 + for sample in self.p_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + ): + + num += 1 + final = sample + if num > 45: + mask += sample["mask"] + return final["sample"], mask / 5 + + def p_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + """ + Generate samples from the model and yield intermediate samples from + each timestep of diffusion. + Arguments are the same as p_sample_loop(). + Returns a generator over dicts, where each dict is the return value of + p_sample(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + if noise is not None: + img = noise + else: + img = th.randn(*shape, device=device) + indices = list(range(self.num_timesteps))[::-1] + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + with th.no_grad(): + out = self.p_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + ) + yield out + img = out["sample"] + + def ddim_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + eta=0.0, + ): + """ + Sample x_{t-1} from the model using DDIM. + Same usage as p_sample(). + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + if cond_fn is not None: + out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) + + # Usually our model outputs epsilon, but we re-derive it + # in case we used x_start or x_prev prediction. + eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) + + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) + sigma = ( + eta + * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) + * th.sqrt(1 - alpha_bar / alpha_bar_prev) + ) + # Equation 12. + noise = th.randn_like(x) + mean_pred = ( + out["pred_xstart"] * th.sqrt(alpha_bar_prev) + + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps + ) + nonzero_mask = ( + (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + sample = mean_pred + nonzero_mask * sigma * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"]} + + def ddim_reverse_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + eta=0.0, + ): + """ + Sample x_{t+1} from the model using DDIM reverse ODE. + """ + assert eta == 0.0, "Reverse ODE only for deterministic path" + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + if cond_fn is not None: + out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) + # Usually our model outputs epsilon, but we re-derive it + # in case we used x_start or x_prev prediction. + eps = ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x + - out["pred_xstart"] + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) + alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) + + # Equation 12. reversed + mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps + + return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} + + def ddim_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + """ + Generate samples from the model using DDIM. + Same usage as p_sample_loop(). + """ + final = None + for sample in self.ddim_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + eta=eta, + ): + final = sample + return final["sample"] + + def ddim_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + """ + Use DDIM to sample from the model and yield intermediate samples from + each timestep of DDIM. + Same usage as p_sample_loop_progressive(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + if noise is not None: + img = noise + else: + img = th.randn(*shape, device=device) + indices = list(range(self.num_timesteps))[::-1] + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + with th.no_grad(): + out = self.ddim_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + eta=eta, + ) + yield out + img = out["sample"] + + def _vb_terms_bpd( + self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None + ): + """ + Get a term for the variational lower-bound. + The resulting units are bits (rather than nats, as one might expect). + This allows for comparison to other papers. + :return: a dict with the following keys: + - 'output': a shape [N] tensor of NLLs or KLs. + - 'pred_xstart': the x_0 predictions. + """ + true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + ) + out = self.p_mean_variance( + model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs + ) + kl = normal_kl( + true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] + ) + kl = mean_flat(kl) / np.log(2.0) + + decoder_nll = -discretized_gaussian_log_likelihood( + x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] + ) + assert decoder_nll.shape == x_start.shape + decoder_nll = mean_flat(decoder_nll) / np.log(2.0) + + # At the first timestep return the decoder NLL, + # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) + output = th.where((t == 0), decoder_nll, kl) + return {"output": output, "pred_xstart": out["pred_xstart"]} + + def training_losses(self, model, x_start, t, model_kwargs=None, noise=None, label_mask=None, mask_resize=None, mask_att=None): + """ + Compute training losses for a single timestep. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param t: a batch of timestep indices. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param noise: if specified, the specific Gaussian noise to try to remove. + :param label_mask: tiao jie sun shi + :return: a dict with the key "loss" containing a tensor of shape [N]. + Some mean or variance settings may also have other keys. + """ + + if model_kwargs is None: + model_kwargs = {} + if noise is None: + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start, t, noise=noise) + + terms = {} + + if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: + terms["loss"] = self._vb_terms_bpd( + model=model, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + model_kwargs=model_kwargs, + )["output"] + if self.loss_type == LossType.RESCALED_KL: + terms["loss"] *= self.num_timesteps + elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: + model_output, att_mask, att_loss = model(x_t, t, **model_kwargs) + + if self.model_var_type in [ + ModelVarType.LEARNED, + ModelVarType.LEARNED_RANGE, + ]: + B, C = x_t.shape[:2] + assert model_output.shape == (B, C * 2, *x_t.shape[2:]) + model_output, model_var_values = th.split(model_output, C, dim=1) + # Learn the variance using the variational bound, but don't let + # it affect our mean prediction. + frozen_out = th.cat([model_output.detach(), model_var_values], dim=1) + terms["vb"] = self._vb_terms_bpd( + model=lambda *args, r=frozen_out: r, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + )["output"] + if self.loss_type == LossType.RESCALED_MSE: + # Divide by 1000 for equivalence with initial implementation. + # Without a factor of 1/1000, the VB term hurts the MSE term. + terms["vb"] *= self.num_timesteps / 1000.0 + + target = { + ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + )[0], + ModelMeanType.START_X: x_start, + ModelMeanType.EPSILON: noise, + }[self.model_mean_type] + assert model_output.shape == target.shape == x_start.shape + + # print(att_loss.shape, mask_att.shape) + loss_defect = mean_flat(((th.mul(target, mask_resize) - th.mul(model_output, mask_resize)) ** 2)) + # loss_back = mean_flat(((th.mul(target, 1 - mask_resize) - th.mul(model_output, 1 - mask_resize)) ** 2)) + rat_loss = (th.sum(th.sum(th.mul(att_loss, 1 - mask_att), dim=-1), dim=-1)) / (th.sum(th.sum(th.mul(att_loss, mask_att), dim=-1), dim=-1) + 0.0001) + + rat_loss[rat_loss > 8] = 8 + rat_loss[rat_loss < 2] = 2 + + loss_att = rat_loss * loss_defect + + terms["mse"] = mean_flat((target - model_output) ** 2) + terms["mask"] = mean_flat((att_mask - label_mask) ** 2) + if "vb" in terms: + terms["loss"] = terms["mse"] + terms["vb"] + 0.2 * terms["mask"] + loss_att + else: + terms["loss"] = terms["mse"] + 0.2 * terms["mask"] + loss_att + else: + raise NotImplementedError(self.loss_type) + + return terms + + def _prior_bpd(self, x_start): + """ + Get the prior KL term for the variational lower-bound, measured in + bits-per-dim. + This term can't be optimized, as it only depends on the encoder. + :param x_start: the [N x C x ...] tensor of inputs. + :return: a batch of [N] KL values (in bits), one per batch element. + """ + batch_size = x_start.shape[0] + t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) + qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) + kl_prior = normal_kl( + mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 + ) + return mean_flat(kl_prior) / np.log(2.0) + + def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): + """ + Compute the entire variational lower-bound, measured in bits-per-dim, + as well as other related quantities. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param clip_denoised: if True, clip denoised samples. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - total_bpd: the total variational lower-bound, per batch element. + - prior_bpd: the prior term in the lower-bound. + - vb: an [N x T] tensor of terms in the lower-bound. + - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. + - mse: an [N x T] tensor of epsilon MSEs for each timestep. + """ + device = x_start.device + batch_size = x_start.shape[0] + + vb = [] + xstart_mse = [] + mse = [] + for t in list(range(self.num_timesteps))[::-1]: + t_batch = th.tensor([t] * batch_size, device=device) + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) + # Calculate VLB term at the current timestep + with th.no_grad(): + out = self._vb_terms_bpd( + model, + x_start=x_start, + x_t=x_t, + t=t_batch, + clip_denoised=clip_denoised, + model_kwargs=model_kwargs, + ) + vb.append(out["output"]) + xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) + eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) + mse.append(mean_flat((eps - noise) ** 2)) + + vb = th.stack(vb, dim=1) + xstart_mse = th.stack(xstart_mse, dim=1) + mse = th.stack(mse, dim=1) + + prior_bpd = self._prior_bpd(x_start) + total_bpd = vb.sum(dim=1) + prior_bpd + return { + "total_bpd": total_bpd, + "prior_bpd": prior_bpd, + "vb": vb, + "xstart_mse": xstart_mse, + "mse": mse, + } + + +def _extract_into_tensor(arr, timesteps, broadcast_shape): + """ + Extract values from a 1-D numpy array for a batch of indices. + :param arr: the 1-D numpy array. + :param timesteps: a tensor of indices into the array to extract. + :param broadcast_shape: a larger shape of K dimensions with the batch + dimension equal to the length of timesteps. + :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. + """ + res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() + while len(res.shape) < len(broadcast_shape): + res = res[..., None] + return res + th.zeros(broadcast_shape, device=timesteps.device) diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/respace.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/respace.py new file mode 100644 index 0000000000000000000000000000000000000000..0a2cc0435d1ace54466585db9043b284973d454e --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/respace.py @@ -0,0 +1,129 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +import numpy as np +import torch as th + +from .gaussian_diffusion import GaussianDiffusion + + +def space_timesteps(num_timesteps, section_counts): + """ + Create a list of timesteps to use from an original diffusion process, + given the number of timesteps we want to take from equally-sized portions + of the original process. + For example, if there's 300 timesteps and the section counts are [10,15,20] + then the first 100 timesteps are strided to be 10 timesteps, the second 100 + are strided to be 15 timesteps, and the final 100 are strided to be 20. + If the stride is a string starting with "ddim", then the fixed striding + from the DDIM paper is used, and only one section is allowed. + :param num_timesteps: the number of diffusion steps in the original + process to divide up. + :param section_counts: either a list of numbers, or a string containing + comma-separated numbers, indicating the step count + per section. As a special case, use "ddimN" where N + is a number of steps to use the striding from the + DDIM paper. + :return: a set of diffusion steps from the original process to use. + """ + if isinstance(section_counts, str): + if section_counts.startswith("ddim"): + desired_count = int(section_counts[len("ddim") :]) + for i in range(1, num_timesteps): + if len(range(0, num_timesteps, i)) == desired_count: + return set(range(0, num_timesteps, i)) + raise ValueError( + f"cannot create exactly {num_timesteps} steps with an integer stride" + ) + section_counts = [int(x) for x in section_counts.split(",")] + size_per = num_timesteps // len(section_counts) + extra = num_timesteps % len(section_counts) + start_idx = 0 + all_steps = [] + for i, section_count in enumerate(section_counts): + size = size_per + (1 if i < extra else 0) + if size < section_count: + raise ValueError( + f"cannot divide section of {size} steps into {section_count}" + ) + if section_count <= 1: + frac_stride = 1 + else: + frac_stride = (size - 1) / (section_count - 1) + cur_idx = 0.0 + taken_steps = [] + for _ in range(section_count): + taken_steps.append(start_idx + round(cur_idx)) + cur_idx += frac_stride + all_steps += taken_steps + start_idx += size + return set(all_steps) + + +class SpacedDiffusion(GaussianDiffusion): + """ + A diffusion process which can skip steps in a base diffusion process. + :param use_timesteps: a collection (sequence or set) of timesteps from the + original diffusion process to retain. + :param kwargs: the kwargs to create the base diffusion process. + """ + + def __init__(self, use_timesteps, **kwargs): + self.use_timesteps = set(use_timesteps) + self.timestep_map = [] + self.original_num_steps = len(kwargs["betas"]) + + base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa + last_alpha_cumprod = 1.0 + new_betas = [] + for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): + if i in self.use_timesteps: + new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) + last_alpha_cumprod = alpha_cumprod + self.timestep_map.append(i) + kwargs["betas"] = np.array(new_betas) + super().__init__(**kwargs) + + def p_mean_variance( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) + + def training_losses( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().training_losses(self._wrap_model(model), *args, **kwargs) + + def condition_mean(self, cond_fn, *args, **kwargs): + return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) + + def condition_score(self, cond_fn, *args, **kwargs): + return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) + + def _wrap_model(self, model): + if isinstance(model, _WrappedModel): + return model + return _WrappedModel( + model, self.timestep_map, self.original_num_steps + ) + + def _scale_timesteps(self, t): + # Scaling is done by the wrapped model. + return t + + +class _WrappedModel: + def __init__(self, model, timestep_map, original_num_steps): + self.model = model + self.timestep_map = timestep_map + # self.rescale_timesteps = rescale_timesteps + self.original_num_steps = original_num_steps + + def __call__(self, x, ts, **kwargs): + map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) + new_ts = map_tensor[ts] + # if self.rescale_timesteps: + # new_ts = new_ts.float() * (1000.0 / self.original_num_steps) + return self.model(x, new_ts, **kwargs) diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/timestep_sampler.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/timestep_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..a3f369847677d8dbaaadb8297691b1be92cf189f --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/diffusion/timestep_sampler.py @@ -0,0 +1,150 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +from abc import ABC, abstractmethod + +import numpy as np +import torch as th +import torch.distributed as dist + + +def create_named_schedule_sampler(name, diffusion): + """ + Create a ScheduleSampler from a library of pre-defined samplers. + :param name: the name of the sampler. + :param diffusion: the diffusion object to sample for. + """ + if name == "uniform": + return UniformSampler(diffusion) + elif name == "loss-second-moment": + return LossSecondMomentResampler(diffusion) + else: + raise NotImplementedError(f"unknown schedule sampler: {name}") + + +class ScheduleSampler(ABC): + """ + A distribution over timesteps in the diffusion process, intended to reduce + variance of the objective. + By default, samplers perform unbiased importance sampling, in which the + objective's mean is unchanged. + However, subclasses may override sample() to change how the resampled + terms are reweighted, allowing for actual changes in the objective. + """ + + @abstractmethod + def weights(self): + """ + Get a numpy array of weights, one per diffusion step. + The weights needn't be normalized, but must be positive. + """ + + def sample(self, batch_size, device): + """ + Importance-sample timesteps for a batch. + :param batch_size: the number of timesteps. + :param device: the torch device to save to. + :return: a tuple (timesteps, weights): + - timesteps: a tensor of timestep indices. + - weights: a tensor of weights to scale the resulting losses. + """ + w = self.weights() + p = w / np.sum(w) + indices_np = np.random.choice(len(p), size=(batch_size,), p=p) + indices = th.from_numpy(indices_np).long().to(device) + weights_np = 1 / (len(p) * p[indices_np]) + weights = th.from_numpy(weights_np).float().to(device) + return indices, weights + + +class UniformSampler(ScheduleSampler): + def __init__(self, diffusion): + self.diffusion = diffusion + self._weights = np.ones([diffusion.num_timesteps]) + + def weights(self): + return self._weights + + +class LossAwareSampler(ScheduleSampler): + def update_with_local_losses(self, local_ts, local_losses): + """ + Update the reweighting using losses from a model. + Call this method from each rank with a batch of timesteps and the + corresponding losses for each of those timesteps. + This method will perform synchronization to make sure all of the ranks + maintain the exact same reweighting. + :param local_ts: an integer Tensor of timesteps. + :param local_losses: a 1D Tensor of losses. + """ + batch_sizes = [ + th.tensor([0], dtype=th.int32, device=local_ts.device) + for _ in range(dist.get_world_size()) + ] + dist.all_gather( + batch_sizes, + th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), + ) + + # Pad all_gather batches to be the maximum batch size. + batch_sizes = [x.item() for x in batch_sizes] + max_bs = max(batch_sizes) + + timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] + loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] + dist.all_gather(timestep_batches, local_ts) + dist.all_gather(loss_batches, local_losses) + timesteps = [ + x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] + ] + losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] + self.update_with_all_losses(timesteps, losses) + + @abstractmethod + def update_with_all_losses(self, ts, losses): + """ + Update the reweighting using losses from a model. + Sub-classes should override this method to update the reweighting + using losses from the model. + This method directly updates the reweighting without synchronizing + between workers. It is called by update_with_local_losses from all + ranks with identical arguments. Thus, it should have deterministic + behavior to maintain state across workers. + :param ts: a list of int timesteps. + :param losses: a list of float losses, one per timestep. + """ + + +class LossSecondMomentResampler(LossAwareSampler): + def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): + self.diffusion = diffusion + self.history_per_term = history_per_term + self.uniform_prob = uniform_prob + self._loss_history = np.zeros( + [diffusion.num_timesteps, history_per_term], dtype=np.float64 + ) + self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) + + def weights(self): + if not self._warmed_up(): + return np.ones([self.diffusion.num_timesteps], dtype=np.float64) + weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) + weights /= np.sum(weights) + weights *= 1 - self.uniform_prob + weights += self.uniform_prob / len(weights) + return weights + + def update_with_all_losses(self, ts, losses): + for t, loss in zip(ts, losses): + if self._loss_counts[t] == self.history_per_term: + # Shift out the oldest loss term. + self._loss_history[t, :-1] = self._loss_history[t, 1:] + self._loss_history[t, -1] = loss + else: + self._loss_history[t, self._loss_counts[t]] = loss + self._loss_counts[t] += 1 + + def _warmed_up(self): + return (self._loss_counts == self.history_per_term).all() diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/models_add_cross_concate.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/models_add_cross_concate.py new file mode 100644 index 0000000000000000000000000000000000000000..c8f6ea523047abd4e78de8d9f409cbc4ac064361 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/models_add_cross_concate.py @@ -0,0 +1,498 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# GLIDE: https://github.com/openai/glide-text2im +# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import numpy as np +import math +from timm.models.vision_transformer import PatchEmbed, Attention, Mlp +from torch import einsum +from einops import rearrange, repeat +from autoencoder import * + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +################################################################################# +# Embedding Layers for Timesteps and Class Labels # +################################################################################# + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb + + +################################################################################# +# Core DiT Model # +################################################################################# + + +class CrossAttention(nn.Module): + def __init__(self, query_dim, heads=8, dropout=0.): + super().__init__() + dim_head = query_dim / heads + + self.scale = dim_head ** -0.5 + self.heads = heads + self.to_q = nn.Linear(query_dim, query_dim, bias=True) + self.to_k = nn.Linear(query_dim, query_dim, bias=True) + self.to_v = nn.Linear(query_dim, query_dim, bias=True) + + def forward(self, x, context=None): + h = self.heads + q = self.to_q(x) + k = self.to_k(context).unsqueeze(1) + v = self.to_v(context).unsqueeze(1) + + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) + sim = einsum('b i d, b j d -> b i j', q, k) * self.scale + + # attention, what we cannot get enough of + attn = sim.softmax(dim=-2) + + out = einsum('b i j, b j d -> b i d', attn, v) + out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + attn_out = rearrange(attn, '(b h) n d -> b n (h d)', h=h) + return out, attn_out + + +class Cross_Norm(nn.Module): + def __init__(self, hidden_size, num_heads): + super().__init__() + self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.cross_attention = CrossAttention(hidden_size, heads=num_heads) + + def forward(self, x, c): + x = self.norm(x) + x = self.cross_attention(x, c) + + return x + + + +class DiTBlock(nn.Module): + """ + A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning. + """ + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + mlp_hidden_dim = int(hidden_size * mlp_ratio) + approx_gelu = lambda: nn.GELU() + self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 6 * hidden_size, bias=True) + ) + + def forward(self, x, c): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) + x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa)) + x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) + return x + + +class FinalLayer(nn.Module): + """ + The final layer of DiT. + """ + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True) + ) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class temp_Adaptive_Mask(nn.Module): + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + self.out_channels = out_channels + self.patch_size = patch_size + self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.mlp = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True), + nn.Linear(2 * hidden_size, hidden_size, bias=True) + ) + + + def forward(self, x): + x = self.norm(x) + x = self.mlp(x) + x = self.linear(x) + h = w = int(x.shape[1] ** 0.5) + assert h * w == x.shape[1] + c = self.out_channels + p = self.patch_size + + x = x.reshape(shape=(x.shape[0], h, w, p, p, c)) + x = torch.einsum('nhwpqc->nchpwq', x) + x = x.reshape(shape=(x.shape[0], c, h * p, h * p)) + return x + + +class DiT(nn.Module): + """ + Diffusion model with a Transformer backbone. + """ + def __init__( + self, + input_size=32, + patch_size=2, + in_channels=4, + hidden_size=1152, + depth=28, + num_heads=16, + mlp_ratio=4.0, + class_dropout_prob=0.1, + num_classes=1000, + learn_sigma=True, + ): + super().__init__() + self.learn_sigma = learn_sigma + self.in_channels = in_channels + self.out_channels = in_channels * 2 if learn_sigma else in_channels + self.patch_size = patch_size + self.num_heads = num_heads + + self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True) + self.t_embedder = TimestepEmbedder(hidden_size) + self.y_embedders = nn.Linear(1024, 1152) + num_patches = self.x_embedder.num_patches + + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False) + + self.blocks = nn.ModuleList([ + DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth) + ]) + + self.cross_defect = nn.ModuleList([ + Cross_Norm(hidden_size, num_heads) for _ in range(10) + ]) + self.adapt_mask = temp_Adaptive_Mask(num_heads*10, patch_size, in_channels) + + self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels) + self.initialize_weights() + + def initialize_weights(self): + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize (and freeze) pos_embed by sin-cos embedding: + pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5)) + self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) + + # Initialize patch_embed like nn.Linear (instead of nn.Conv2d): + w = self.x_embedder.proj.weight.data + nn.init.xavier_uniform_(w.view([w.shape[0], -1])) + nn.init.constant_(self.x_embedder.proj.bias, 0) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.final_layer.linear.weight, 0) + nn.init.constant_(self.final_layer.linear.bias, 0) + + def unpatchify(self, x): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.out_channels + p = self.x_embedder.patch_size[0] + h = w = int(x.shape[1] ** 0.5) + assert h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], h, w, p, p, c)) + x = torch.einsum('nhwpqc->nchpwq', x) + imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p)) + return imgs + + def forward(self, x, t, y): + """ + Forward pass of DiT. + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + t: (N,) tensor of diffusion timesteps + y: (N,) tensor of class labels + """ + x = self.x_embedder(x) + self.pos_embed # (N, T, D), where T = H * W / patch_size ** 2 + t = self.t_embedder(t) # (N, D) + y_defect = self.y_embedders(y[0]) + y_class = self.y_embedders(y[1]) # (N, D) + y_all = self.y_embedders(y[2]) + att_map = [] + loss_att = 0 + for i in range(28): + block = self.blocks[i] + if i < 10: + c = t + y_class + x = block(x, c) # (N, T, D) + elif i < 20: + cross = self.cross_defect[i - 10] + c = t + y_defect + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaLN_modulation(c).chunk(6, + dim=1) + x = x + gate_msa.unsqueeze(1) * block.attn(modulate(block.norm1(x), shift_msa, scale_msa)) + + cross_att, att_weight = cross(x, c) + loss_att += att_weight + att_map.append(att_weight) + x = x + cross_att + + x = x + gate_mlp.unsqueeze(1) * block.mlp(modulate(block.norm2(x), shift_mlp, scale_mlp)) + + elif i < 28: + c = t + y_all + x = block(x, c) + + x = self.final_layer(x, c) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x) # (N, out_channels, H, W) + att_map = torch.cat(att_map, dim=-1) + att_mask = self.adapt_mask(att_map) + + return x, att_mask, loss_att.resize(x.shape[0], x.shape[2]//2, x.shape[3]//2, 16).mean(dim=-1) + + def forward_free_2(self, x, t, y, mask_temp=None): + """ + Forward pass of DiT. + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + t: (N,) tensor of diffusion timesteps + y: (N,) tensor of class labels + """ + x = self.x_embedder(x) + self.pos_embed # (N, T, D), where T = H * W / patch_size ** 2 + t = self.t_embedder(t) # (N, D) + y_defect = self.y_embedders(torch.cat([y[0][0], y[1][0]], dim=0)) + y_class = self.y_embedders(torch.cat([y[0][1], y[1][1]], dim=0)) # (N, D) + y_all = self.y_embedders(torch.cat([y[0][2], y[1][2]], dim=0)) + att_map = [] + loss_att = 0 + for i in range(28): + block = self.blocks[i] + if i < 10: + c = t + y_class + x = block(x, c) # (N, T, D) + elif i < 20: + cross = self.cross_defect[i - 10] + c = t + y_defect + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaLN_modulation(c).chunk(6, + dim=1) + x = x + gate_msa.unsqueeze(1) * block.attn(modulate(block.norm1(x), shift_msa, scale_msa)) + cross_att, att_weight = cross(x, c) + att_map.append(att_weight) + loss_att+=att_weight + x = x + cross_att + x = x + gate_mlp.unsqueeze(1) * block.mlp(modulate(block.norm2(x), shift_mlp, scale_mlp)) + + elif i < 28: + c = t + y_all + x = block(x, c) + + x = self.final_layer(x, c) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x) # (N, out_channels, H, W) + att_map = torch.cat(att_map, dim=-1) + att_mask = self.adapt_mask(att_map) + + return x, att_mask, loss_att + + def forward_with_cfg_2(self, x, t, y, cfg_scale): + """ + Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. + """ + # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb + half = x[: len(x) // 2] + combined = torch.cat([half, half], dim=0) + model_out, mask, _ = self.forward_free_2(combined, t, y) + # For exact reproducibility reasons, we apply classifier-free guidance on only + # three channels by default. The standard approach to cfg applies it to all channels. + # This can be done by uncommenting the following line and commenting-out the line following that. + # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] + eps, rest = model_out[:, :3], model_out[:, 3:] + cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) + half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) + eps = torch.cat([half_eps, half_eps], dim=0) + return torch.cat([eps, rest], dim=1), mask, _ + + def forward_free_3(self, x, t, y, mask_temp=None): + """ + Forward pass of DiT. + x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + t: (N,) tensor of diffusion timesteps + y: (N,) tensor of class labels + """ + x = self.x_embedder(x) + self.pos_embed # (N, T, D), where T = H * W / patch_size ** 2 + t = self.t_embedder(t) # (N, D) + y_defect = self.y_embedders(torch.cat([y[0][0], y[1][0], y[2][0]], dim=0)) + y_class = self.y_embedders(torch.cat([y[0][1], y[1][1], y[2][1]], dim=0)) # (N, D) + y_all = self.y_embedders(torch.cat([y[0][2], y[1][2], y[2][2]], dim=0)) + att_map = [] + att_loss = 0 + for i in range(28): + block = self.blocks[i] + if i < 10: + c = t + y_class + x = block(x, c) # (N, T, D) + elif i < 20: + cross = self.cross_defect[i - 10] + c = t + y_defect + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaLN_modulation(c).chunk(6, + dim=1) + x = x + gate_msa.unsqueeze(1) * block.attn(modulate(block.norm1(x), shift_msa, scale_msa)) + cross_att, att_weight = cross(x, c) + att_map.append(att_weight) + att_loss += att_weight + x = x + cross_att + x = x + gate_mlp.unsqueeze(1) * block.mlp(modulate(block.norm2(x), shift_mlp, scale_mlp)) + + elif i < 28: + c = t + y_all + x = block(x, c) + + x = self.final_layer(x, c) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x) # (N, out_channels, H, W) + att_map = torch.cat(att_map, dim=-1) + att_mask = self.adapt_mask(att_map) + + return x, att_mask, att_loss + + def forward_with_cfg_3(self, x, t, y, cfg_scale): + """ + Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. + """ + # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb + half = x[: len(x) // 3] + combined = torch.cat([half, half, half], dim=0) + model_out, mask, _ = self.forward_free_3(combined, t, y) + # For exact reproducibility reasons, we apply classifier-free guidance on only + # three channels by default. The standard approach to cfg applies it to all channels. + # This can be done by uncommenting the following line and commenting-out the line following that. + # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] + eps, rest = model_out[:, :3], model_out[:, 3:] + cond_eps, uncond_eps_defect, uncond_eps = torch.split(eps, len(eps) // 3, dim=0) + half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps_defect) + cfg_scale * (uncond_eps_defect - uncond_eps) + eps = torch.cat([half_eps, half_eps, half_eps], dim=0) + return torch.cat([eps, rest], dim=1), mask, _ + +################################################################################# +# Sine/Cosine Positional Embedding Functions # +################################################################################# +# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py + +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): + """ + grid_size: int of the grid height and width + return: + pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + grid_h = np.arange(grid_size, dtype=np.float32) + grid_w = np.arange(grid_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size, grid_size]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token and extra_tokens > 0: + pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + assert embed_dim % 2 == 0 + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) + + emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) + return emb + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2. + omega = 1. / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/test.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/test.py new file mode 100644 index 0000000000000000000000000000000000000000..2616b5ed46d948e6deb818e21b45249a71f030e8 --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/test.py @@ -0,0 +1,198 @@ +import os +import argparse +import torch +import numpy as np +from torchvision.utils import save_image +from diffusers.models import AutoencoderKL +import clip.clip as clip + +from models_add_cross_concate import DiT +from diffusion import create_diffusion +from autoencoder import * + +# Enable TF32 for fast execution on modern NVIDIA GPUs +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True + + +def rgb_to_gray(tensor): + r, g, b = tensor[:, 0], tensor[:, 1], tensor[:, 2] + gray = 0.299 * r + 0.587 * g + 0.114 * b + return gray + + +def iterative_thresholding_batch(gray_tensor): + gray_np = gray_tensor.detach().cpu().numpy() + binarized = np.zeros_like(gray_np, dtype=np.uint8) + + for i in range(gray_np.shape[0]): + img = gray_np[i] + T = img.mean() + prev_T = -1 + + while abs(T - prev_T) > 1e-4: + prev_T = T + G1 = img[img >= T] + G2 = img[img < T] + m1 = G1.mean() if G1.size > 0 else 0 + m2 = G2.mean() if G2.size > 0 else 0 + T = (m1 + m2) / 2 + + binarized[i] = (img >= T).astype(np.uint8) + + return torch.from_numpy(binarized).to(gray_tensor.device) + + +def binarize_tensor_iterative(x): + gray = rgb_to_gray(x) + binary = iterative_thresholding_batch(gray) + return binary.unsqueeze(1) + + +def get_label(data_path): + """Safely extracts defect labels from 4-level dataset hierarchy.""" + label_list1 = [] + if not os.path.exists(data_path): + return label_list1 + + for name_class in os.listdir(data_path): + img_dir = os.path.join(data_path, name_class, 'img') + if os.path.exists(img_dir) and os.path.isdir(img_dir): + for class_object in os.listdir(img_dir): + defect_dir = os.path.join(img_dir, class_object) + if os.path.isdir(defect_dir) and class_object != 'good': + label_list1.append(f"{class_object} {name_class}") + return label_list1 + + +def gen(args): + data_path = args.data + label_list = get_label(data_path) + + if not label_list: + print(f"โŒ No valid defect subfolders found in {data_path}. Please check directory structure.") + return + + print(f"๐Ÿ“‹ Found defect categories to generate: {label_list}") + + image_size = args.imagesize + device = "cuda" if torch.cuda.is_available() else "cpu" + latent_size = image_size // 8 + + # 1. Load CLIP model + model_clip, _ = clip.load('RN50', device) + + # 2. Setup DiT architecture and weights + model = DiT( + depth=28, hidden_size=1152, patch_size=2, + num_heads=16, input_size=latent_size, num_classes=1000 + ).to(device) + + print(f"๐Ÿ“ฆ Loading checkpoint from: {args.ckpt}") + checkpoint = torch.load(args.ckpt, map_location=device) + if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: + model.load_state_dict(checkpoint['model_state_dict']) + else: + model.load_state_dict(checkpoint) + + model.eval() + + # 3. Setup VAE and Diffusion pipeline + diffusion = create_diffusion(timestep_respacing="50") + vae = AutoencoderKL.from_pretrained(args.vae).to(device) + + os.makedirs(args.out_dir, exist_ok=True) + num_img = args.batchsize + + # 4. Generate specified number of output batches (replaces infinite while loop) + for sample_round in range(args.num_samples): + print(f"\n๐Ÿ”„ --- Generating Batch {sample_round + 1}/{args.num_samples} ---") + + for c in label_list: + defect_name, class_name = c.split()[0], c.split()[1] + print(f"๐ŸŽจ Generating defect: '{defect_name}' on object: '{class_name}'...") + + # Prepare text embeddings for dual-branch CFG + y_null_product = torch.cat([clip.tokenize("a photo of good industry")] * num_img).to(device) + y_null_good = torch.cat([clip.tokenize(f"a photo of good {class_name}")] * num_img).to(device) + + with torch.no_grad(): + y_null_product = model_clip.encode_text(y_null_product) + y_null_good = model_clip.encode_text(y_null_good) + + y_null_product = (y_null_product / y_null_product.norm(dim=-1, keepdim=True)).float() + y_null_good = (y_null_good / y_null_good.norm(dim=-1, keepdim=True)).float() + + only_good = torch.cat([clip.tokenize("a photo of good")] * num_img).to(device) + defect = torch.cat([clip.tokenize(f"a photo of {defect_name}")] * num_img).to(device) + classes = torch.cat([clip.tokenize(f"a photo of {class_name}")] * num_img).to(device) + classes_industry = torch.cat([clip.tokenize("a photo of industry")] * num_img).to(device) + y_all = torch.cat([clip.tokenize(f"a photo of {c}")] * num_img).to(device) + + with torch.no_grad(): + only_good = model_clip.encode_text(only_good) + defect = model_clip.encode_text(defect) + classes = model_clip.encode_text(classes) + classes_industry = model_clip.encode_text(classes_industry) + y_all = model_clip.encode_text(y_all) + + only_good = (only_good / only_good.norm(dim=-1, keepdim=True)).float() + defect = (defect / defect.norm(dim=-1, keepdim=True)).float() + classes_industry = (classes_industry / classes_industry.norm(dim=-1, keepdim=True)).float() + classes = (classes / classes.norm(dim=-1, keepdim=True)).float() + y_all = (y_all / y_all.norm(dim=-1, keepdim=True)).float() + + y_defect_class = [defect, classes, y_all] + y_good_class = [only_good, classes, y_null_good] + + z = torch.randn(num_img, 4, latent_size, latent_size, device=device) + z = torch.cat([z, z], 0) + + y = [y_defect_class, y_good_class] + + for num in np.arange(0.5, 3.0, 0.5): + model_kwargs = dict(y=y, cfg_scale=float(num)) + + with torch.no_grad(): + samples, cross = diffusion.p_sample_loop( + model.forward_with_cfg_2, + z.shape, + z, + clip_denoised=False, + model_kwargs=model_kwargs, + progress=False, + device=device + ) + + img_gen, _ = samples.chunk(2, dim=0) + mask_gen, _ = cross.chunk(2, dim=0) + + with torch.no_grad(): + img_gen = vae.decode(img_gen / 0.18215).sample + mask_gen = vae.decode(mask_gen / 0.18215).sample + + # Save generated images and binarized masks + img_path = os.path.join(args.out_dir, f"{class_name}_{defect_name}_cfg{num:.1f}_b{sample_round}.png") + mask_path = os.path.join(args.out_dir, f"{class_name}_{defect_name}_cfg{num:.1f}_b{sample_round}_mask.png") + + save_image(img_gen, img_path, nrow=2, normalize=True) + + mask_gen = binarize_tensor_iterative(mask_gen) + mask_gen = (mask_gen * 255).to(torch.uint8).float() / 255.0 + save_image(mask_gen, mask_path, nrow=2, normalize=True) + + print(f"\nโœจ Generation complete! Synthetic pairs saved to: {os.path.abspath(args.out_dir)}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--batchsize", type=int, default=2) + parser.add_argument("--num_samples", type=int, default=1, help="Number of sampling passes to run.") + parser.add_argument("--data", type=str, required=True) + parser.add_argument("--imagesize", type=int, choices=[256, 512], default=512) + parser.add_argument("--ckpt", type=str, required=True, help="Path to fine-tuned checkpoint.") + parser.add_argument("--vae", type=str, required=True, help="Path to VAE checkpoint.") + parser.add_argument("--out_dir", type=str, default="./generated_results", help="Directory to save generated samples.") + + args = parser.parse_args() + gen(args) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/engine/DefectDiffu/train.py b/ArtiAgent - DefectDiffu/engine/DefectDiffu/train.py new file mode 100644 index 0000000000000000000000000000000000000000..e9efd0e07002a742491f0b6cce65f828ea3e311f --- /dev/null +++ b/ArtiAgent - DefectDiffu/engine/DefectDiffu/train.py @@ -0,0 +1,231 @@ +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) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/LICENSE b/ArtiAgent - DefectDiffu/src/GroundingDINO/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..b1395e94b016dd1b95b4c7e3ed493e1d0b342917 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2020 - present, Facebook, Inc + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/README.md b/ArtiAgent - DefectDiffu/src/GroundingDINO/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b6610df03d409633e572ef49d67a445d35a63967 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/README.md @@ -0,0 +1,163 @@ +# Grounding DINO + +--- + +[![arXiv](https://img.shields.io/badge/arXiv-2303.05499-b31b1b.svg)](https://arxiv.org/abs/2303.05499) +[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/wxWDt5UiwY8) +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/zero-shot-object-detection-with-grounding-dino.ipynb) +[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/cMa77r3YrDk) +[![HuggingFace space](https://img.shields.io/badge/๐Ÿค—-HuggingFace%20Space-cyan.svg)](https://huggingface.co/spaces/ShilongLiu/Grounding_DINO_demo) + +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/zero-shot-object-detection-on-mscoco)](https://paperswithcode.com/sota/zero-shot-object-detection-on-mscoco?p=grounding-dino-marrying-dino-with-grounded) \ +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/zero-shot-object-detection-on-odinw)](https://paperswithcode.com/sota/zero-shot-object-detection-on-odinw?p=grounding-dino-marrying-dino-with-grounded) \ +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/object-detection-on-coco-minival)](https://paperswithcode.com/sota/object-detection-on-coco-minival?p=grounding-dino-marrying-dino-with-grounded) \ +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/object-detection-on-coco)](https://paperswithcode.com/sota/object-detection-on-coco?p=grounding-dino-marrying-dino-with-grounded) + + + +Official PyTorch implementation of [Grounding DINO](https://arxiv.org/abs/2303.05499), a stronger open-set object detector. Code is available now! + + +## Highlight + +- **Open-Set Detection.** Detect **everything** with language! +- **High Performancce.** COCO zero-shot **52.5 AP** (training without COCO data!). COCO fine-tune **63.0 AP**. +- **Flexible.** Collaboration with Stable Diffusion for Image Editting. + +## News +[2023/03/28] A YouTube [video](https://youtu.be/cMa77r3YrDk) about Grounding DINO and basic object detection prompt engineering. [[SkalskiP](https://github.com/SkalskiP)] \ +[2023/03/28] Add a [demo](https://huggingface.co/spaces/ShilongLiu/Grounding_DINO_demo) on Hugging Face Space! \ +[2023/03/27] Support CPU-only mode. Now the model can run on machines without GPUs.\ +[2023/03/25] A [demo](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/zero-shot-object-detection-with-grounding-dino.ipynb) for Grounding DINO is available at Colab. [[SkalskiP](https://github.com/SkalskiP)] \ +[2023/03/22] Code is available Now! + +
+ +Description + +ODinW +
+ + + +## TODO + +- [x] Release inference code and demo. +- [x] Release checkpoints. +- [ ] Grounding DINO with Stable Diffusion and GLIGEN demos. +- [ ] Release training codes. + +## Install + +If you have a CUDA environment, please make sure the environment variable `CUDA_HOME` is set. It will be compiled under CPU-only mode if no CUDA available. + +```bash +pip install -e . +``` + +## Demo + +```bash +CUDA_VISIBLE_DEVICES=6 python demo/inference_on_a_image.py \ + -c /path/to/config \ + -p /path/to/checkpoint \ + -i .asset/cats.png \ + -o "outputs/0" \ + -t "cat ear." \ + [--cpu-only] # open it for cpu mode +``` +See the `demo/inference_on_a_image.py` for more details. + +**Web UI** + +We also provide a demo code to integrate Grounding DINO with Gradio Web UI. See the file `demo/gradio_app.py` for more details. + +## Checkpoints + + + + + + + + + + + + + + + + + + + + + + + + + +
namebackboneDatabox AP on COCOCheckpointConfig
1GroundingDINO-TSwin-TO365,GoldG,Cap4M48.4 (zero-shot) / 57.2 (fine-tune)Github link | HF linklink
+ +## Results + +
+ +COCO Object Detection Results + +COCO +
+ +
+ +ODinW Object Detection Results + +ODinW +
+ +
+ +Marrying Grounding DINO with Stable Diffusion for Image Editing + +GD_SD +
+ +
+ +Marrying Grounding DINO with GLIGEN for more Detailed Image Editing + +GD_GLIGEN +
+ +## Model + +Includes: a text backbone, an image backbone, a feature enhancer, a language-guided query selection, and a cross-modality decoder. + +![arch](.asset/arch.png) + + +## Acknowledgement + +Our model is related to [DINO](https://github.com/IDEA-Research/DINO) and [GLIP](https://github.com/microsoft/GLIP). Thanks for their great work! + +We also thank great previous work including DETR, Deformable DETR, SMCA, Conditional DETR, Anchor DETR, Dynamic DETR, DAB-DETR, DN-DETR, etc. More related work are available at [Awesome Detection Transformer](https://github.com/IDEACVR/awesome-detection-transformer). A new toolbox [detrex](https://github.com/IDEA-Research/detrex) is available as well. + +Thanks [Stable Diffusion](https://github.com/Stability-AI/StableDiffusion) and [GLIGEN](https://github.com/gligen/GLIGEN) for their awesome models. + + +## Citation + +If you find our work helpful for your research, please consider citing the following BibTeX entry. + +```bibtex +@inproceedings{ShilongLiu2023GroundingDM, + title={Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection}, + author={Shilong Liu and Zhaoyang Zeng and Tianhe Ren and Feng Li and Hao Zhang and Jie Yang and Chunyuan Li and Jianwei Yang and Hang Su and Jun Zhu and Lei Zhang}, + year={2023} +} +``` + + + + diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/__init__.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5963a5e6797a32718001851e7a31bc82f911ff0c Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinB.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinB.py new file mode 100644 index 0000000000000000000000000000000000000000..f490c4bbd598a35de43d36ceafcbd769e7ff21bf --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinB.py @@ -0,0 +1,43 @@ +batch_size = 1 +modelname = "groundingdino" +backbone = "swin_B_384_22k" +position_embedding = "sine" +pe_temperatureH = 20 +pe_temperatureW = 20 +return_interm_indices = [1, 2, 3] +backbone_freeze_keywords = None +enc_layers = 6 +dec_layers = 6 +pre_norm = False +dim_feedforward = 2048 +hidden_dim = 256 +dropout = 0.0 +nheads = 8 +num_queries = 900 +query_dim = 4 +num_patterns = 0 +num_feature_levels = 4 +enc_n_points = 4 +dec_n_points = 4 +two_stage_type = "standard" +two_stage_bbox_embed_share = False +two_stage_class_embed_share = False +transformer_activation = "relu" +dec_pred_bbox_embed_share = True +dn_box_noise_scale = 1.0 +dn_label_noise_ratio = 0.5 +dn_label_coef = 1.0 +dn_bbox_coef = 1.0 +embed_init_tgt = True +dn_labelbook_size = 2000 +max_text_len = 256 +text_encoder_type = "bert-base-uncased" +use_text_enhancer = True +use_fusion_layer = True +use_checkpoint = True +use_transformer_ckpt = True +use_text_cross_attention = True +text_dropout = 0.0 +fusion_dropout = 0.0 +fusion_droppath = 0.1 +sub_sentence_present = True diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py new file mode 100644 index 0000000000000000000000000000000000000000..9158d5f6260ec74bded95377d382387430d7cd70 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py @@ -0,0 +1,43 @@ +batch_size = 1 +modelname = "groundingdino" +backbone = "swin_T_224_1k" +position_embedding = "sine" +pe_temperatureH = 20 +pe_temperatureW = 20 +return_interm_indices = [1, 2, 3] +backbone_freeze_keywords = None +enc_layers = 6 +dec_layers = 6 +pre_norm = False +dim_feedforward = 2048 +hidden_dim = 256 +dropout = 0.0 +nheads = 8 +num_queries = 900 +query_dim = 4 +num_patterns = 0 +num_feature_levels = 4 +enc_n_points = 4 +dec_n_points = 4 +two_stage_type = "standard" +two_stage_bbox_embed_share = False +two_stage_class_embed_share = False +transformer_activation = "relu" +dec_pred_bbox_embed_share = True +dn_box_noise_scale = 1.0 +dn_label_noise_ratio = 0.5 +dn_label_coef = 1.0 +dn_bbox_coef = 1.0 +embed_init_tgt = True +dn_labelbook_size = 2000 +max_text_len = 256 +text_encoder_type = "bert-base-uncased" +use_text_enhancer = True +use_fusion_layer = True +use_checkpoint = True +use_transformer_ckpt = True +use_text_cross_attention = True +text_dropout = 0.0 +fusion_dropout = 0.0 +fusion_droppath = 0.1 +sub_sentence_present = True diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__init__.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..217e2897203cb3ae56638ee6c9df19512e2cf0c3 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__pycache__/transforms.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__pycache__/transforms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..292ecd04e8c8a5a73e2248e2cf999de15cc3c092 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/__pycache__/transforms.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/transforms.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..91cf9269e4b31008a3ddca34a19b038a9b399991 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/datasets/transforms.py @@ -0,0 +1,311 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Transforms and data augmentation for both image + bbox. +""" +import os +import random + +import PIL +import torch +import torchvision.transforms as T +import torchvision.transforms.functional as F + +from groundingdino.util.box_ops import box_xyxy_to_cxcywh +from groundingdino.util.misc import interpolate + + +def crop(image, target, region): + cropped_image = F.crop(image, *region) + + target = target.copy() + i, j, h, w = region + + # should we do something wrt the original size? + target["size"] = torch.tensor([h, w]) + + fields = ["labels", "area", "iscrowd", "positive_map"] + + if "boxes" in target: + boxes = target["boxes"] + max_size = torch.as_tensor([w, h], dtype=torch.float32) + cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) + cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) + cropped_boxes = cropped_boxes.clamp(min=0) + area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) + target["boxes"] = cropped_boxes.reshape(-1, 4) + target["area"] = area + fields.append("boxes") + + if "masks" in target: + # FIXME should we update the area here if there are no boxes? + target["masks"] = target["masks"][:, i : i + h, j : j + w] + fields.append("masks") + + # remove elements for which the boxes or masks that have zero area + if "boxes" in target or "masks" in target: + # favor boxes selection when defining which elements to keep + # this is compatible with previous implementation + if "boxes" in target: + cropped_boxes = target["boxes"].reshape(-1, 2, 2) + keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) + else: + keep = target["masks"].flatten(1).any(1) + + for field in fields: + if field in target: + target[field] = target[field][keep] + + if os.environ.get("IPDB_SHILONG_DEBUG", None) == "INFO": + # for debug and visualization only. + if "strings_positive" in target: + target["strings_positive"] = [ + _i for _i, _j in zip(target["strings_positive"], keep) if _j + ] + + return cropped_image, target + + +def hflip(image, target): + flipped_image = F.hflip(image) + + w, h = image.size + + target = target.copy() + if "boxes" in target: + boxes = target["boxes"] + boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor( + [w, 0, w, 0] + ) + target["boxes"] = boxes + + if "masks" in target: + target["masks"] = target["masks"].flip(-1) + + return flipped_image, target + + +def resize(image, target, size, max_size=None): + # size can be min_size (scalar) or (w, h) tuple + + def get_size_with_aspect_ratio(image_size, size, max_size=None): + w, h = image_size + if max_size is not None: + min_original_size = float(min((w, h))) + max_original_size = float(max((w, h))) + if max_original_size / min_original_size * size > max_size: + size = int(round(max_size * min_original_size / max_original_size)) + + if (w <= h and w == size) or (h <= w and h == size): + return (h, w) + + if w < h: + ow = size + oh = int(size * h / w) + else: + oh = size + ow = int(size * w / h) + + return (oh, ow) + + def get_size(image_size, size, max_size=None): + if isinstance(size, (list, tuple)): + return size[::-1] + else: + return get_size_with_aspect_ratio(image_size, size, max_size) + + size = get_size(image.size, size, max_size) + rescaled_image = F.resize(image, size) + + if target is None: + return rescaled_image, None + + ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size)) + ratio_width, ratio_height = ratios + + target = target.copy() + if "boxes" in target: + boxes = target["boxes"] + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height] + ) + target["boxes"] = scaled_boxes + + if "area" in target: + area = target["area"] + scaled_area = area * (ratio_width * ratio_height) + target["area"] = scaled_area + + h, w = size + target["size"] = torch.tensor([h, w]) + + if "masks" in target: + target["masks"] = ( + interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0] > 0.5 + ) + + return rescaled_image, target + + +def pad(image, target, padding): + # assumes that we only pad on the bottom right corners + padded_image = F.pad(image, (0, 0, padding[0], padding[1])) + if target is None: + return padded_image, None + target = target.copy() + # should we do something wrt the original size? + target["size"] = torch.tensor(padded_image.size[::-1]) + if "masks" in target: + target["masks"] = torch.nn.functional.pad(target["masks"], (0, padding[0], 0, padding[1])) + return padded_image, target + + +class ResizeDebug(object): + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + return resize(img, target, self.size) + + +class RandomCrop(object): + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + region = T.RandomCrop.get_params(img, self.size) + return crop(img, target, region) + + +class RandomSizeCrop(object): + def __init__(self, min_size: int, max_size: int, respect_boxes: bool = False): + # respect_boxes: True to keep all boxes + # False to tolerence box filter + self.min_size = min_size + self.max_size = max_size + self.respect_boxes = respect_boxes + + def __call__(self, img: PIL.Image.Image, target: dict): + init_boxes = len(target["boxes"]) + max_patience = 10 + for i in range(max_patience): + w = random.randint(self.min_size, min(img.width, self.max_size)) + h = random.randint(self.min_size, min(img.height, self.max_size)) + region = T.RandomCrop.get_params(img, [h, w]) + result_img, result_target = crop(img, target, region) + if ( + not self.respect_boxes + or len(result_target["boxes"]) == init_boxes + or i == max_patience - 1 + ): + return result_img, result_target + return result_img, result_target + + +class CenterCrop(object): + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + image_width, image_height = img.size + crop_height, crop_width = self.size + crop_top = int(round((image_height - crop_height) / 2.0)) + crop_left = int(round((image_width - crop_width) / 2.0)) + return crop(img, target, (crop_top, crop_left, crop_height, crop_width)) + + +class RandomHorizontalFlip(object): + def __init__(self, p=0.5): + self.p = p + + def __call__(self, img, target): + if random.random() < self.p: + return hflip(img, target) + return img, target + + +class RandomResize(object): + def __init__(self, sizes, max_size=None): + assert isinstance(sizes, (list, tuple)) + self.sizes = sizes + self.max_size = max_size + + def __call__(self, img, target=None): + size = random.choice(self.sizes) + return resize(img, target, size, self.max_size) + + +class RandomPad(object): + def __init__(self, max_pad): + self.max_pad = max_pad + + def __call__(self, img, target): + pad_x = random.randint(0, self.max_pad) + pad_y = random.randint(0, self.max_pad) + return pad(img, target, (pad_x, pad_y)) + + +class RandomSelect(object): + """ + Randomly selects between transforms1 and transforms2, + with probability p for transforms1 and (1 - p) for transforms2 + """ + + def __init__(self, transforms1, transforms2, p=0.5): + self.transforms1 = transforms1 + self.transforms2 = transforms2 + self.p = p + + def __call__(self, img, target): + if random.random() < self.p: + return self.transforms1(img, target) + return self.transforms2(img, target) + + +class ToTensor(object): + def __call__(self, img, target): + return F.to_tensor(img), target + + +class RandomErasing(object): + def __init__(self, *args, **kwargs): + self.eraser = T.RandomErasing(*args, **kwargs) + + def __call__(self, img, target): + return self.eraser(img), target + + +class Normalize(object): + def __init__(self, mean, std): + self.mean = mean + self.std = std + + def __call__(self, image, target=None): + image = F.normalize(image, mean=self.mean, std=self.std) + if target is None: + return image, None + target = target.copy() + h, w = image.shape[-2:] + if "boxes" in target: + boxes = target["boxes"] + boxes = box_xyxy_to_cxcywh(boxes) + boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) + target["boxes"] = boxes + return image, target + + +class Compose(object): + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, image, target): + for t in self.transforms: + image, target = t(image, target) + return image, target + + def __repr__(self): + format_string = self.__class__.__name__ + "(" + for t in self.transforms: + format_string += "\n" + format_string += " {0}".format(t) + format_string += "\n)" + return format_string diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__init__.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2af819d61d589cfec2e0ca46612a7456f42b831a --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__init__.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copied from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +from .groundingdino import build_groundingdino diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0140700cc85de677d5fad29bc7bfa8a5e31f9aa4 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/groundingdino.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/groundingdino.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78b1c51c7da0ecea65b2a8fa4e6bb736d25cb1ea Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/groundingdino.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/__init__.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..76e4b272b479a26c63d120c818c140870cd8c287 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/__init__.py @@ -0,0 +1 @@ +from .backbone import build_backbone diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/backbone.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..c8340c723fad8e07e2fc62daaa3912487498814b --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/backbone.py @@ -0,0 +1,221 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copied from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +""" +Backbone modules. +""" + +from typing import Dict, List + +import torch +import torch.nn.functional as F +import torchvision +from torch import nn +from torchvision.models._utils import IntermediateLayerGetter + +from groundingdino.util.misc import NestedTensor, clean_state_dict, is_main_process + +from .position_encoding import build_position_encoding +from .swin_transformer import build_swin_transformer + + +class FrozenBatchNorm2d(torch.nn.Module): + """ + BatchNorm2d where the batch statistics and the affine parameters are fixed. + + Copy-paste from torchvision.misc.ops with added eps before rqsrt, + without which any other models than torchvision.models.resnet[18,34,50,101] + produce nans. + """ + + def __init__(self, n): + super(FrozenBatchNorm2d, self).__init__() + self.register_buffer("weight", torch.ones(n)) + self.register_buffer("bias", torch.zeros(n)) + self.register_buffer("running_mean", torch.zeros(n)) + self.register_buffer("running_var", torch.ones(n)) + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + num_batches_tracked_key = prefix + "num_batches_tracked" + if num_batches_tracked_key in state_dict: + del state_dict[num_batches_tracked_key] + + super(FrozenBatchNorm2d, self)._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + def forward(self, x): + # move reshapes to the beginning + # to make it fuser-friendly + w = self.weight.reshape(1, -1, 1, 1) + b = self.bias.reshape(1, -1, 1, 1) + rv = self.running_var.reshape(1, -1, 1, 1) + rm = self.running_mean.reshape(1, -1, 1, 1) + eps = 1e-5 + scale = w * (rv + eps).rsqrt() + bias = b - rm * scale + return x * scale + bias + + +class BackboneBase(nn.Module): + def __init__( + self, + backbone: nn.Module, + train_backbone: bool, + num_channels: int, + return_interm_indices: list, + ): + super().__init__() + for name, parameter in backbone.named_parameters(): + if ( + not train_backbone + or "layer2" not in name + and "layer3" not in name + and "layer4" not in name + ): + parameter.requires_grad_(False) + + return_layers = {} + for idx, layer_index in enumerate(return_interm_indices): + return_layers.update( + {"layer{}".format(5 - len(return_interm_indices) + idx): "{}".format(layer_index)} + ) + + # if len: + # if use_stage1_feature: + # return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} + # else: + # return_layers = {"layer2": "0", "layer3": "1", "layer4": "2"} + # else: + # return_layers = {'layer4': "0"} + self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) + self.num_channels = num_channels + + def forward(self, tensor_list: NestedTensor): + xs = self.body(tensor_list.tensors) + out: Dict[str, NestedTensor] = {} + for name, x in xs.items(): + m = tensor_list.mask + assert m is not None + mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] + out[name] = NestedTensor(x, mask) + # import ipdb; ipdb.set_trace() + return out + + +class Backbone(BackboneBase): + """ResNet backbone with frozen BatchNorm.""" + + def __init__( + self, + name: str, + train_backbone: bool, + dilation: bool, + return_interm_indices: list, + batch_norm=FrozenBatchNorm2d, + ): + if name in ["resnet18", "resnet34", "resnet50", "resnet101"]: + backbone = getattr(torchvision.models, name)( + replace_stride_with_dilation=[False, False, dilation], + pretrained=is_main_process(), + norm_layer=batch_norm, + ) + else: + raise NotImplementedError("Why you can get here with name {}".format(name)) + # num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 + assert name not in ("resnet18", "resnet34"), "Only resnet50 and resnet101 are available." + assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]] + num_channels_all = [256, 512, 1024, 2048] + num_channels = num_channels_all[4 - len(return_interm_indices) :] + super().__init__(backbone, train_backbone, num_channels, return_interm_indices) + + +class Joiner(nn.Sequential): + def __init__(self, backbone, position_embedding): + super().__init__(backbone, position_embedding) + + def forward(self, tensor_list: NestedTensor): + xs = self[0](tensor_list) + out: List[NestedTensor] = [] + pos = [] + for name, x in xs.items(): + out.append(x) + # position encoding + pos.append(self[1](x).to(x.tensors.dtype)) + + return out, pos + + +def build_backbone(args): + """ + Useful args: + - backbone: backbone name + - lr_backbone: + - dilation + - return_interm_indices: available: [0,1,2,3], [1,2,3], [3] + - backbone_freeze_keywords: + - use_checkpoint: for swin only for now + + """ + position_embedding = build_position_encoding(args) + train_backbone = True + if not train_backbone: + raise ValueError("Please set lr_backbone > 0") + return_interm_indices = args.return_interm_indices + assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]] + args.backbone_freeze_keywords + use_checkpoint = getattr(args, "use_checkpoint", False) + + if args.backbone in ["resnet50", "resnet101"]: + backbone = Backbone( + args.backbone, + train_backbone, + args.dilation, + return_interm_indices, + batch_norm=FrozenBatchNorm2d, + ) + bb_num_channels = backbone.num_channels + elif args.backbone in [ + "swin_T_224_1k", + "swin_B_224_22k", + "swin_B_384_22k", + "swin_L_224_22k", + "swin_L_384_22k", + ]: + pretrain_img_size = int(args.backbone.split("_")[-2]) + backbone = build_swin_transformer( + args.backbone, + pretrain_img_size=pretrain_img_size, + out_indices=tuple(return_interm_indices), + dilation=False, + use_checkpoint=use_checkpoint, + ) + + bb_num_channels = backbone.num_features[4 - len(return_interm_indices) :] + else: + raise NotImplementedError("Unknown backbone {}".format(args.backbone)) + + assert len(bb_num_channels) == len( + return_interm_indices + ), f"len(bb_num_channels) {len(bb_num_channels)} != len(return_interm_indices) {len(return_interm_indices)}" + + model = Joiner(backbone, position_embedding) + model.num_channels = bb_num_channels + assert isinstance( + bb_num_channels, List + ), "bb_num_channels is expected to be a List but {}".format(type(bb_num_channels)) + # import ipdb; ipdb.set_trace() + return model diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/position_encoding.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/position_encoding.py new file mode 100644 index 0000000000000000000000000000000000000000..eac7e896bbe85a670824bfe8ef487d0535d5bd99 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/position_encoding.py @@ -0,0 +1,186 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copied from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +""" +Various positional encodings for the transformer. +""" +import math + +import torch +from torch import nn + +from groundingdino.util.misc import NestedTensor + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + mask = tensor_list.mask + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + # if os.environ.get("SHILONG_AMP", None) == '1': + # eps = 1e-4 + # else: + # eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos + + +class PositionEmbeddingSineHW(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__( + self, num_pos_feats=64, temperatureH=10000, temperatureW=10000, normalize=False, scale=None + ): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperatureH = temperatureH + self.temperatureW = temperatureW + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + mask = tensor_list.mask + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + + # import ipdb; ipdb.set_trace() + + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_tx = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_tx = self.temperatureW ** (2 * (torch.div(dim_tx, 2, rounding_mode='floor')) / self.num_pos_feats) + pos_x = x_embed[:, :, :, None] / dim_tx + + dim_ty = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_ty = self.temperatureH ** (2 * (torch.div(dim_ty, 2, rounding_mode='floor')) / self.num_pos_feats) + pos_y = y_embed[:, :, :, None] / dim_ty + + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + + # import ipdb; ipdb.set_trace() + + return pos + + +class PositionEmbeddingLearned(nn.Module): + """ + Absolute pos embedding, learned. + """ + + def __init__(self, num_pos_feats=256): + super().__init__() + self.row_embed = nn.Embedding(50, num_pos_feats) + self.col_embed = nn.Embedding(50, num_pos_feats) + self.reset_parameters() + + def reset_parameters(self): + nn.init.uniform_(self.row_embed.weight) + nn.init.uniform_(self.col_embed.weight) + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + h, w = x.shape[-2:] + i = torch.arange(w, device=x.device) + j = torch.arange(h, device=x.device) + x_emb = self.col_embed(i) + y_emb = self.row_embed(j) + pos = ( + torch.cat( + [ + x_emb.unsqueeze(0).repeat(h, 1, 1), + y_emb.unsqueeze(1).repeat(1, w, 1), + ], + dim=-1, + ) + .permute(2, 0, 1) + .unsqueeze(0) + .repeat(x.shape[0], 1, 1, 1) + ) + return pos + + +def build_position_encoding(args): + N_steps = args.hidden_dim // 2 + if args.position_embedding in ("v2", "sine"): + # TODO find a better way of exposing other arguments + position_embedding = PositionEmbeddingSineHW( + N_steps, + temperatureH=args.pe_temperatureH, + temperatureW=args.pe_temperatureW, + normalize=True, + ) + elif args.position_embedding in ("v3", "learned"): + position_embedding = PositionEmbeddingLearned(N_steps) + else: + raise ValueError(f"not supported {args.position_embedding}") + + return position_embedding diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..fa8837e4001e41dfed6af99e6619f8b03b989824 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py @@ -0,0 +1,802 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# -------------------------------------------------------- +# modified from https://github.com/SwinTransformer/Swin-Transformer-Object-Detection/blob/master/mmdet/models/backbones/swin_transformer.py +# -------------------------------------------------------- + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +from groundingdino.util.misc import NestedTensor + + +class Mlp(nn.Module): + """Multilayer perceptron.""" + + def __init__( + self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0 + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=0.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """Forward function. + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """Swin Transformer Block. + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__( + self, + dim, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, + window_size=to_2tuple(self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop + ) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """Patch Merging Layer + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """A basic Swin Transformer layer for one stage. + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__( + self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + ): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + Hp = int(np.ceil(H / self.window_size)) * self.window_size + Wp = int(np.ceil(W / self.window_size)) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device, dtype=x.dtype) # 1 Hp Wp 1 + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( + attn_mask == 0, float(0.0) + ) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding + Args: + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + dilation (bool): if True, the output size if 16x downsample, ow 32x downsample. + """ + + def __init__( + self, + pretrain_img_size=224, + patch_size=4, + in_chans=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + dilation=False, + use_checkpoint=False, + ): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + self.dilation = dilation + + # if use_checkpoint: + # print("use_checkpoint!!!!!!!!!!!!!!!!!!!!!!!!") + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [ + pretrain_img_size[0] // patch_size[0], + pretrain_img_size[1] // patch_size[1], + ] + + self.absolute_pos_embed = nn.Parameter( + torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]) + ) + trunc_normal_(self.absolute_pos_embed, std=0.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + # prepare downsample list + downsamplelist = [PatchMerging for i in range(self.num_layers)] + downsamplelist[-1] = None + num_features = [int(embed_dim * 2**i) for i in range(self.num_layers)] + if self.dilation: + downsamplelist[-2] = None + num_features[-1] = int(embed_dim * 2 ** (self.num_layers - 1)) // 2 + for i_layer in range(self.num_layers): + layer = BasicLayer( + # dim=int(embed_dim * 2 ** i_layer), + dim=num_features[i_layer], + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + # downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + downsample=downsamplelist[i_layer], + use_checkpoint=use_checkpoint, + ) + self.layers.append(layer) + + # num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f"norm{i_layer}" + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + # def init_weights(self, pretrained=None): + # """Initialize the weights in backbone. + # Args: + # pretrained (str, optional): Path to pre-trained weights. + # Defaults to None. + # """ + + # def _init_weights(m): + # if isinstance(m, nn.Linear): + # trunc_normal_(m.weight, std=.02) + # if isinstance(m, nn.Linear) and m.bias is not None: + # nn.init.constant_(m.bias, 0) + # elif isinstance(m, nn.LayerNorm): + # nn.init.constant_(m.bias, 0) + # nn.init.constant_(m.weight, 1.0) + + # if isinstance(pretrained, str): + # self.apply(_init_weights) + # logger = get_root_logger() + # load_checkpoint(self, pretrained, strict=False, logger=logger) + # elif pretrained is None: + # self.apply(_init_weights) + # else: + # raise TypeError('pretrained must be a str or None') + + def forward_raw(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate( + self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic" + ) + x = (x + absolute_pos_embed).flatten(2).transpose(1, 2) # B Wh*Ww C + else: + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + + outs = [] + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + # import ipdb; ipdb.set_trace() + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + # in: + # torch.Size([2, 3, 1024, 1024]) + # outs: + # [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \ + # torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])] + return tuple(outs) + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate( + self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic" + ) + x = (x + absolute_pos_embed).flatten(2).transpose(1, 2) # B Wh*Ww C + else: + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + + outs = [] + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + # in: + # torch.Size([2, 3, 1024, 1024]) + # out: + # [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \ + # torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])] + + # collect for nesttensors + outs_dict = {} + for idx, out_i in enumerate(outs): + m = tensor_list.mask + assert m is not None + mask = F.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0] + outs_dict[idx] = NestedTensor(out_i, mask) + + return outs_dict + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + + +def build_swin_transformer(modelname, pretrain_img_size, **kw): + assert modelname in [ + "swin_T_224_1k", + "swin_B_224_22k", + "swin_B_384_22k", + "swin_L_224_22k", + "swin_L_384_22k", + ] + + model_para_dict = { + "swin_T_224_1k": dict( + embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7 + ), + "swin_B_224_22k": dict( + embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=7 + ), + "swin_B_384_22k": dict( + embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12 + ), + "swin_L_224_22k": dict( + embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=7 + ), + "swin_L_384_22k": dict( + embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12 + ), + } + kw_cgf = model_para_dict[modelname] + kw_cgf.update(kw) + model = SwinTransformer(pretrain_img_size=pretrain_img_size, **kw_cgf) + return model + + +if __name__ == "__main__": + model = build_swin_transformer("swin_L_384_22k", 384, dilation=True) + x = torch.rand(2, 3, 1024, 1024) + y = model.forward_raw(x) + import ipdb + + ipdb.set_trace() + x = torch.rand(2, 3, 384, 384) + y = model.forward_raw(x) diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/bertwarper.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/bertwarper.py new file mode 100644 index 0000000000000000000000000000000000000000..f0cf9779b270e1aead32845006f8b881fcba37ad --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/bertwarper.py @@ -0,0 +1,273 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from torch import Tensor, nn +from torchvision.ops.boxes import nms +from transformers import BertConfig, BertModel, BertPreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions + + +class BertModelWarper(nn.Module): + def __init__(self, bert_model): + super().__init__() + # self.bert = bert_modelc + + self.config = bert_model.config + self.embeddings = bert_model.embeddings + self.encoder = bert_model.encoder + self.pooler = bert_model.pooler + + self.get_extended_attention_mask = bert_model.get_extended_attention_mask + self.invert_attention_mask = bert_model.invert_attention_mask + self.get_head_mask = bert_model.get_head_mask + + def forward( + self, + input_ids=None, + attention_mask=None, + token_type_ids=None, + position_ids=None, + head_mask=None, + inputs_embeds=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + r""" + encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). + """ + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + batch_size, seq_length = input_shape + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + batch_size, seq_length = input_shape + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + # past_key_values_length + past_key_values_length = ( + past_key_values[0][0].shape[2] if past_key_values is not None else 0 + ) + + if attention_mask is None: + attention_mask = torch.ones( + ((batch_size, seq_length + past_key_values_length)), device=device + ) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask( + attention_mask, input_shape, device + ) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + +class TextEncoderShell(nn.Module): + def __init__(self, text_encoder): + super().__init__() + self.text_encoder = text_encoder + self.config = self.text_encoder.config + + def forward(self, **kw): + # feed into text encoder + return self.text_encoder(**kw) + + +def generate_masks_with_special_tokens(tokenized, special_tokens_list, tokenizer): + """Generate attention mask between each pair of special tokens + Args: + input_ids (torch.Tensor): input ids. Shape: [bs, num_token] + special_tokens_mask (list): special tokens mask. + Returns: + torch.Tensor: attention mask between each special tokens. + """ + input_ids = tokenized["input_ids"] + bs, num_token = input_ids.shape + # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens + special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool() + for special_token in special_tokens_list: + special_tokens_mask |= input_ids == special_token + + # idxs: each row is a list of indices of special tokens + idxs = torch.nonzero(special_tokens_mask) + + # generate attention mask and positional ids + attention_mask = ( + torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1) + ) + position_ids = torch.zeros((bs, num_token), device=input_ids.device) + previous_col = 0 + for i in range(idxs.shape[0]): + row, col = idxs[i] + if (col == 0) or (col == num_token - 1): + attention_mask[row, col, col] = True + position_ids[row, col] = 0 + else: + attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True + position_ids[row, previous_col + 1 : col + 1] = torch.arange( + 0, col - previous_col, device=input_ids.device + ) + + previous_col = col + + # # padding mask + # padding_mask = tokenized['attention_mask'] + # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool() + + return attention_mask, position_ids.to(torch.long) + + +def generate_masks_with_special_tokens_and_transfer_map(tokenized, special_tokens_list, tokenizer): + """Generate attention mask between each pair of special tokens + Args: + input_ids (torch.Tensor): input ids. Shape: [bs, num_token] + special_tokens_mask (list): special tokens mask. + Returns: + torch.Tensor: attention mask between each special tokens. + """ + input_ids = tokenized["input_ids"] + bs, num_token = input_ids.shape + # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens + special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool() + for special_token in special_tokens_list: + special_tokens_mask |= input_ids == special_token + + # idxs: each row is a list of indices of special tokens + idxs = torch.nonzero(special_tokens_mask) + + # generate attention mask and positional ids + attention_mask = ( + torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1) + ) + position_ids = torch.zeros((bs, num_token), device=input_ids.device) + cate_to_token_mask_list = [[] for _ in range(bs)] + previous_col = 0 + for i in range(idxs.shape[0]): + row, col = idxs[i] + if (col == 0) or (col == num_token - 1): + attention_mask[row, col, col] = True + position_ids[row, col] = 0 + else: + attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True + position_ids[row, previous_col + 1 : col + 1] = torch.arange( + 0, col - previous_col, device=input_ids.device + ) + c2t_maski = torch.zeros((num_token), device=input_ids.device).bool() + c2t_maski[previous_col + 1 : col] = True + cate_to_token_mask_list[row].append(c2t_maski) + previous_col = col + + cate_to_token_mask_list = [ + torch.stack(cate_to_token_mask_listi, dim=0) + for cate_to_token_mask_listi in cate_to_token_mask_list + ] + + # # padding mask + # padding_mask = tokenized['attention_mask'] + # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool() + + return attention_mask, position_ids.to(torch.long), cate_to_token_mask_list diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn.h b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn.h new file mode 100644 index 0000000000000000000000000000000000000000..c7408eba007b424194618baa63726657e36875e3 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn.h @@ -0,0 +1,64 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once + +#include "ms_deform_attn_cpu.h" + +#ifdef WITH_CUDA +#include "ms_deform_attn_cuda.h" +#endif + +namespace groundingdino { + +at::Tensor +ms_deform_attn_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + if (value.type().is_cuda()) + { +#ifdef WITH_CUDA + return ms_deform_attn_cuda_forward( + value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step); +#else + AT_ERROR("Not compiled with GPU support"); +#endif + } + AT_ERROR("Not implemented on the CPU"); +} + +std::vector +ms_deform_attn_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + if (value.type().is_cuda()) + { +#ifdef WITH_CUDA + return ms_deform_attn_cuda_backward( + value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step); +#else + AT_ERROR("Not compiled with GPU support"); +#endif + } + AT_ERROR("Not implemented on the CPU"); +} + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.cpp b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.cpp new file mode 100644 index 0000000000000000000000000000000000000000..551243fdadfd1682b5dc6628623b67a79b3f6c74 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.cpp @@ -0,0 +1,43 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include + +#include +#include + +namespace groundingdino { + +at::Tensor +ms_deform_attn_cpu_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ERROR("Not implement on cpu"); +} + +std::vector +ms_deform_attn_cpu_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + AT_ERROR("Not implement on cpu"); +} + +} // namespace groundingdino diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.h b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.h new file mode 100644 index 0000000000000000000000000000000000000000..b2b88e8c46f19b6db0933163e57ccdb51180f517 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.h @@ -0,0 +1,35 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once +#include + +namespace groundingdino { + +at::Tensor +ms_deform_attn_cpu_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step); + +std::vector +ms_deform_attn_cpu_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step); + +} // namespace groundingdino diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..d04fae8a9a45c11e4e74f3035e94762796da4096 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu @@ -0,0 +1,156 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include +#include "ms_deform_im2col_cuda.cuh" + +#include +#include +#include +#include + +namespace groundingdino { + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto output = at::zeros({batch, num_query, num_heads, channels}, value.options()); + + const int batch_n = im2col_step_; + auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto columns = output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] { + ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + columns.data()); + + })); + } + + output = output.view({batch, num_query, num_heads*channels}); + + return output; +} + + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto grad_value = at::zeros_like(value); + auto grad_sampling_loc = at::zeros_like(sampling_loc); + auto grad_attn_weight = at::zeros_like(attn_weight); + + const int batch_n = im2col_step_; + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto grad_output_g = grad_output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] { + ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(), + grad_output_g.data(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + grad_value.data() + n * im2col_step_ * per_value_size, + grad_sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + grad_attn_weight.data() + n * im2col_step_ * per_attn_weight_size); + + })); + } + + return { + grad_value, grad_sampling_loc, grad_attn_weight + }; +} + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.h b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..ad1311a78f61303616504eb991aaa9c4a93d9948 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.h @@ -0,0 +1,33 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once +#include + +namespace groundingdino { + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step); + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step); + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_im2col_cuda.cuh b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_im2col_cuda.cuh new file mode 100644 index 0000000000000000000000000000000000000000..6bc2acb7aea0eab2e9e91e769a16861e1652c284 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_im2col_cuda.cuh @@ -0,0 +1,1327 @@ +/*! +************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************** +* Modified from DCN (https://github.com/msracver/Deformable-ConvNets) +* Copyright (c) 2018 Microsoft +************************************************************************** +*/ + +#include +#include +#include + +#include +#include + +#include + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +const int CUDA_NUM_THREADS = 1024; +inline int GET_BLOCKS(const int N, const int num_threads) +{ + return (N + num_threads - 1) / num_threads; +} + + +template +__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + } + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + return val; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + *grad_attn_weight = top_grad * val; + *grad_sampling_loc = width * grad_w_weight * top_grad_value; + *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + atomicAdd(grad_attn_weight, top_grad * val); + atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value); + atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value); +} + + +template +__global__ void ms_deformable_im2col_gpu_kernel(const int n, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *data_col) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + scalar_t *data_col_ptr = data_col + index; + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + scalar_t col = 0; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride); + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight; + } + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + } + } + *data_col_ptr = col; + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockSize; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockSize/2; s>0; s>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockDim.x; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]); + atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]); + atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]); + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear_gm( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + grad_sampling_loc, grad_attn_weight); + } + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +void ms_deformable_im2col_cuda(cudaStream_t stream, + const scalar_t* data_value, + const int64_t* data_spatial_shapes, + const int64_t* data_level_start_index, + const scalar_t* data_sampling_loc, + const scalar_t* data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* data_col) +{ + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + const int num_threads = CUDA_NUM_THREADS; + ms_deformable_im2col_gpu_kernel + <<>>( + num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, + batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +template +void ms_deformable_col2im_cuda(cudaStream_t stream, + const scalar_t* grad_col, + const scalar_t* data_value, + const int64_t * data_spatial_shapes, + const int64_t * data_level_start_index, + const scalar_t * data_sampling_loc, + const scalar_t * data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels; + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + if (channels > 1024) + { + if ((channels & 1023) == 0) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_gm + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + else{ + switch(channels) + { + case 1: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 2: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 4: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 8: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 16: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 32: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 64: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 128: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 256: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 512: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 1024: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + default: + if (channels < 64) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err)); + } + +} \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/cuda_version.cu b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/cuda_version.cu new file mode 100644 index 0000000000000000000000000000000000000000..64569e34ffb250964de27e33e7a53f3822270b9e --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/cuda_version.cu @@ -0,0 +1,7 @@ +#include + +namespace groundingdino { +int get_cudart_version() { + return CUDART_VERSION; +} +} // namespace groundingdino diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/vision.cpp b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/vision.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c1f2c50c82909bbd5492c163d634af77a3ba1781 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/vision.cpp @@ -0,0 +1,58 @@ +// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved + +#include "MsDeformAttn/ms_deform_attn.h" + +namespace groundingdino { + +#ifdef WITH_CUDA +extern int get_cudart_version(); +#endif + +std::string get_cuda_version() { +#ifdef WITH_CUDA + std::ostringstream oss; + + // copied from + // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/detail/CUDAHooks.cpp#L231 + auto printCudaStyleVersion = [&](int v) { + oss << (v / 1000) << "." << (v / 10 % 100); + if (v % 10 != 0) { + oss << "." << (v % 10); + } + }; + printCudaStyleVersion(get_cudart_version()); + return oss.str(); +#else + return std::string("not available"); +#endif +} + +// similar to +// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Version.cpp +std::string get_compiler_version() { + std::ostringstream ss; +#if defined(__GNUC__) +#ifndef __clang__ + { ss << "GCC " << __GNUC__ << "." << __GNUC_MINOR__; } +#endif +#endif + +#if defined(__clang_major__) + { + ss << "clang " << __clang_major__ << "." << __clang_minor__ << "." + << __clang_patchlevel__; + } +#endif + +#if defined(_MSC_VER) + { ss << "MSVC " << _MSC_FULL_VER; } +#endif + return ss.str(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward"); + m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward"); +} + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..2753b3ddee43c7a9fe28d1824db5d786e7e1ad59 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py @@ -0,0 +1,297 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.models.layers import DropPath + + +class FeatureResizer(nn.Module): + """ + This class takes as input a set of embeddings of dimension C1 and outputs a set of + embedding of dimension C2, after a linear transformation, dropout and normalization (LN). + """ + + def __init__(self, input_feat_size, output_feat_size, dropout, do_ln=True): + super().__init__() + self.do_ln = do_ln + # Object feature encoding + self.fc = nn.Linear(input_feat_size, output_feat_size, bias=True) + self.layer_norm = nn.LayerNorm(output_feat_size, eps=1e-12) + self.dropout = nn.Dropout(dropout) + + def forward(self, encoder_features): + x = self.fc(encoder_features) + if self.do_ln: + x = self.layer_norm(x) + output = self.dropout(x) + return output + + +def l1norm(X, dim, eps=1e-8): + """L1-normalize columns of X""" + norm = torch.abs(X).sum(dim=dim, keepdim=True) + eps + X = torch.div(X, norm) + return X + + +def l2norm(X, dim, eps=1e-8): + """L2-normalize columns of X""" + norm = torch.pow(X, 2).sum(dim=dim, keepdim=True).sqrt() + eps + X = torch.div(X, norm) + return X + + +def func_attention(query, context, smooth=1, raw_feature_norm="softmax", eps=1e-8): + """ + query: (n_context, queryL, d) + context: (n_context, sourceL, d) + """ + batch_size_q, queryL = query.size(0), query.size(1) + batch_size, sourceL = context.size(0), context.size(1) + + # Get attention + # --> (batch, d, queryL) + queryT = torch.transpose(query, 1, 2) + + # (batch, sourceL, d)(batch, d, queryL) + # --> (batch, sourceL, queryL) + attn = torch.bmm(context, queryT) + if raw_feature_norm == "softmax": + # --> (batch*sourceL, queryL) + attn = attn.view(batch_size * sourceL, queryL) + attn = nn.Softmax()(attn) + # --> (batch, sourceL, queryL) + attn = attn.view(batch_size, sourceL, queryL) + elif raw_feature_norm == "l2norm": + attn = l2norm(attn, 2) + elif raw_feature_norm == "clipped_l2norm": + attn = nn.LeakyReLU(0.1)(attn) + attn = l2norm(attn, 2) + else: + raise ValueError("unknown first norm type:", raw_feature_norm) + # --> (batch, queryL, sourceL) + attn = torch.transpose(attn, 1, 2).contiguous() + # --> (batch*queryL, sourceL) + attn = attn.view(batch_size * queryL, sourceL) + attn = nn.Softmax()(attn * smooth) + # --> (batch, queryL, sourceL) + attn = attn.view(batch_size, queryL, sourceL) + # --> (batch, sourceL, queryL) + attnT = torch.transpose(attn, 1, 2).contiguous() + + # --> (batch, d, sourceL) + contextT = torch.transpose(context, 1, 2) + # (batch x d x sourceL)(batch x sourceL x queryL) + # --> (batch, d, queryL) + weightedContext = torch.bmm(contextT, attnT) + # --> (batch, queryL, d) + weightedContext = torch.transpose(weightedContext, 1, 2) + + return weightedContext, attnT + + +class BiMultiHeadAttention(nn.Module): + def __init__(self, v_dim, l_dim, embed_dim, num_heads, dropout=0.1, cfg=None): + super(BiMultiHeadAttention, self).__init__() + + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.v_dim = v_dim + self.l_dim = l_dim + + assert ( + self.head_dim * self.num_heads == self.embed_dim + ), f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})." + self.scale = self.head_dim ** (-0.5) + self.dropout = dropout + + self.v_proj = nn.Linear(self.v_dim, self.embed_dim) + self.l_proj = nn.Linear(self.l_dim, self.embed_dim) + self.values_v_proj = nn.Linear(self.v_dim, self.embed_dim) + self.values_l_proj = nn.Linear(self.l_dim, self.embed_dim) + + self.out_v_proj = nn.Linear(self.embed_dim, self.v_dim) + self.out_l_proj = nn.Linear(self.embed_dim, self.l_dim) + + self.stable_softmax_2d = True + self.clamp_min_for_underflow = True + self.clamp_max_for_overflow = True + + self._reset_parameters() + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def _reset_parameters(self): + nn.init.xavier_uniform_(self.v_proj.weight) + self.v_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.l_proj.weight) + self.l_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.values_v_proj.weight) + self.values_v_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.values_l_proj.weight) + self.values_l_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.out_v_proj.weight) + self.out_v_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.out_l_proj.weight) + self.out_l_proj.bias.data.fill_(0) + + def forward(self, v, l, attention_mask_v=None, attention_mask_l=None): + """_summary_ + + Args: + v (_type_): bs, n_img, dim + l (_type_): bs, n_text, dim + attention_mask_v (_type_, optional): _description_. bs, n_img + attention_mask_l (_type_, optional): _description_. bs, n_text + + Returns: + _type_: _description_ + """ + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + bsz, tgt_len, _ = v.size() + + query_states = self.v_proj(v) * self.scale + key_states = self._shape(self.l_proj(l), -1, bsz) + value_v_states = self._shape(self.values_v_proj(v), -1, bsz) + value_l_states = self._shape(self.values_l_proj(l), -1, bsz) + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + key_states = key_states.view(*proj_shape) + value_v_states = value_v_states.view(*proj_shape) + value_l_states = value_l_states.view(*proj_shape) + + src_len = key_states.size(1) + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) # bs*nhead, nimg, ntxt + + if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is {attn_weights.size()}" + ) + + if self.stable_softmax_2d: + attn_weights = attn_weights - attn_weights.max() + + if self.clamp_min_for_underflow: + attn_weights = torch.clamp( + attn_weights, min=-50000 + ) # Do not increase -50000, data type half has quite limited range + if self.clamp_max_for_overflow: + attn_weights = torch.clamp( + attn_weights, max=50000 + ) # Do not increase 50000, data type half has quite limited range + + attn_weights_T = attn_weights.transpose(1, 2) + attn_weights_l = attn_weights_T - torch.max(attn_weights_T, dim=-1, keepdim=True)[0] + if self.clamp_min_for_underflow: + attn_weights_l = torch.clamp( + attn_weights_l, min=-50000 + ) # Do not increase -50000, data type half has quite limited range + if self.clamp_max_for_overflow: + attn_weights_l = torch.clamp( + attn_weights_l, max=50000 + ) # Do not increase 50000, data type half has quite limited range + + # mask vison for language + if attention_mask_v is not None: + attention_mask_v = ( + attention_mask_v[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1) + ) + attn_weights_l.masked_fill_(attention_mask_v, float("-inf")) + + attn_weights_l = attn_weights_l.softmax(dim=-1) + + # mask language for vision + if attention_mask_l is not None: + attention_mask_l = ( + attention_mask_l[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1) + ) + attn_weights.masked_fill_(attention_mask_l, float("-inf")) + attn_weights_v = attn_weights.softmax(dim=-1) + + attn_probs_v = F.dropout(attn_weights_v, p=self.dropout, training=self.training) + attn_probs_l = F.dropout(attn_weights_l, p=self.dropout, training=self.training) + + attn_output_v = torch.bmm(attn_probs_v, value_l_states) + attn_output_l = torch.bmm(attn_probs_l, value_v_states) + + if attn_output_v.size() != (bsz * self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output_v` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is {attn_output_v.size()}" + ) + + if attn_output_l.size() != (bsz * self.num_heads, src_len, self.head_dim): + raise ValueError( + f"`attn_output_l` should be of size {(bsz, self.num_heads, src_len, self.head_dim)}, but is {attn_output_l.size()}" + ) + + attn_output_v = attn_output_v.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output_v = attn_output_v.transpose(1, 2) + attn_output_v = attn_output_v.reshape(bsz, tgt_len, self.embed_dim) + + attn_output_l = attn_output_l.view(bsz, self.num_heads, src_len, self.head_dim) + attn_output_l = attn_output_l.transpose(1, 2) + attn_output_l = attn_output_l.reshape(bsz, src_len, self.embed_dim) + + attn_output_v = self.out_v_proj(attn_output_v) + attn_output_l = self.out_l_proj(attn_output_l) + + return attn_output_v, attn_output_l + + +# Bi-Direction MHA (text->image, image->text) +class BiAttentionBlock(nn.Module): + def __init__( + self, + v_dim, + l_dim, + embed_dim, + num_heads, + dropout=0.1, + drop_path=0.0, + init_values=1e-4, + cfg=None, + ): + """ + Inputs: + embed_dim - Dimensionality of input and attention feature vectors + hidden_dim - Dimensionality of hidden layer in feed-forward network + (usually 2-4x larger than embed_dim) + num_heads - Number of heads to use in the Multi-Head Attention block + dropout - Amount of dropout to apply in the feed-forward network + """ + super(BiAttentionBlock, self).__init__() + + # pre layer norm + self.layer_norm_v = nn.LayerNorm(v_dim) + self.layer_norm_l = nn.LayerNorm(l_dim) + self.attn = BiMultiHeadAttention( + v_dim=v_dim, l_dim=l_dim, embed_dim=embed_dim, num_heads=num_heads, dropout=dropout + ) + + # add layer scale for training stability + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.gamma_v = nn.Parameter(init_values * torch.ones((v_dim)), requires_grad=True) + self.gamma_l = nn.Parameter(init_values * torch.ones((l_dim)), requires_grad=True) + + def forward(self, v, l, attention_mask_v=None, attention_mask_l=None): + v = self.layer_norm_v(v) + l = self.layer_norm_l(l) + delta_v, delta_l = self.attn( + v, l, attention_mask_v=attention_mask_v, attention_mask_l=attention_mask_l + ) + # v, l = v + delta_v, l + delta_l + v = v + self.drop_path(self.gamma_v * delta_v) + l = l + self.drop_path(self.gamma_l * delta_l) + return v, l + + # def forward(self, v:List[torch.Tensor], l, attention_mask_v=None, attention_mask_l=None) diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/groundingdino.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/groundingdino.py new file mode 100644 index 0000000000000000000000000000000000000000..134cadaac010a2b23fe407fa304c71a08d6da206 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/groundingdino.py @@ -0,0 +1,398 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR model and criterion classes. +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Modified from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ +# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR) +# Copyright (c) 2020 SenseTime. All Rights Reserved. +# ------------------------------------------------------------------------ +import copy +from typing import List + +import torch +import torch.nn.functional as F +from torch import nn +from torchvision.ops.boxes import nms +from transformers import AutoTokenizer, BertModel, BertTokenizer, RobertaModel, RobertaTokenizerFast + +from groundingdino.util import box_ops, get_tokenlizer +from groundingdino.util.misc import ( + NestedTensor, + accuracy, + get_world_size, + interpolate, + inverse_sigmoid, + is_dist_avail_and_initialized, + nested_tensor_from_tensor_list, +) +from groundingdino.util.utils import get_phrases_from_posmap +from groundingdino.util.visualizer import COCOVisualizer +from groundingdino.util.vl_utils import create_positive_map_from_span + +from ..registry import MODULE_BUILD_FUNCS +from .backbone import build_backbone +from .bertwarper import ( + BertModelWarper, + generate_masks_with_special_tokens, + generate_masks_with_special_tokens_and_transfer_map, +) +from .transformer import build_transformer +from .utils import MLP, ContrastiveEmbed, sigmoid_focal_loss + + +class GroundingDINO(nn.Module): + """This is the Cross-Attention Detector module that performs object detection""" + + def __init__( + self, + backbone, + transformer, + num_queries, + bert_base_uncased_path, + aux_loss=False, + iter_update=False, + query_dim=2, + num_feature_levels=1, + nheads=8, + # two stage + two_stage_type="no", # ['no', 'standard'] + dec_pred_bbox_embed_share=True, + two_stage_class_embed_share=True, + two_stage_bbox_embed_share=True, + num_patterns=0, + dn_number=100, + dn_box_noise_scale=0.4, + dn_label_noise_ratio=0.5, + dn_labelbook_size=100, + text_encoder_type="bert-base-uncased", + sub_sentence_present=True, + max_text_len=256, + ): + """Initializes the model. + Parameters: + backbone: torch module of the backbone to be used. See backbone.py + transformer: torch module of the transformer architecture. See transformer.py + num_queries: number of object queries, ie detection slot. This is the maximal number of objects + Conditional DETR can detect in a single image. For COCO, we recommend 100 queries. + aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. + """ + super().__init__() + self.num_queries = num_queries + self.transformer = transformer + self.hidden_dim = hidden_dim = transformer.d_model + self.num_feature_levels = num_feature_levels + self.nheads = nheads + self.max_text_len = 256 + self.sub_sentence_present = sub_sentence_present + + # setting query dim + self.query_dim = query_dim + assert query_dim == 4 + + # for dn training + self.num_patterns = num_patterns + self.dn_number = dn_number + self.dn_box_noise_scale = dn_box_noise_scale + self.dn_label_noise_ratio = dn_label_noise_ratio + self.dn_labelbook_size = dn_labelbook_size + + # bert + self.tokenizer = get_tokenlizer.get_tokenlizer(text_encoder_type, bert_base_uncased_path) + self.bert = get_tokenlizer.get_pretrained_language_model(text_encoder_type, bert_base_uncased_path) + self.bert.pooler.dense.weight.requires_grad_(False) + self.bert.pooler.dense.bias.requires_grad_(False) + self.bert = BertModelWarper(bert_model=self.bert) + + self.feat_map = nn.Linear(self.bert.config.hidden_size, self.hidden_dim, bias=True) + nn.init.constant_(self.feat_map.bias.data, 0) + nn.init.xavier_uniform_(self.feat_map.weight.data) + # freeze + + # special tokens + self.specical_tokens = self.tokenizer.convert_tokens_to_ids(["[CLS]", "[SEP]", ".", "?"]) + + # prepare input projection layers + if num_feature_levels > 1: + num_backbone_outs = len(backbone.num_channels) + input_proj_list = [] + for _ in range(num_backbone_outs): + in_channels = backbone.num_channels[_] + input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, hidden_dim, kernel_size=1), + nn.GroupNorm(32, hidden_dim), + ) + ) + for _ in range(num_feature_levels - num_backbone_outs): + input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1), + nn.GroupNorm(32, hidden_dim), + ) + ) + in_channels = hidden_dim + self.input_proj = nn.ModuleList(input_proj_list) + else: + assert two_stage_type == "no", "two_stage_type should be no if num_feature_levels=1 !!!" + self.input_proj = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1), + nn.GroupNorm(32, hidden_dim), + ) + ] + ) + + self.backbone = backbone + self.aux_loss = aux_loss + self.box_pred_damping = box_pred_damping = None + + self.iter_update = iter_update + assert iter_update, "Why not iter_update?" + + # prepare pred layers + self.dec_pred_bbox_embed_share = dec_pred_bbox_embed_share + # prepare class & box embed + _class_embed = ContrastiveEmbed() + + _bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3) + nn.init.constant_(_bbox_embed.layers[-1].weight.data, 0) + nn.init.constant_(_bbox_embed.layers[-1].bias.data, 0) + + if dec_pred_bbox_embed_share: + box_embed_layerlist = [_bbox_embed for i in range(transformer.num_decoder_layers)] + else: + box_embed_layerlist = [ + copy.deepcopy(_bbox_embed) for i in range(transformer.num_decoder_layers) + ] + class_embed_layerlist = [_class_embed for i in range(transformer.num_decoder_layers)] + self.bbox_embed = nn.ModuleList(box_embed_layerlist) + self.class_embed = nn.ModuleList(class_embed_layerlist) + self.transformer.decoder.bbox_embed = self.bbox_embed + self.transformer.decoder.class_embed = self.class_embed + + # two stage + self.two_stage_type = two_stage_type + assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format( + two_stage_type + ) + if two_stage_type != "no": + if two_stage_bbox_embed_share: + assert dec_pred_bbox_embed_share + self.transformer.enc_out_bbox_embed = _bbox_embed + else: + self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed) + + if two_stage_class_embed_share: + assert dec_pred_bbox_embed_share + self.transformer.enc_out_class_embed = _class_embed + else: + self.transformer.enc_out_class_embed = copy.deepcopy(_class_embed) + + self.refpoint_embed = None + + self._reset_parameters() + + def _reset_parameters(self): + # init input_proj + for proj in self.input_proj: + nn.init.xavier_uniform_(proj[0].weight, gain=1) + nn.init.constant_(proj[0].bias, 0) + + def init_ref_points(self, use_num_queries): + self.refpoint_embed = nn.Embedding(use_num_queries, self.query_dim) + + def forward(self, samples: NestedTensor, targets: List = None, **kw): + """The forward expects a NestedTensor, which consists of: + - samples.tensor: batched images, of shape [batch_size x 3 x H x W] + - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels + + It returns a dict with the following elements: + - "pred_logits": the classification logits (including no-object) for all queries. + Shape= [batch_size x num_queries x num_classes] + - "pred_boxes": The normalized boxes coordinates for all queries, represented as + (center_x, center_y, width, height). These values are normalized in [0, 1], + relative to the size of each individual image (disregarding possible padding). + See PostProcess for information on how to retrieve the unnormalized bounding box. + - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of + dictionnaries containing the two above keys for each decoder layer. + """ + if targets is None: + captions = kw["captions"] + else: + captions = [t["caption"] for t in targets] + len(captions) + + # encoder texts + tokenized = self.tokenizer(captions, padding="longest", return_tensors="pt").to( + samples.device + ) + ( + text_self_attention_masks, + position_ids, + cate_to_token_mask_list, + ) = generate_masks_with_special_tokens_and_transfer_map( + tokenized, self.specical_tokens, self.tokenizer + ) + + if text_self_attention_masks.shape[1] > self.max_text_len: + text_self_attention_masks = text_self_attention_masks[ + :, : self.max_text_len, : self.max_text_len + ] + position_ids = position_ids[:, : self.max_text_len] + tokenized["input_ids"] = tokenized["input_ids"][:, : self.max_text_len] + tokenized["attention_mask"] = tokenized["attention_mask"][:, : self.max_text_len] + tokenized["token_type_ids"] = tokenized["token_type_ids"][:, : self.max_text_len] + + # extract text embeddings + if self.sub_sentence_present: + tokenized_for_encoder = {k: v for k, v in tokenized.items() if k != "attention_mask"} + tokenized_for_encoder["attention_mask"] = text_self_attention_masks + tokenized_for_encoder["position_ids"] = position_ids + else: + # import ipdb; ipdb.set_trace() + tokenized_for_encoder = tokenized + + bert_output = self.bert(**tokenized_for_encoder) # bs, 195, 768 + + encoded_text = self.feat_map(bert_output["last_hidden_state"]) # bs, 195, d_model + text_token_mask = tokenized.attention_mask.bool() # bs, 195 + # text_token_mask: True for nomask, False for mask + # text_self_attention_masks: True for nomask, False for mask + + if encoded_text.shape[1] > self.max_text_len: + encoded_text = encoded_text[:, : self.max_text_len, :] + text_token_mask = text_token_mask[:, : self.max_text_len] + position_ids = position_ids[:, : self.max_text_len] + text_self_attention_masks = text_self_attention_masks[ + :, : self.max_text_len, : self.max_text_len + ] + + text_dict = { + "encoded_text": encoded_text, # bs, 195, d_model + "text_token_mask": text_token_mask, # bs, 195 + "position_ids": position_ids, # bs, 195 + "text_self_attention_masks": text_self_attention_masks, # bs, 195,195 + } + + # import ipdb; ipdb.set_trace() + + if isinstance(samples, (list, torch.Tensor)): + samples = nested_tensor_from_tensor_list(samples) + features, poss = self.backbone(samples) + + srcs = [] + masks = [] + for l, feat in enumerate(features): + src, mask = feat.decompose() + srcs.append(self.input_proj[l](src)) + masks.append(mask) + assert mask is not None + if self.num_feature_levels > len(srcs): + _len_srcs = len(srcs) + for l in range(_len_srcs, self.num_feature_levels): + if l == _len_srcs: + src = self.input_proj[l](features[-1].tensors) + else: + src = self.input_proj[l](srcs[-1]) + m = samples.mask + mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0] + pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype) + srcs.append(src) + masks.append(mask) + poss.append(pos_l) + + input_query_bbox = input_query_label = attn_mask = dn_meta = None + hs, reference, hs_enc, ref_enc, init_box_proposal = self.transformer( + srcs, masks, input_query_bbox, poss, input_query_label, attn_mask, text_dict + ) + + # deformable-detr-like anchor update + outputs_coord_list = [] + for dec_lid, (layer_ref_sig, layer_bbox_embed, layer_hs) in enumerate( + zip(reference[:-1], self.bbox_embed, hs) + ): + layer_delta_unsig = layer_bbox_embed(layer_hs) + layer_outputs_unsig = layer_delta_unsig + inverse_sigmoid(layer_ref_sig) + layer_outputs_unsig = layer_outputs_unsig.sigmoid() + outputs_coord_list.append(layer_outputs_unsig) + outputs_coord_list = torch.stack(outputs_coord_list) + + # output + outputs_class = torch.stack( + [ + layer_cls_embed(layer_hs, text_dict) + for layer_cls_embed, layer_hs in zip(self.class_embed, hs) + ] + ) + out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord_list[-1]} + + # # for intermediate outputs + # if self.aux_loss: + # out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord_list) + + # # for encoder output + # if hs_enc is not None: + # # prepare intermediate outputs + # interm_coord = ref_enc[-1] + # interm_class = self.transformer.enc_out_class_embed(hs_enc[-1], text_dict) + # out['interm_outputs'] = {'pred_logits': interm_class, 'pred_boxes': interm_coord} + # out['interm_outputs_for_matching_pre'] = {'pred_logits': interm_class, 'pred_boxes': init_box_proposal} + + return out + + @torch.jit.unused + def _set_aux_loss(self, outputs_class, outputs_coord): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [ + {"pred_logits": a, "pred_boxes": b} + for a, b in zip(outputs_class[:-1], outputs_coord[:-1]) + ] + + +@MODULE_BUILD_FUNCS.registe_with_name(module_name="groundingdino") +def build_groundingdino(args): + + backbone = build_backbone(args) + transformer = build_transformer(args) + + dn_labelbook_size = args.dn_labelbook_size + dec_pred_bbox_embed_share = args.dec_pred_bbox_embed_share + sub_sentence_present = args.sub_sentence_present + bert_base_uncased_path = args.bert_base_uncased_path if 'bert_base_uncased_path' in args else None + + model = GroundingDINO( + backbone, + transformer, + num_queries=args.num_queries, + bert_base_uncased_path=bert_base_uncased_path, + aux_loss=True, + iter_update=True, + query_dim=4, + num_feature_levels=args.num_feature_levels, + nheads=args.nheads, + dec_pred_bbox_embed_share=dec_pred_bbox_embed_share, + two_stage_type=args.two_stage_type, + two_stage_bbox_embed_share=args.two_stage_bbox_embed_share, + two_stage_class_embed_share=args.two_stage_class_embed_share, + num_patterns=args.num_patterns, + dn_number=0, + dn_box_noise_scale=args.dn_box_noise_scale, + dn_label_noise_ratio=args.dn_label_noise_ratio, + dn_labelbook_size=dn_labelbook_size, + text_encoder_type=args.text_encoder_type, + sub_sentence_present=sub_sentence_present, + max_text_len=args.max_text_len, + ) + + return model diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/ms_deform_attn.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/ms_deform_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..489d501bef364020212306d81e9b85c8daa27491 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/ms_deform_attn.py @@ -0,0 +1,413 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Deformable DETR +# Copyright (c) 2020 SenseTime. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------------------------------ +# Modified from: +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/functions/ms_deform_attn_func.py +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/modules/ms_deform_attn.py +# https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/multi_scale_deform_attn.py +# ------------------------------------------------------------------------------------------------ + +import math +import warnings +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Function +from torch.autograd.function import once_differentiable +from torch.nn.init import constant_, xavier_uniform_ + +try: + from groundingdino import _C +except: + warnings.warn("Failed to load custom C++ ops. Running on CPU mode Only!") + + +# helpers +def _is_power_of_2(n): + if (not isinstance(n, int)) or (n < 0): + raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n))) + return (n & (n - 1) == 0) and n != 0 + + +class MultiScaleDeformableAttnFunction(Function): + @staticmethod + def forward( + ctx, + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + im2col_step, + ): + ctx.im2col_step = im2col_step + output = _C.ms_deform_attn_forward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ctx.im2col_step, + ) + ctx.save_for_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ) + return output + + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + ( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ) = ctx.saved_tensors + grad_value, grad_sampling_loc, grad_attn_weight = _C.ms_deform_attn_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + grad_output, + ctx.im2col_step, + ) + + return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None + + +def multi_scale_deformable_attn_pytorch( + value: torch.Tensor, + value_spatial_shapes: torch.Tensor, + sampling_locations: torch.Tensor, + attention_weights: torch.Tensor, +) -> torch.Tensor: + + bs, _, num_heads, embed_dims = value.shape + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level, (H_, W_) in enumerate(value_spatial_shapes): + # bs, H_*W_, num_heads, embed_dims -> + # bs, H_*W_, num_heads*embed_dims -> + # bs, num_heads*embed_dims, H_*W_ -> + # bs*num_heads, embed_dims, H_, W_ + value_l_ = ( + value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_) + ) + # bs, num_queries, num_heads, num_points, 2 -> + # bs, num_heads, num_queries, num_points, 2 -> + # bs*num_heads, num_queries, num_points, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1) + # bs*num_heads, embed_dims, num_queries, num_points + sampling_value_l_ = F.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + # (bs, num_queries, num_heads, num_levels, num_points) -> + # (bs, num_heads, num_queries, num_levels, num_points) -> + # (bs, num_heads, 1, num_queries, num_levels*num_points) + attention_weights = attention_weights.transpose(1, 2).reshape( + bs * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(bs, num_heads * embed_dims, num_queries) + ) + return output.transpose(1, 2).contiguous() + + +class MultiScaleDeformableAttention(nn.Module): + """Multi-Scale Deformable Attention Module used in Deformable-DETR + + `Deformable DETR: Deformable Transformers for End-to-End Object Detection. + `_. + + Args: + embed_dim (int): The embedding dimension of Attention. Default: 256. + num_heads (int): The number of attention heads. Default: 8. + num_levels (int): The number of feature map used in Attention. Default: 4. + num_points (int): The number of sampling points for each query + in each head. Default: 4. + img2col_steps (int): The step used in image_to_column. Defualt: 64. + dropout (float): Dropout layer used in output. Default: 0.1. + batch_first (bool): if ``True``, then the input and output tensor will be + provided as `(bs, n, embed_dim)`. Default: False. `(n, bs, embed_dim)` + """ + + def __init__( + self, + embed_dim: int = 256, + num_heads: int = 8, + num_levels: int = 4, + num_points: int = 4, + img2col_step: int = 64, + batch_first: bool = False, + ): + super().__init__() + if embed_dim % num_heads != 0: + raise ValueError( + "embed_dim must be divisible by num_heads, but got {} and {}".format( + embed_dim, num_heads + ) + ) + head_dim = embed_dim // num_heads + + self.batch_first = batch_first + + if not _is_power_of_2(head_dim): + warnings.warn( + """ + You'd better set d_model in MSDeformAttn to make sure that + each dim of the attention head a power of 2, which is more efficient. + """ + ) + + self.im2col_step = img2col_step + self.embed_dim = embed_dim + self.num_heads = num_heads + self.num_levels = num_levels + self.num_points = num_points + self.sampling_offsets = nn.Linear(embed_dim, num_heads * num_levels * num_points * 2) + self.attention_weights = nn.Linear(embed_dim, num_heads * num_levels * num_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + + self.init_weights() + + def _reset_parameters(self): + return self.init_weights() + + def init_weights(self): + """ + Default initialization for Parameters of Module. + """ + constant_(self.sampling_offsets.weight.data, 0.0) + thetas = torch.arange(self.num_heads, dtype=torch.float32) * ( + 2.0 * math.pi / self.num_heads + ) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(self.num_heads, 1, 1, 2) + .repeat(1, self.num_levels, self.num_points, 1) + ) + for i in range(self.num_points): + grid_init[:, :, i, :] *= i + 1 + with torch.no_grad(): + self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1)) + constant_(self.attention_weights.weight.data, 0.0) + constant_(self.attention_weights.bias.data, 0.0) + xavier_uniform_(self.value_proj.weight.data) + constant_(self.value_proj.bias.data, 0.0) + xavier_uniform_(self.output_proj.weight.data) + constant_(self.output_proj.bias.data, 0.0) + + def freeze_sampling_offsets(self): + print("Freeze sampling offsets") + self.sampling_offsets.weight.requires_grad = False + self.sampling_offsets.bias.requires_grad = False + + def freeze_attention_weights(self): + print("Freeze attention weights") + self.attention_weights.weight.requires_grad = False + self.attention_weights.bias.requires_grad = False + + def forward( + self, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + value: Optional[torch.Tensor] = None, + query_pos: Optional[torch.Tensor] = None, + key_padding_mask: Optional[torch.Tensor] = None, + reference_points: Optional[torch.Tensor] = None, + spatial_shapes: Optional[torch.Tensor] = None, + level_start_index: Optional[torch.Tensor] = None, + **kwargs + ) -> torch.Tensor: + + """Forward Function of MultiScaleDeformableAttention + + Args: + query (torch.Tensor): Query embeddings with shape + `(num_query, bs, embed_dim)` + key (torch.Tensor): Key embeddings with shape + `(num_key, bs, embed_dim)` + value (torch.Tensor): Value embeddings with shape + `(num_key, bs, embed_dim)` + query_pos (torch.Tensor): The position embedding for `query`. Default: None. + key_padding_mask (torch.Tensor): ByteTensor for `query`, with shape `(bs, num_key)`, + indicating which elements within `key` to be ignored in attention. + reference_points (torch.Tensor): The normalized reference points + with shape `(bs, num_query, num_levels, 2)`, + all elements is range in [0, 1], top-left (0, 0), + bottom-right (1, 1), including padding are. + or `(N, Length_{query}, num_levels, 4)`, add additional + two dimensions `(h, w)` to form reference boxes. + spatial_shapes (torch.Tensor): Spatial shape of features in different levels. + With shape `(num_levels, 2)`, last dimension represents `(h, w)`. + level_start_index (torch.Tensor): The start index of each level. A tensor with + shape `(num_levels, )` which can be represented as + `[0, h_0 * w_0, h_0 * w_0 + h_1 * w_1, ...]`. + + Returns: + torch.Tensor: forward results with shape `(num_query, bs, embed_dim)` + """ + + if value is None: + value = query + + if query_pos is not None: + query = query + query_pos + + if not self.batch_first: + # change to (bs, num_query ,embed_dims) + query = query.permute(1, 0, 2) + value = value.permute(1, 0, 2) + + bs, num_query, _ = query.shape + bs, num_value, _ = value.shape + + assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value + + value = self.value_proj(value) + if key_padding_mask is not None: + value = value.masked_fill(key_padding_mask[..., None], float(0)) + value = value.view(bs, num_value, self.num_heads, -1) + sampling_offsets = self.sampling_offsets(query).view( + bs, num_query, self.num_heads, self.num_levels, self.num_points, 2 + ) + attention_weights = self.attention_weights(query).view( + bs, num_query, self.num_heads, self.num_levels * self.num_points + ) + attention_weights = attention_weights.softmax(-1) + attention_weights = attention_weights.view( + bs, + num_query, + self.num_heads, + self.num_levels, + self.num_points, + ) + + # bs, num_query, num_heads, num_levels, num_points, 2 + if reference_points.shape[-1] == 2: + offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif reference_points.shape[-1] == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets + / self.num_points + * reference_points[:, :, None, :, None, 2:] + * 0.5 + ) + else: + raise ValueError( + "Last dim of reference_points must be 2 or 4, but get {} instead.".format( + reference_points.shape[-1] + ) + ) + + if torch.cuda.is_available() and value.is_cuda: + halffloat = False + if value.dtype == torch.float16: + halffloat = True + value = value.float() + sampling_locations = sampling_locations.float() + attention_weights = attention_weights.float() + + output = MultiScaleDeformableAttnFunction.apply( + value, + spatial_shapes, + level_start_index, + sampling_locations, + attention_weights, + self.im2col_step, + ) + + if halffloat: + output = output.half() + else: + output = multi_scale_deformable_attn_pytorch( + value, spatial_shapes, sampling_locations, attention_weights + ) + + output = self.output_proj(output) + + if not self.batch_first: + output = output.permute(1, 0, 2) + + return output + + +def create_dummy_class(klass, dependency, message=""): + """ + When a dependency of a class is not available, create a dummy class which throws ImportError + when used. + + Args: + klass (str): name of the class. + dependency (str): name of the dependency. + message: extra message to print + Returns: + class: a class object + """ + err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, klass) + if message: + err = err + " " + message + + class _DummyMetaClass(type): + # throw error on class attribute access + def __getattr__(_, __): # noqa: B902 + raise ImportError(err) + + class _Dummy(object, metaclass=_DummyMetaClass): + # throw error on constructor + def __init__(self, *args, **kwargs): + raise ImportError(err) + + return _Dummy + + +def create_dummy_func(func, dependency, message=""): + """ + When a dependency of a function is not available, create a dummy function which throws + ImportError when used. + + Args: + func (str): name of the function. + dependency (str or list[str]): name(s) of the dependency. + message: extra message to print + Returns: + function: a function object + """ + err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, func) + if message: + err = err + " " + message + + if isinstance(dependency, (list, tuple)): + dependency = ",".join(dependency) + + def _dummy(*args, **kwargs): + raise ImportError(err) + + return _dummy diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..d554215ecfaa7ad5a7661fa50757e5de713f0b32 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer.py @@ -0,0 +1,960 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR Transformer class. +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Modified from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +from typing import Optional + +import torch +import torch.utils.checkpoint as checkpoint +from torch import Tensor, nn + +from groundingdino.util.misc import inverse_sigmoid + +from .fuse_modules import BiAttentionBlock +from .ms_deform_attn import MultiScaleDeformableAttention as MSDeformAttn +from .transformer_vanilla import TransformerEncoderLayer +from .utils import ( + MLP, + _get_activation_fn, + _get_clones, + gen_encoder_output_proposals, + gen_sineembed_for_position, + get_sine_pos_embed, +) + + +class Transformer(nn.Module): + def __init__( + self, + d_model=256, + nhead=8, + num_queries=300, + num_encoder_layers=6, + num_unicoder_layers=0, + num_decoder_layers=6, + dim_feedforward=2048, + dropout=0.0, + activation="relu", + normalize_before=False, + return_intermediate_dec=False, + query_dim=4, + num_patterns=0, + # for deformable encoder + num_feature_levels=1, + enc_n_points=4, + dec_n_points=4, + # init query + learnable_tgt_init=False, + # two stage + two_stage_type="no", # ['no', 'standard', 'early', 'combine', 'enceachlayer', 'enclayer1'] + embed_init_tgt=False, + # for text + use_text_enhancer=False, + use_fusion_layer=False, + use_checkpoint=False, + use_transformer_ckpt=False, + use_text_cross_attention=False, + text_dropout=0.1, + fusion_dropout=0.1, + fusion_droppath=0.0, + ): + super().__init__() + self.num_feature_levels = num_feature_levels + self.num_encoder_layers = num_encoder_layers + self.num_unicoder_layers = num_unicoder_layers + self.num_decoder_layers = num_decoder_layers + self.num_queries = num_queries + assert query_dim == 4 + + # choose encoder layer type + encoder_layer = DeformableTransformerEncoderLayer( + d_model, dim_feedforward, dropout, activation, num_feature_levels, nhead, enc_n_points + ) + + if use_text_enhancer: + text_enhance_layer = TransformerEncoderLayer( + d_model=d_model, + nhead=nhead // 2, + dim_feedforward=dim_feedforward // 2, + dropout=text_dropout, + ) + else: + text_enhance_layer = None + + if use_fusion_layer: + feature_fusion_layer = BiAttentionBlock( + v_dim=d_model, + l_dim=d_model, + embed_dim=dim_feedforward // 2, + num_heads=nhead // 2, + dropout=fusion_dropout, + drop_path=fusion_droppath, + ) + else: + feature_fusion_layer = None + + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + assert encoder_norm is None + self.encoder = TransformerEncoder( + encoder_layer, + num_encoder_layers, + d_model=d_model, + num_queries=num_queries, + text_enhance_layer=text_enhance_layer, + feature_fusion_layer=feature_fusion_layer, + use_checkpoint=use_checkpoint, + use_transformer_ckpt=use_transformer_ckpt, + ) + + # choose decoder layer type + decoder_layer = DeformableTransformerDecoderLayer( + d_model, + dim_feedforward, + dropout, + activation, + num_feature_levels, + nhead, + dec_n_points, + use_text_cross_attention=use_text_cross_attention, + ) + + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder( + decoder_layer, + num_decoder_layers, + decoder_norm, + return_intermediate=return_intermediate_dec, + d_model=d_model, + query_dim=query_dim, + num_feature_levels=num_feature_levels, + ) + + self.d_model = d_model + self.nhead = nhead + self.dec_layers = num_decoder_layers + self.num_queries = num_queries # useful for single stage model only + self.num_patterns = num_patterns + if not isinstance(num_patterns, int): + Warning("num_patterns should be int but {}".format(type(num_patterns))) + self.num_patterns = 0 + + if num_feature_levels > 1: + if self.num_encoder_layers > 0: + self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model)) + else: + self.level_embed = None + + self.learnable_tgt_init = learnable_tgt_init + assert learnable_tgt_init, "why not learnable_tgt_init" + self.embed_init_tgt = embed_init_tgt + if (two_stage_type != "no" and embed_init_tgt) or (two_stage_type == "no"): + self.tgt_embed = nn.Embedding(self.num_queries, d_model) + nn.init.normal_(self.tgt_embed.weight.data) + else: + self.tgt_embed = None + + # for two stage + self.two_stage_type = two_stage_type + assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format( + two_stage_type + ) + if two_stage_type == "standard": + # anchor selection at the output of encoder + self.enc_output = nn.Linear(d_model, d_model) + self.enc_output_norm = nn.LayerNorm(d_model) + self.two_stage_wh_embedding = None + + if two_stage_type == "no": + self.init_ref_points(num_queries) # init self.refpoint_embed + + self.enc_out_class_embed = None + self.enc_out_bbox_embed = None + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + for m in self.modules(): + if isinstance(m, MSDeformAttn): + m._reset_parameters() + if self.num_feature_levels > 1 and self.level_embed is not None: + nn.init.normal_(self.level_embed) + + def get_valid_ratio(self, mask): + _, H, W = mask.shape + valid_H = torch.sum(~mask[:, :, 0], 1) + valid_W = torch.sum(~mask[:, 0, :], 1) + valid_ratio_h = valid_H.float() / H + valid_ratio_w = valid_W.float() / W + valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1) + return valid_ratio + + def init_ref_points(self, use_num_queries): + self.refpoint_embed = nn.Embedding(use_num_queries, 4) + + def forward(self, srcs, masks, refpoint_embed, pos_embeds, tgt, attn_mask=None, text_dict=None): + """ + Input: + - srcs: List of multi features [bs, ci, hi, wi] + - masks: List of multi masks [bs, hi, wi] + - refpoint_embed: [bs, num_dn, 4]. None in infer + - pos_embeds: List of multi pos embeds [bs, ci, hi, wi] + - tgt: [bs, num_dn, d_model]. None in infer + + """ + # prepare input for encoder + src_flatten = [] + mask_flatten = [] + lvl_pos_embed_flatten = [] + spatial_shapes = [] + for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)): + bs, c, h, w = src.shape + spatial_shape = (h, w) + spatial_shapes.append(spatial_shape) + + src = src.flatten(2).transpose(1, 2) # bs, hw, c + mask = mask.flatten(1) # bs, hw + pos_embed = pos_embed.flatten(2).transpose(1, 2) # bs, hw, c + if self.num_feature_levels > 1 and self.level_embed is not None: + lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1) + else: + lvl_pos_embed = pos_embed + lvl_pos_embed_flatten.append(lvl_pos_embed) + src_flatten.append(src) + mask_flatten.append(mask) + src_flatten = torch.cat(src_flatten, 1) # bs, \sum{hxw}, c + mask_flatten = torch.cat(mask_flatten, 1) # bs, \sum{hxw} + lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) # bs, \sum{hxw}, c + spatial_shapes = torch.as_tensor( + spatial_shapes, dtype=torch.long, device=src_flatten.device + ) + level_start_index = torch.cat( + (spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]) + ) + valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1).to(src.dtype) + + # two stage + enc_topk_proposals = enc_refpoint_embed = None + + ######################################################### + # Begin Encoder + ######################################################### + memory, memory_text = self.encoder( + src_flatten, + pos=lvl_pos_embed_flatten, + level_start_index=level_start_index, + spatial_shapes=spatial_shapes, + valid_ratios=valid_ratios, + key_padding_mask=mask_flatten, + memory_text=text_dict["encoded_text"], + text_attention_mask=~text_dict["text_token_mask"], + # we ~ the mask . False means use the token; True means pad the token + position_ids=text_dict["position_ids"], + text_self_attention_masks=text_dict["text_self_attention_masks"], + ) + ######################################################### + # End Encoder + # - memory: bs, \sum{hw}, c + # - mask_flatten: bs, \sum{hw} + # - lvl_pos_embed_flatten: bs, \sum{hw}, c + # - enc_intermediate_output: None or (nenc+1, bs, nq, c) or (nenc, bs, nq, c) + # - enc_intermediate_refpoints: None or (nenc+1, bs, nq, c) or (nenc, bs, nq, c) + ######################################################### + text_dict["encoded_text"] = memory_text + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # if memory.isnan().any() | memory.isinf().any(): + # import ipdb; ipdb.set_trace() + + if self.two_stage_type == "standard": + output_memory, output_proposals = gen_encoder_output_proposals( + memory, mask_flatten, spatial_shapes + ) + output_memory = self.enc_output_norm(self.enc_output(output_memory)) + + if text_dict is not None: + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory, text_dict) + else: + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory) + + topk_logits = enc_outputs_class_unselected.max(-1)[0] + enc_outputs_coord_unselected = ( + self.enc_out_bbox_embed(output_memory) + output_proposals + ) # (bs, \sum{hw}, 4) unsigmoid + topk = self.num_queries + + topk_proposals = torch.topk(topk_logits, topk, dim=1)[1] # bs, nq + + # gather boxes + refpoint_embed_undetach = torch.gather( + enc_outputs_coord_unselected, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4) + ) # unsigmoid + refpoint_embed_ = refpoint_embed_undetach.detach() + init_box_proposal = torch.gather( + output_proposals, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4) + ).sigmoid() # sigmoid + + # gather tgt + tgt_undetach = torch.gather( + output_memory, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, self.d_model) + ) + if self.embed_init_tgt: + tgt_ = ( + self.tgt_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, d_model + else: + tgt_ = tgt_undetach.detach() + + if refpoint_embed is not None: + refpoint_embed = torch.cat([refpoint_embed, refpoint_embed_], dim=1) + tgt = torch.cat([tgt, tgt_], dim=1) + else: + refpoint_embed, tgt = refpoint_embed_, tgt_ + + elif self.two_stage_type == "no": + tgt_ = ( + self.tgt_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, d_model + refpoint_embed_ = ( + self.refpoint_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, 4 + + if refpoint_embed is not None: + refpoint_embed = torch.cat([refpoint_embed, refpoint_embed_], dim=1) + tgt = torch.cat([tgt, tgt_], dim=1) + else: + refpoint_embed, tgt = refpoint_embed_, tgt_ + + if self.num_patterns > 0: + tgt_embed = tgt.repeat(1, self.num_patterns, 1) + refpoint_embed = refpoint_embed.repeat(1, self.num_patterns, 1) + tgt_pat = self.patterns.weight[None, :, :].repeat_interleave( + self.num_queries, 1 + ) # 1, n_q*n_pat, d_model + tgt = tgt_embed + tgt_pat + + init_box_proposal = refpoint_embed_.sigmoid() + + else: + raise NotImplementedError("unknown two_stage_type {}".format(self.two_stage_type)) + ######################################################### + # End preparing tgt + # - tgt: bs, NQ, d_model + # - refpoint_embed(unsigmoid): bs, NQ, d_model + ######################################################### + + ######################################################### + # Begin Decoder + ######################################################### + hs, references = self.decoder( + tgt=tgt.transpose(0, 1), + memory=memory.transpose(0, 1), + memory_key_padding_mask=mask_flatten, + pos=lvl_pos_embed_flatten.transpose(0, 1), + refpoints_unsigmoid=refpoint_embed.transpose(0, 1), + level_start_index=level_start_index, + spatial_shapes=spatial_shapes, + valid_ratios=valid_ratios, + tgt_mask=attn_mask, + memory_text=text_dict["encoded_text"], + text_attention_mask=~text_dict["text_token_mask"], + # we ~ the mask . False means use the token; True means pad the token + ) + ######################################################### + # End Decoder + # hs: n_dec, bs, nq, d_model + # references: n_dec+1, bs, nq, query_dim + ######################################################### + + ######################################################### + # Begin postprocess + ######################################################### + if self.two_stage_type == "standard": + hs_enc = tgt_undetach.unsqueeze(0) + ref_enc = refpoint_embed_undetach.sigmoid().unsqueeze(0) + else: + hs_enc = ref_enc = None + ######################################################### + # End postprocess + # hs_enc: (n_enc+1, bs, nq, d_model) or (1, bs, nq, d_model) or (n_enc, bs, nq, d_model) or None + # ref_enc: (n_enc+1, bs, nq, query_dim) or (1, bs, nq, query_dim) or (n_enc, bs, nq, d_model) or None + ######################################################### + + return hs, references, hs_enc, ref_enc, init_box_proposal + # hs: (n_dec, bs, nq, d_model) + # references: sigmoid coordinates. (n_dec+1, bs, bq, 4) + # hs_enc: (n_enc+1, bs, nq, d_model) or (1, bs, nq, d_model) or None + # ref_enc: sigmoid coordinates. \ + # (n_enc+1, bs, nq, query_dim) or (1, bs, nq, query_dim) or None + + +class TransformerEncoder(nn.Module): + def __init__( + self, + encoder_layer, + num_layers, + d_model=256, + num_queries=300, + enc_layer_share=False, + text_enhance_layer=None, + feature_fusion_layer=None, + use_checkpoint=False, + use_transformer_ckpt=False, + ): + """_summary_ + + Args: + encoder_layer (_type_): _description_ + num_layers (_type_): _description_ + norm (_type_, optional): _description_. Defaults to None. + d_model (int, optional): _description_. Defaults to 256. + num_queries (int, optional): _description_. Defaults to 300. + enc_layer_share (bool, optional): _description_. Defaults to False. + + """ + super().__init__() + # prepare layers + self.layers = [] + self.text_layers = [] + self.fusion_layers = [] + if num_layers > 0: + self.layers = _get_clones(encoder_layer, num_layers, layer_share=enc_layer_share) + + if text_enhance_layer is not None: + self.text_layers = _get_clones( + text_enhance_layer, num_layers, layer_share=enc_layer_share + ) + if feature_fusion_layer is not None: + self.fusion_layers = _get_clones( + feature_fusion_layer, num_layers, layer_share=enc_layer_share + ) + else: + self.layers = [] + del encoder_layer + + if text_enhance_layer is not None: + self.text_layers = [] + del text_enhance_layer + if feature_fusion_layer is not None: + self.fusion_layers = [] + del feature_fusion_layer + + self.query_scale = None + self.num_queries = num_queries + self.num_layers = num_layers + self.d_model = d_model + + self.use_checkpoint = use_checkpoint + self.use_transformer_ckpt = use_transformer_ckpt + + @staticmethod + def get_reference_points(spatial_shapes, valid_ratios, device): + reference_points_list = [] + for lvl, (H_, W_) in enumerate(spatial_shapes): + + ref_y, ref_x = torch.meshgrid( + torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device), + torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device), + ) + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + return reference_points + + def forward( + self, + # for images + src: Tensor, + pos: Tensor, + spatial_shapes: Tensor, + level_start_index: Tensor, + valid_ratios: Tensor, + key_padding_mask: Tensor, + # for texts + memory_text: Tensor = None, + text_attention_mask: Tensor = None, + pos_text: Tensor = None, + text_self_attention_masks: Tensor = None, + position_ids: Tensor = None, + ): + """ + Input: + - src: [bs, sum(hi*wi), 256] + - pos: pos embed for src. [bs, sum(hi*wi), 256] + - spatial_shapes: h,w of each level [num_level, 2] + - level_start_index: [num_level] start point of level in sum(hi*wi). + - valid_ratios: [bs, num_level, 2] + - key_padding_mask: [bs, sum(hi*wi)] + + - memory_text: bs, n_text, 256 + - text_attention_mask: bs, n_text + False for no padding; True for padding + - pos_text: bs, n_text, 256 + + - position_ids: bs, n_text + Intermedia: + - reference_points: [bs, sum(hi*wi), num_level, 2] + Outpus: + - output: [bs, sum(hi*wi), 256] + """ + + output = src + + # preparation and reshape + if self.num_layers > 0: + reference_points = self.get_reference_points( + spatial_shapes, valid_ratios, device=src.device + ) + + if self.text_layers: + # generate pos_text + bs, n_text, text_dim = memory_text.shape + if pos_text is None and position_ids is None: + pos_text = ( + torch.arange(n_text, device=memory_text.device) + .float() + .unsqueeze(0) + .unsqueeze(-1) + .repeat(bs, 1, 1) + ) + pos_text = get_sine_pos_embed(pos_text, num_pos_feats=256, exchange_xy=False) + if position_ids is not None: + pos_text = get_sine_pos_embed( + position_ids[..., None], num_pos_feats=256, exchange_xy=False + ) + pos_text = pos_text.to(src.dtype) + + # main process + for layer_id, layer in enumerate(self.layers): + # if output.isnan().any() or memory_text.isnan().any(): + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + if self.fusion_layers: + if self.use_checkpoint: + output, memory_text = checkpoint.checkpoint( + self.fusion_layers[layer_id], + output, + memory_text, + key_padding_mask, + text_attention_mask, + ) + else: + output, memory_text = self.fusion_layers[layer_id]( + v=output, + l=memory_text, + attention_mask_v=key_padding_mask, + attention_mask_l=text_attention_mask, + ) + + if self.text_layers: + memory_text = self.text_layers[layer_id]( + src=memory_text.transpose(0, 1), + src_mask=~text_self_attention_masks, # note we use ~ for mask here + src_key_padding_mask=text_attention_mask, + pos=(pos_text.transpose(0, 1) if pos_text is not None else None), + ).transpose(0, 1) + + # main process + if self.use_transformer_ckpt: + output = checkpoint.checkpoint( + layer, + output, + pos, + reference_points, + spatial_shapes, + level_start_index, + key_padding_mask, + ) + else: + output = layer( + src=output, + pos=pos, + reference_points=reference_points, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + key_padding_mask=key_padding_mask, + ) + + return output, memory_text + + +class TransformerDecoder(nn.Module): + def __init__( + self, + decoder_layer, + num_layers, + norm=None, + return_intermediate=False, + d_model=256, + query_dim=4, + num_feature_levels=1, + ): + super().__init__() + if num_layers > 0: + self.layers = _get_clones(decoder_layer, num_layers) + else: + self.layers = [] + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + assert return_intermediate, "support return_intermediate only" + self.query_dim = query_dim + assert query_dim in [2, 4], "query_dim should be 2/4 but {}".format(query_dim) + self.num_feature_levels = num_feature_levels + + self.ref_point_head = MLP(query_dim // 2 * d_model, d_model, d_model, 2) + self.query_pos_sine_scale = None + + self.query_scale = None + self.bbox_embed = None + self.class_embed = None + + self.d_model = d_model + + self.ref_anchor_head = None + + def forward( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + refpoints_unsigmoid: Optional[Tensor] = None, # num_queries, bs, 2 + # for memory + level_start_index: Optional[Tensor] = None, # num_levels + spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + valid_ratios: Optional[Tensor] = None, + # for text + memory_text: Optional[Tensor] = None, + text_attention_mask: Optional[Tensor] = None, + ): + """ + Input: + - tgt: nq, bs, d_model + - memory: hw, bs, d_model + - pos: hw, bs, d_model + - refpoints_unsigmoid: nq, bs, 2/4 + - valid_ratios/spatial_shapes: bs, nlevel, 2 + """ + output = tgt + + intermediate = [] + reference_points = refpoints_unsigmoid.sigmoid() + ref_points = [reference_points] + + for layer_id, layer in enumerate(self.layers): + + if reference_points.shape[-1] == 4: + reference_points_input = ( + reference_points[:, :, None] + * torch.cat([valid_ratios, valid_ratios], -1)[None, :] + ) # nq, bs, nlevel, 4 + else: + assert reference_points.shape[-1] == 2 + reference_points_input = reference_points[:, :, None] * valid_ratios[None, :] + query_sine_embed = gen_sineembed_for_position( + reference_points_input[:, :, 0, :] + ) # nq, bs, 256*2 + + # conditional query + raw_query_pos = self.ref_point_head(query_sine_embed) # nq, bs, 256 + pos_scale = self.query_scale(output) if self.query_scale is not None else 1 + query_pos = pos_scale * raw_query_pos + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # if query_pos.isnan().any() | query_pos.isinf().any(): + # import ipdb; ipdb.set_trace() + + # main process + output = layer( + tgt=output, + tgt_query_pos=query_pos, + tgt_query_sine_embed=query_sine_embed, + tgt_key_padding_mask=tgt_key_padding_mask, + tgt_reference_points=reference_points_input, + memory_text=memory_text, + text_attention_mask=text_attention_mask, + memory=memory, + memory_key_padding_mask=memory_key_padding_mask, + memory_level_start_index=level_start_index, + memory_spatial_shapes=spatial_shapes, + memory_pos=pos, + self_attn_mask=tgt_mask, + cross_attn_mask=memory_mask, + ) + if output.isnan().any() | output.isinf().any(): + print(f"output layer_id {layer_id} is nan") + try: + num_nan = output.isnan().sum().item() + num_inf = output.isinf().sum().item() + print(f"num_nan {num_nan}, num_inf {num_inf}") + except Exception as e: + print(e) + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # import ipdb; ipdb.set_trace() + + # iter update + if self.bbox_embed is not None: + # box_holder = self.bbox_embed(output) + # box_holder[..., :self.query_dim] += inverse_sigmoid(reference_points) + # new_reference_points = box_holder[..., :self.query_dim].sigmoid() + + reference_before_sigmoid = inverse_sigmoid(reference_points) + delta_unsig = self.bbox_embed[layer_id](output) + outputs_unsig = delta_unsig + reference_before_sigmoid + new_reference_points = outputs_unsig.sigmoid() + + reference_points = new_reference_points.detach() + # if layer_id != self.num_layers - 1: + ref_points.append(new_reference_points) + + intermediate.append(self.norm(output)) + + return [ + [itm_out.transpose(0, 1) for itm_out in intermediate], + [itm_refpoint.transpose(0, 1) for itm_refpoint in ref_points], + ] + + +class DeformableTransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + ): + super().__init__() + + # self attention + self.self_attn = MSDeformAttn( + embed_dim=d_model, + num_levels=n_levels, + num_heads=n_heads, + num_points=n_points, + batch_first=True, + ) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn) + self.dropout2 = nn.Dropout(dropout) + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout3 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(d_model) + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, src): + src2 = self.linear2(self.dropout2(self.activation(self.linear1(src)))) + src = src + self.dropout3(src2) + src = self.norm2(src) + return src + + def forward( + self, src, pos, reference_points, spatial_shapes, level_start_index, key_padding_mask=None + ): + # self attention + # import ipdb; ipdb.set_trace() + src2 = self.self_attn( + query=self.with_pos_embed(src, pos), + reference_points=reference_points, + value=src, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + key_padding_mask=key_padding_mask, + ) + src = src + self.dropout1(src2) + src = self.norm1(src) + + # ffn + src = self.forward_ffn(src) + + return src + + +class DeformableTransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + use_text_feat_guide=False, + use_text_cross_attention=False, + ): + super().__init__() + + # cross attention + self.cross_attn = MSDeformAttn( + embed_dim=d_model, + num_levels=n_levels, + num_heads=n_heads, + num_points=n_points, + batch_first=True, + ) + self.dropout1 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm1 = nn.LayerNorm(d_model) + + # cross attention text + if use_text_cross_attention: + self.ca_text = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.catext_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.catext_norm = nn.LayerNorm(d_model) + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.dropout2 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm2 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn, batch_dim=1) + self.dropout3 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout4 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm3 = nn.LayerNorm(d_model) + + self.key_aware_proj = None + self.use_text_feat_guide = use_text_feat_guide + assert not use_text_feat_guide + self.use_text_cross_attention = use_text_cross_attention + + def rm_self_attn_modules(self): + self.self_attn = None + self.dropout2 = None + self.norm2 = None + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, tgt): + with torch.cuda.amp.autocast(enabled=False): + tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout4(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward( + self, + # for tgt + tgt: Optional[Tensor], # nq, bs, d_model + tgt_query_pos: Optional[Tensor] = None, # pos for query. MLP(Sine(pos)) + tgt_query_sine_embed: Optional[Tensor] = None, # pos for query. Sine(pos) + tgt_key_padding_mask: Optional[Tensor] = None, + tgt_reference_points: Optional[Tensor] = None, # nq, bs, 4 + memory_text: Optional[Tensor] = None, # bs, num_token, d_model + text_attention_mask: Optional[Tensor] = None, # bs, num_token + # for memory + memory: Optional[Tensor] = None, # hw, bs, d_model + memory_key_padding_mask: Optional[Tensor] = None, + memory_level_start_index: Optional[Tensor] = None, # num_levels + memory_spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + memory_pos: Optional[Tensor] = None, # pos for memory + # sa + self_attn_mask: Optional[Tensor] = None, # mask used for self-attention + cross_attn_mask: Optional[Tensor] = None, # mask used for cross-attention + ): + """ + Input: + - tgt/tgt_query_pos: nq, bs, d_model + - + """ + assert cross_attn_mask is None + + # self attention + if self.self_attn is not None: + # import ipdb; ipdb.set_trace() + q = k = self.with_pos_embed(tgt, tgt_query_pos) + tgt2 = self.self_attn(q, k, tgt, attn_mask=self_attn_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + + if self.use_text_cross_attention: + tgt2 = self.ca_text( + self.with_pos_embed(tgt, tgt_query_pos), + memory_text.transpose(0, 1), + memory_text.transpose(0, 1), + key_padding_mask=text_attention_mask, + )[0] + tgt = tgt + self.catext_dropout(tgt2) + tgt = self.catext_norm(tgt) + + tgt2 = self.cross_attn( + query=self.with_pos_embed(tgt, tgt_query_pos).transpose(0, 1), + reference_points=tgt_reference_points.transpose(0, 1).contiguous(), + value=memory.transpose(0, 1), + spatial_shapes=memory_spatial_shapes, + level_start_index=memory_level_start_index, + key_padding_mask=memory_key_padding_mask, + ).transpose(0, 1) + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + + # ffn + tgt = self.forward_ffn(tgt) + + return tgt + + +def build_transformer(args): + return Transformer( + d_model=args.hidden_dim, + dropout=args.dropout, + nhead=args.nheads, + num_queries=args.num_queries, + dim_feedforward=args.dim_feedforward, + num_encoder_layers=args.enc_layers, + num_decoder_layers=args.dec_layers, + normalize_before=args.pre_norm, + return_intermediate_dec=True, + query_dim=args.query_dim, + activation=args.transformer_activation, + num_patterns=args.num_patterns, + num_feature_levels=args.num_feature_levels, + enc_n_points=args.enc_n_points, + dec_n_points=args.dec_n_points, + learnable_tgt_init=True, + # two stage + two_stage_type=args.two_stage_type, # ['no', 'standard', 'early'] + embed_init_tgt=args.embed_init_tgt, + use_text_enhancer=args.use_text_enhancer, + use_fusion_layer=args.use_fusion_layer, + use_checkpoint=args.use_checkpoint, + use_transformer_ckpt=args.use_transformer_ckpt, + use_text_cross_attention=args.use_text_cross_attention, + text_dropout=args.text_dropout, + fusion_dropout=args.fusion_dropout, + fusion_droppath=args.fusion_droppath, + ) diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer_vanilla.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer_vanilla.py new file mode 100644 index 0000000000000000000000000000000000000000..10c0920c1a217af5bb3e1b13077568035ab3b7b5 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer_vanilla.py @@ -0,0 +1,123 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copyright (c) Aishwarya Kamath & Nicolas Carion. Licensed under the Apache License 2.0. All Rights Reserved +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +DETR Transformer class. + +Copy-paste from torch.nn.Transformer with modifications: + * positional encodings are passed in MHattention + * extra LN at the end of encoder is removed + * decoder returns a stack of activations from all decoding layers +""" +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from .utils import ( + MLP, + _get_activation_fn, + _get_clones, + gen_encoder_output_proposals, + gen_sineembed_for_position, + sigmoid_focal_loss, +) + + +class TextTransformer(nn.Module): + def __init__(self, num_layers, d_model=256, nheads=8, dim_feedforward=2048, dropout=0.1): + super().__init__() + self.num_layers = num_layers + self.d_model = d_model + self.nheads = nheads + self.dim_feedforward = dim_feedforward + self.norm = None + + single_encoder_layer = TransformerEncoderLayer( + d_model=d_model, nhead=nheads, dim_feedforward=dim_feedforward, dropout=dropout + ) + self.layers = _get_clones(single_encoder_layer, num_layers) + + def forward(self, memory_text: torch.Tensor, text_attention_mask: torch.Tensor): + """ + + Args: + text_attention_mask: bs, num_token + memory_text: bs, num_token, d_model + + Raises: + RuntimeError: _description_ + + Returns: + output: bs, num_token, d_model + """ + + output = memory_text.transpose(0, 1) + + for layer in self.layers: + output = layer(output, src_key_padding_mask=text_attention_mask) + + if self.norm is not None: + output = self.norm(output) + + return output.transpose(0, 1) + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + self.nhead = nhead + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward( + self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + # repeat attn mask + if src_mask.dim() == 3 and src_mask.shape[0] == src.shape[1]: + # bs, num_q, num_k + src_mask = src_mask.repeat(self.nhead, 1, 1) + + q = k = self.with_pos_embed(src, pos) + + src2 = self.self_attn(q, k, value=src, attn_mask=src_mask)[0] + + # src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/utils.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..27da9bbd88843598238467951c8339d5f92c95a4 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/GroundingDINO/utils.py @@ -0,0 +1,270 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +import copy +import math + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +def _get_clones(module, N, layer_share=False): + # import ipdb; ipdb.set_trace() + if layer_share: + return nn.ModuleList([module for i in range(N)]) + else: + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def get_sine_pos_embed( + pos_tensor: torch.Tensor, + num_pos_feats: int = 128, + temperature: int = 10000, + exchange_xy: bool = True, +): + """generate sine position embedding from a position tensor + Args: + pos_tensor (torch.Tensor): shape: [..., n]. + num_pos_feats (int): projected shape for each float in the tensor. + temperature (int): temperature in the sine/cosine function. + exchange_xy (bool, optional): exchange pos x and pos y. \ + For example, input tensor is [x,y], the results will be [pos(y), pos(x)]. Defaults to True. + Returns: + pos_embed (torch.Tensor): shape: [..., n*num_pos_feats]. + """ + scale = 2 * math.pi + dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos_tensor.device) + dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats) + + def sine_func(x: torch.Tensor): + sin_x = x * scale / dim_t + sin_x = torch.stack((sin_x[..., 0::2].sin(), sin_x[..., 1::2].cos()), dim=3).flatten(2) + return sin_x + + pos_res = [sine_func(x) for x in pos_tensor.split([1] * pos_tensor.shape[-1], dim=-1)] + if exchange_xy: + pos_res[0], pos_res[1] = pos_res[1], pos_res[0] + pos_res = torch.cat(pos_res, dim=-1) + return pos_res + + +def gen_encoder_output_proposals( + memory: Tensor, memory_padding_mask: Tensor, spatial_shapes: Tensor, learnedwh=None +): + """ + Input: + - memory: bs, \sum{hw}, d_model + - memory_padding_mask: bs, \sum{hw} + - spatial_shapes: nlevel, 2 + - learnedwh: 2 + Output: + - output_memory: bs, \sum{hw}, d_model + - output_proposals: bs, \sum{hw}, 4 + """ + N_, S_, C_ = memory.shape + proposals = [] + _cur = 0 + for lvl, (H_, W_) in enumerate(spatial_shapes): + mask_flatten_ = memory_padding_mask[:, _cur : (_cur + H_ * W_)].view(N_, H_, W_, 1) + valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1) + valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1) + + # import ipdb; ipdb.set_trace() + + grid_y, grid_x = torch.meshgrid( + torch.linspace(0, H_ - 1, H_, dtype=torch.float32, device=memory.device), + torch.linspace(0, W_ - 1, W_, dtype=torch.float32, device=memory.device), + ) + grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) # H_, W_, 2 + + scale = torch.cat([valid_W.unsqueeze(-1), valid_H.unsqueeze(-1)], 1).view(N_, 1, 1, 2) + grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale + + if learnedwh is not None: + # import ipdb; ipdb.set_trace() + wh = torch.ones_like(grid) * learnedwh.sigmoid() * (2.0**lvl) + else: + wh = torch.ones_like(grid) * 0.05 * (2.0**lvl) + + # scale = torch.cat([W_[None].unsqueeze(-1), H_[None].unsqueeze(-1)], 1).view(1, 1, 1, 2).repeat(N_, 1, 1, 1) + # grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale + # wh = torch.ones_like(grid) / scale + proposal = torch.cat((grid, wh), -1).view(N_, -1, 4) + proposals.append(proposal) + _cur += H_ * W_ + # import ipdb; ipdb.set_trace() + output_proposals = torch.cat(proposals, 1) + output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all( + -1, keepdim=True + ) + output_proposals = torch.log(output_proposals / (1 - output_proposals)) # unsigmoid + output_proposals = output_proposals.masked_fill(memory_padding_mask.unsqueeze(-1), float("inf")) + output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf")) + + output_memory = memory + output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float(0)) + output_memory = output_memory.masked_fill(~output_proposals_valid, float(0)) + + # output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float('inf')) + # output_memory = output_memory.masked_fill(~output_proposals_valid, float('inf')) + + output_proposals = output_proposals.to(output_memory.dtype) + return output_memory, output_proposals + + +class RandomBoxPerturber: + def __init__( + self, x_noise_scale=0.2, y_noise_scale=0.2, w_noise_scale=0.2, h_noise_scale=0.2 + ) -> None: + self.noise_scale = torch.Tensor( + [x_noise_scale, y_noise_scale, w_noise_scale, h_noise_scale] + ) + + def __call__(self, refanchors: Tensor) -> Tensor: + nq, bs, query_dim = refanchors.shape + device = refanchors.device + + noise_raw = torch.rand_like(refanchors) + noise_scale = self.noise_scale.to(device)[:query_dim] + + new_refanchors = refanchors * (1 + (noise_raw - 0.5) * noise_scale) + return new_refanchors.clamp_(0, 1) + + +def sigmoid_focal_loss( + inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2, no_reduction=False +): + """ + Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. + Args: + inputs: A float tensor of arbitrary shape. + The predictions for each example. + targets: A float tensor with the same shape as inputs. Stores the binary + classification label for each element in inputs + (0 for the negative class and 1 for the positive class). + alpha: (optional) Weighting factor in range (0,1) to balance + positive vs negative examples. Default = -1 (no weighting). + gamma: Exponent of the modulating factor (1 - p_t) to + balance easy vs hard examples. + Returns: + Loss tensor + """ + prob = inputs.sigmoid() + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = prob * targets + (1 - prob) * (1 - targets) + loss = ce_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * loss + + if no_reduction: + return loss + + return loss.mean(1).sum() / num_boxes + + +class MLP(nn.Module): + """Very simple multi-layer perceptron (also called FFN)""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +def _get_activation_fn(activation, d_model=256, batch_dim=0): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + if activation == "prelu": + return nn.PReLU() + if activation == "selu": + return F.selu + + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def gen_sineembed_for_position(pos_tensor): + # n_query, bs, _ = pos_tensor.size() + # sineembed_tensor = torch.zeros(n_query, bs, 256) + scale = 2 * math.pi + dim_t = torch.arange(128, dtype=torch.float32, device=pos_tensor.device) + dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode='floor')) / 128) + x_embed = pos_tensor[:, :, 0] * scale + y_embed = pos_tensor[:, :, 1] * scale + pos_x = x_embed[:, :, None] / dim_t + pos_y = y_embed[:, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) + pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) + if pos_tensor.size(-1) == 2: + pos = torch.cat((pos_y, pos_x), dim=2) + elif pos_tensor.size(-1) == 4: + w_embed = pos_tensor[:, :, 2] * scale + pos_w = w_embed[:, :, None] / dim_t + pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2) + + h_embed = pos_tensor[:, :, 3] * scale + pos_h = h_embed[:, :, None] / dim_t + pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2) + + pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) + else: + raise ValueError("Unknown pos_tensor shape(-1):{}".format(pos_tensor.size(-1))) + pos = pos.to(pos_tensor.dtype) + return pos + + +class ContrastiveEmbed(nn.Module): + def __init__(self, max_text_len=256): + """ + Args: + max_text_len: max length of text. + """ + super().__init__() + self.max_text_len = max_text_len + + def forward(self, x, text_dict): + """_summary_ + + Args: + x (_type_): _description_ + text_dict (_type_): _description_ + { + 'encoded_text': encoded_text, # bs, 195, d_model + 'text_token_mask': text_token_mask, # bs, 195 + # True for used tokens. False for padding tokens + } + Returns: + _type_: _description_ + """ + assert isinstance(text_dict, dict) + + y = text_dict["encoded_text"] + text_token_mask = text_dict["text_token_mask"] + + res = x @ y.transpose(-1, -2) + res.masked_fill_(~text_token_mask[:, None, :], float("-inf")) + + # padding to max_text_len + new_res = torch.full((*res.shape[:-1], self.max_text_len), float("-inf"), device=res.device, dtype=res.dtype) + new_res[..., : res.shape[-1]] = res + + return new_res diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/__init__.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e3413961d1d184b99835eb1e919b052d70298bc6 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/__init__.py @@ -0,0 +1,18 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +from .GroundingDINO import build_groundingdino + + +def build_model(args): + # we use register to maintain models from catdet6 on. + from .registry import MODULE_BUILD_FUNCS + + assert args.modelname in MODULE_BUILD_FUNCS._module_dict + build_func = MODULE_BUILD_FUNCS.get(args.modelname) + model = build_func(args) + return model diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87900c238577a114de73990d868beb39462cfd3f Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/registry.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..2d22a59eec79a2a19b83fa1779f2adaf5753aec6 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/models/registry.py @@ -0,0 +1,66 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# -*- coding: utf-8 -*- +# @Author: Yihao Chen +# @Date: 2021-08-16 16:03:17 +# @Last Modified by: Shilong Liu +# @Last Modified time: 2022-01-23 15:26 +# modified from mmcv + +import inspect +from functools import partial + + +class Registry(object): + def __init__(self, name): + self._name = name + self._module_dict = dict() + + def __repr__(self): + format_str = self.__class__.__name__ + "(name={}, items={})".format( + self._name, list(self._module_dict.keys()) + ) + return format_str + + def __len__(self): + return len(self._module_dict) + + @property + def name(self): + return self._name + + @property + def module_dict(self): + return self._module_dict + + def get(self, key): + return self._module_dict.get(key, None) + + def registe_with_name(self, module_name=None, force=False): + return partial(self.register, module_name=module_name, force=force) + + def register(self, module_build_function, module_name=None, force=False): + """Register a module build function. + Args: + module (:obj:`nn.Module`): Module to be registered. + """ + if not inspect.isfunction(module_build_function): + raise TypeError( + "module_build_function must be a function, but got {}".format( + type(module_build_function) + ) + ) + if module_name is None: + module_name = module_build_function.__name__ + if not force and module_name in self._module_dict: + raise KeyError("{} is already registered in {}".format(module_name, self.name)) + self._module_dict[module_name] = module_build_function + + return module_build_function + + +MODULE_BUILD_FUNCS = Registry("model build functions") diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__init__.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..168f9979a4623806934b0ff1102ac166704e7dec --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84df51d9a1a19906b2a987924e8b0abbc5972f54 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/box_ops.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/box_ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a10f34511df4ee7bbccf9bec819f36450442e3b Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/box_ops.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/get_tokenlizer.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/get_tokenlizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c11731abbe3497787c4ee851dedf435720df20c Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/get_tokenlizer.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/inference.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2da245fa6fed294f378c6e24d8a4be294c7ebe3 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/inference.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/misc.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/misc.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3930c4130c24e20cdda9df1cbc8cd6fc67f563be Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/misc.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/slconfig.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/slconfig.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dff34dbdf701d3ce47930726c8f0d41eaed8bbb Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/slconfig.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/utils.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1c171fd6ec49a7cdd25a58a33108002f926cbdd Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/__pycache__/utils.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/box_ops.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/box_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..781068d294e576954edb4bd07b6e0f30e4e1bcd9 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/box_ops.py @@ -0,0 +1,140 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Utilities for bounding box manipulation and GIoU. +""" +import torch +from torchvision.ops.boxes import box_area + + +def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.stack(b, dim=-1) + + +def box_xyxy_to_cxcywh(x): + x0, y0, x1, y1 = x.unbind(-1) + b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)] + return torch.stack(b, dim=-1) + + +# modified from torchvision to also return the union +def box_iou(boxes1, boxes2): + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + # import ipdb; ipdb.set_trace() + lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] + rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] + + wh = (rb - lt).clamp(min=0) # [N,M,2] + inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] + + union = area1[:, None] + area2 - inter + + iou = inter / (union + 1e-6) + return iou, union + + +def generalized_box_iou(boxes1, boxes2): + """ + Generalized IoU from https://giou.stanford.edu/ + + The boxes should be in [x0, y0, x1, y1] format + + Returns a [N, M] pairwise matrix, where N = len(boxes1) + and M = len(boxes2) + """ + # degenerate boxes gives inf / nan results + # so do an early check + assert (boxes1[:, 2:] >= boxes1[:, :2]).all() + assert (boxes2[:, 2:] >= boxes2[:, :2]).all() + # except: + # import ipdb; ipdb.set_trace() + iou, union = box_iou(boxes1, boxes2) + + lt = torch.min(boxes1[:, None, :2], boxes2[:, :2]) + rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:]) + + wh = (rb - lt).clamp(min=0) # [N,M,2] + area = wh[:, :, 0] * wh[:, :, 1] + + return iou - (area - union) / (area + 1e-6) + + +# modified from torchvision to also return the union +def box_iou_pairwise(boxes1, boxes2): + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + lt = torch.max(boxes1[:, :2], boxes2[:, :2]) # [N,2] + rb = torch.min(boxes1[:, 2:], boxes2[:, 2:]) # [N,2] + + wh = (rb - lt).clamp(min=0) # [N,2] + inter = wh[:, 0] * wh[:, 1] # [N] + + union = area1 + area2 - inter + + iou = inter / union + return iou, union + + +def generalized_box_iou_pairwise(boxes1, boxes2): + """ + Generalized IoU from https://giou.stanford.edu/ + + Input: + - boxes1, boxes2: N,4 + Output: + - giou: N, 4 + """ + # degenerate boxes gives inf / nan results + # so do an early check + assert (boxes1[:, 2:] >= boxes1[:, :2]).all() + assert (boxes2[:, 2:] >= boxes2[:, :2]).all() + assert boxes1.shape == boxes2.shape + iou, union = box_iou_pairwise(boxes1, boxes2) # N, 4 + + lt = torch.min(boxes1[:, :2], boxes2[:, :2]) + rb = torch.max(boxes1[:, 2:], boxes2[:, 2:]) + + wh = (rb - lt).clamp(min=0) # [N,2] + area = wh[:, 0] * wh[:, 1] + + return iou - (area - union) / area + + +def masks_to_boxes(masks): + """Compute the bounding boxes around the provided masks + + The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. + + Returns a [N, 4] tensors, with the boxes in xyxy format + """ + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device) + + h, w = masks.shape[-2:] + + y = torch.arange(0, h, dtype=torch.float) + x = torch.arange(0, w, dtype=torch.float) + y, x = torch.meshgrid(y, x) + + x_mask = masks * x.unsqueeze(0) + x_max = x_mask.flatten(1).max(-1)[0] + x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + y_mask = masks * y.unsqueeze(0) + y_max = y_mask.flatten(1).max(-1)[0] + y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + return torch.stack([x_min, y_min, x_max, y_max], 1) + + +if __name__ == "__main__": + x = torch.rand(5, 4) + y = torch.rand(3, 4) + iou, union = box_iou(x, y) + import ipdb + + ipdb.set_trace() diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/get_tokenlizer.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/get_tokenlizer.py new file mode 100644 index 0000000000000000000000000000000000000000..b7b5d72aef873453361cc019427ba31b08c11798 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/get_tokenlizer.py @@ -0,0 +1,37 @@ +from transformers import AutoTokenizer, BertModel, BertTokenizer, RobertaModel, RobertaTokenizerFast + + +def get_tokenlizer(text_encoder_type, bert_base_uncased_path): + if not isinstance(text_encoder_type, str): + # print("text_encoder_type is not a str") + if hasattr(text_encoder_type, "text_encoder_type"): + text_encoder_type = text_encoder_type.text_encoder_type + elif text_encoder_type.get("text_encoder_type", False): + text_encoder_type = text_encoder_type.get("text_encoder_type") + else: + raise ValueError( + "Unknown type of text_encoder_type: {}".format(type(text_encoder_type)) + ) + + # solve huggingface connect issue + if is_bert_model_use_local_path(bert_base_uncased_path) and text_encoder_type == "bert-base-uncased": + print("use local bert model path: {}".format(bert_base_uncased_path)) + return AutoTokenizer.from_pretrained(bert_base_uncased_path) + + print("final text_encoder_type: {}".format(text_encoder_type)) + + tokenizer = AutoTokenizer.from_pretrained(text_encoder_type) + return tokenizer + + +def get_pretrained_language_model(text_encoder_type, bert_base_uncased_path): + if text_encoder_type == "bert-base-uncased": + if is_bert_model_use_local_path(bert_base_uncased_path): + return BertModel.from_pretrained(bert_base_uncased_path) + return BertModel.from_pretrained(text_encoder_type) + if text_encoder_type == "roberta-base": + return RobertaModel.from_pretrained(text_encoder_type) + raise ValueError("Unknown text_encoder_type {}".format(text_encoder_type)) + +def is_bert_model_use_local_path(bert_base_uncased_path): + return bert_base_uncased_path is not None and len(bert_base_uncased_path) > 0 diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/inference.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..fe8077e103e5e269ad175ab056699136cf1c74b1 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/inference.py @@ -0,0 +1,259 @@ +from typing import Tuple, List + +import re +import cv2 +import numpy as np +import supervision as sv +import torch +from PIL import Image +from torchvision.ops import box_convert + +import groundingdino.datasets.transforms as T +from groundingdino.models import build_model +from groundingdino.util.misc import clean_state_dict +from groundingdino.util.slconfig import SLConfig +from groundingdino.util.utils import get_phrases_from_posmap + +# ---------------------------------------------------------------------------------------------------------------------- +# OLD API +# ---------------------------------------------------------------------------------------------------------------------- + + +def preprocess_caption(caption: str) -> str: + result = caption.lower().strip() + if result.endswith("."): + return result + return result + "." + + +def load_model(model_config_path: str, model_checkpoint_path: str, device: str = "cuda"): + args = SLConfig.fromfile(model_config_path) + args.device = device + model = build_model(args) + checkpoint = torch.load(model_checkpoint_path, map_location="cpu") + model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) + model.eval() + return model + + +def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]: + transform = T.Compose( + [ + T.RandomResize([800], max_size=1333), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + image_source = Image.open(image_path).convert("RGB") + image = np.asarray(image_source) + image_transformed, _ = transform(image_source, None) + return image, image_transformed + + +def predict( + model, + image: torch.Tensor, + caption: str, + box_threshold: float, + text_threshold: float, + device: str = "cuda" +) -> Tuple[torch.Tensor, torch.Tensor, List[str]]: + caption = preprocess_caption(caption=caption) + + model = model.to(device) + image = image.to(device) + + with torch.no_grad(): + outputs = model(image[None], captions=[caption]) + + prediction_logits = outputs["pred_logits"].cpu().sigmoid()[0] # prediction_logits.shape = (nq, 256) + prediction_boxes = outputs["pred_boxes"].cpu()[0] # prediction_boxes.shape = (nq, 4) + + mask = prediction_logits.max(dim=1)[0] > box_threshold + logits = prediction_logits[mask] # logits.shape = (n, 256) + boxes = prediction_boxes[mask] # boxes.shape = (n, 4) + + tokenizer = model.tokenizer + tokenized = tokenizer(caption) + + phrases = [ + get_phrases_from_posmap(logit > text_threshold, tokenized, tokenizer).replace('.', '') + for logit + in logits + ] + + return boxes, logits.max(dim=1)[0], phrases + + +def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str]) -> np.ndarray: + h, w, _ = image_source.shape + boxes = boxes * torch.Tensor([w, h, w, h]) + xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() + detections = sv.Detections(xyxy=xyxy) + + labels = [ + f"{phrase} {logit:.2f}" + for phrase, logit + in zip(phrases, logits) + ] + + box_annotator = sv.BoxAnnotator() + # box_annotator = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX) + annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR) + annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels) + return annotated_frame + + +# ---------------------------------------------------------------------------------------------------------------------- +# NEW API +# ---------------------------------------------------------------------------------------------------------------------- + + +class Model: + + def __init__( + self, + model_config_path: str, + model_checkpoint_path: str, + device: str = "cuda" + ): + + self.model = load_model( + model_config_path=model_config_path, + model_checkpoint_path=model_checkpoint_path, + device=device + ).to(device) + self.device = device + + def predict_with_caption( + self, + image: np.ndarray, + caption: str, + box_threshold: float = 0.35, + text_threshold: float = 0.25 + ) -> Tuple[sv.Detections, List[str]]: + """ + import cv2 + + image = cv2.imread(IMAGE_PATH) + + model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH) + detections, labels = model.predict_with_caption( + image=image, + caption=caption, + box_threshold=BOX_THRESHOLD, + text_threshold=TEXT_THRESHOLD + ) + + import supervision as sv + + box_annotator = sv.BoxAnnotator() + annotated_image = box_annotator.annotate(scene=image, detections=detections, labels=labels) + """ + processed_image = Model.preprocess_image(image_bgr=image).to(self.device) + boxes, logits, phrases = predict( + model=self.model, + image=processed_image, + caption=caption, + box_threshold=box_threshold, + text_threshold=text_threshold, + device=self.device) + source_h, source_w, _ = image.shape + detections = Model.post_process_result( + source_h=source_h, + source_w=source_w, + boxes=boxes, + logits=logits) + return detections, phrases + + def predict_with_classes( + self, + image: np.ndarray, + classes: List[str], + box_threshold: float, + text_threshold: float + ) -> sv.Detections: + """ + import cv2 + + image = cv2.imread(IMAGE_PATH) + + model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH) + detections = model.predict_with_classes( + image=image, + classes=CLASSES, + box_threshold=BOX_THRESHOLD, + text_threshold=TEXT_THRESHOLD + ) + + + import supervision as sv + + box_annotator = sv.BoxAnnotator() + annotated_image = box_annotator.annotate(scene=image, detections=detections) + """ + caption = ". ".join(classes) + processed_image = Model.preprocess_image(image_bgr=image).to(self.device) + boxes, logits, phrases = predict( + model=self.model, + image=processed_image, + caption=caption, + box_threshold=box_threshold, + text_threshold=text_threshold, + device=self.device) + source_h, source_w, _ = image.shape + detections = Model.post_process_result( + source_h=source_h, + source_w=source_w, + boxes=boxes, + logits=logits) + class_id = Model.phrases2classes(phrases=phrases, classes=classes) + detections.class_id = class_id + return detections + + @staticmethod + def preprocess_image(image_bgr: np.ndarray) -> torch.Tensor: + transform = T.Compose( + [ + T.RandomResize([800], max_size=1333), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + image_pillow = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)) + image_transformed, _ = transform(image_pillow, None) + return image_transformed + + @staticmethod + def post_process_result( + source_h: int, + source_w: int, + boxes: torch.Tensor, + logits: torch.Tensor + ) -> sv.Detections: + boxes = boxes * torch.Tensor([source_w, source_h, source_w, source_h]) + xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() + confidence = logits.numpy() + return sv.Detections(xyxy=xyxy, confidence=confidence) + + @staticmethod + def phrases2classes(phrases: List[str], classes: List[str]) -> np.ndarray: + class_ids = [] + for phrase in phrases: + try: + # class_ids.append(classes.index(phrase)) + class_ids.append(Model.find_index(phrase, classes)) + except ValueError: + class_ids.append(None) + return np.array(class_ids) + + @staticmethod + def find_index(string, lst): + # if meet string like "lake river" will only keep "lake" + # this is an hack implementation for visualization which will be updated in the future + string = string.lower().split()[0] + for i, s in enumerate(lst): + if string in s.lower(): + return i + print("There's a wrong phrase happen, this is because of our post-process merged wrong tokens, which will be modified in the future. We will assign it with a random label at this time.") + return 0 \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/logger.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..18145f54c927abd59b95f3fa6e6da8002bc2ce97 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/logger.py @@ -0,0 +1,93 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +import functools +import logging +import os +import sys + +from termcolor import colored + + +class _ColorfulFormatter(logging.Formatter): + def __init__(self, *args, **kwargs): + self._root_name = kwargs.pop("root_name") + "." + self._abbrev_name = kwargs.pop("abbrev_name", "") + if len(self._abbrev_name): + self._abbrev_name = self._abbrev_name + "." + super(_ColorfulFormatter, self).__init__(*args, **kwargs) + + def formatMessage(self, record): + record.name = record.name.replace(self._root_name, self._abbrev_name) + log = super(_ColorfulFormatter, self).formatMessage(record) + if record.levelno == logging.WARNING: + prefix = colored("WARNING", "red", attrs=["blink"]) + elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL: + prefix = colored("ERROR", "red", attrs=["blink", "underline"]) + else: + return log + return prefix + " " + log + + +# so that calling setup_logger multiple times won't add many handlers +@functools.lru_cache() +def setup_logger(output=None, distributed_rank=0, *, color=True, name="imagenet", abbrev_name=None): + """ + Initialize the detectron2 logger and set its verbosity level to "INFO". + + Args: + output (str): a file name or a directory to save log. If None, will not save log file. + If ends with ".txt" or ".log", assumed to be a file name. + Otherwise, logs will be saved to `output/log.txt`. + name (str): the root module name of this logger + + Returns: + logging.Logger: a logger + """ + logger = logging.getLogger(name) + logger.setLevel(logging.DEBUG) + logger.propagate = False + + if abbrev_name is None: + abbrev_name = name + + plain_formatter = logging.Formatter( + "[%(asctime)s.%(msecs)03d]: %(message)s", datefmt="%m/%d %H:%M:%S" + ) + # stdout logging: master only + if distributed_rank == 0: + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(logging.DEBUG) + if color: + formatter = _ColorfulFormatter( + colored("[%(asctime)s.%(msecs)03d]: ", "green") + "%(message)s", + datefmt="%m/%d %H:%M:%S", + root_name=name, + abbrev_name=str(abbrev_name), + ) + else: + formatter = plain_formatter + ch.setFormatter(formatter) + logger.addHandler(ch) + + # file logging: all workers + if output is not None: + if output.endswith(".txt") or output.endswith(".log"): + filename = output + else: + filename = os.path.join(output, "log.txt") + if distributed_rank > 0: + filename = filename + f".rank{distributed_rank}" + os.makedirs(os.path.dirname(filename), exist_ok=True) + + fh = logging.StreamHandler(_cached_log_stream(filename)) + fh.setLevel(logging.DEBUG) + fh.setFormatter(plain_formatter) + logger.addHandler(fh) + + return logger + + +# cache the opened file object, so that different calls to `setup_logger` +# with the same file name can safely write to the same file. +@functools.lru_cache(maxsize=None) +def _cached_log_stream(filename): + return open(filename, "a") diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/misc.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..d64b84ef24bea0c98e76824feb1903f6bfebe7a5 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/misc.py @@ -0,0 +1,717 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Misc functions, including distributed helpers. + +Mostly copy-paste from torchvision references. +""" +import colorsys +import datetime +import functools +import io +import json +import os +import pickle +import subprocess +import time +from collections import OrderedDict, defaultdict, deque +from typing import List, Optional + +import numpy as np +import torch +import torch.distributed as dist + +# needed due to empty tensor bug in pytorch and torchvision 0.5 +import torchvision +from torch import Tensor + +__torchvision_need_compat_flag = float(torchvision.__version__.split(".")[1]) < 7 +if __torchvision_need_compat_flag: + from torchvision.ops import _new_empty_tensor + from torchvision.ops.misc import _output_size + + +class SmoothedValue(object): + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda") + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + d = torch.tensor(list(self.deque)) + if d.shape[0] == 0: + return 0 + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + if os.environ.get("SHILONG_AMP", None) == "1": + eps = 1e-4 + else: + eps = 1e-6 + return self.total / (self.count + eps) + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +@functools.lru_cache() +def _get_global_gloo_group(): + """ + Return a process group based on gloo backend, containing all the ranks + The result is cached. + """ + + if dist.get_backend() == "nccl": + return dist.new_group(backend="gloo") + + return dist.group.WORLD + + +def all_gather_cpu(data): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + + world_size = get_world_size() + if world_size == 1: + return [data] + + cpu_group = _get_global_gloo_group() + + buffer = io.BytesIO() + torch.save(data, buffer) + data_view = buffer.getbuffer() + device = "cuda" if cpu_group is None else "cpu" + tensor = torch.ByteTensor(data_view).to(device) + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device=device, dtype=torch.long) + size_list = [torch.tensor([0], device=device, dtype=torch.long) for _ in range(world_size)] + if cpu_group is None: + dist.all_gather(size_list, local_size) + else: + print("gathering on cpu") + dist.all_gather(size_list, local_size, group=cpu_group) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + assert isinstance(local_size.item(), int) + local_size = int(local_size.item()) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=device)) + if local_size != max_size: + padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device=device) + tensor = torch.cat((tensor, padding), dim=0) + if cpu_group is None: + dist.all_gather(tensor_list, tensor) + else: + dist.all_gather(tensor_list, tensor, group=cpu_group) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + tensor = torch.split(tensor, [size, max_size - size], dim=0)[0] + buffer = io.BytesIO(tensor.cpu().numpy()) + obj = torch.load(buffer) + data_list.append(obj) + + return data_list + + +def all_gather(data): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + + if os.getenv("CPU_REDUCE") == "1": + return all_gather_cpu(data) + + world_size = get_world_size() + if world_size == 1: + return [data] + + # serialized to a Tensor + buffer = pickle.dumps(data) + storage = torch.ByteStorage.from_buffer(buffer) + tensor = torch.ByteTensor(storage).to("cuda") + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device="cuda") + size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)] + dist.all_gather(size_list, local_size) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda")) + if local_size != max_size: + padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda") + tensor = torch.cat((tensor, padding), dim=0) + dist.all_gather(tensor_list, tensor) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + buffer = tensor.cpu().numpy().tobytes()[:size] + data_list.append(pickle.loads(buffer)) + + return data_list + + +def reduce_dict(input_dict, average=True): + """ + Args: + input_dict (dict): all the values will be reduced + average (bool): whether to do average or sum + Reduce the values in the dictionary from all processes so that all processes + have the averaged results. Returns a dict with the same fields as + input_dict, after reduction. + """ + world_size = get_world_size() + if world_size < 2: + return input_dict + with torch.no_grad(): + names = [] + values = [] + # sort the keys so that they are consistent across processes + for k in sorted(input_dict.keys()): + names.append(k) + values.append(input_dict[k]) + values = torch.stack(values, dim=0) + dist.all_reduce(values) + if average: + values /= world_size + reduced_dict = {k: v for k, v in zip(names, values)} + return reduced_dict + + +class MetricLogger(object): + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + # print(name, str(meter)) + # import ipdb;ipdb.set_trace() + if meter.count > 0: + loss_str.append("{}: {}".format(name, str(meter))) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None, logger=None): + if logger is None: + print_func = print + else: + print_func = logger.info + + i = 0 + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(len(iterable)))) + "d" + if torch.cuda.is_available(): + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + "max mem: {memory:.0f}", + ] + ) + else: + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + ) + MB = 1024.0 * 1024.0 + for obj in iterable: + data_time.update(time.time() - end) + yield obj + # import ipdb; ipdb.set_trace() + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == len(iterable) - 1: + eta_seconds = iter_time.global_avg * (len(iterable) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print_func( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + else: + print_func( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print_func( + "{} Total time: {} ({:.4f} s / it)".format( + header, total_time_str, total_time / len(iterable) + ) + ) + + +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + subprocess.check_output(["git", "diff"], cwd=cwd) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommited changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +def collate_fn(batch): + # import ipdb; ipdb.set_trace() + batch = list(zip(*batch)) + batch[0] = nested_tensor_from_tensor_list(batch[0]) + return tuple(batch) + + +def _max_by_axis(the_list): + # type: (List[List[int]]) -> List[int] + maxes = the_list[0] + for sublist in the_list[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + +class NestedTensor(object): + def __init__(self, tensors, mask: Optional[Tensor]): + self.tensors = tensors + self.mask = mask + if mask == "auto": + self.mask = torch.zeros_like(tensors).to(tensors.device) + if self.mask.dim() == 3: + self.mask = self.mask.sum(0).to(bool) + elif self.mask.dim() == 4: + self.mask = self.mask.sum(1).to(bool) + else: + raise ValueError( + "tensors dim must be 3 or 4 but {}({})".format( + self.tensors.dim(), self.tensors.shape + ) + ) + + def imgsize(self): + res = [] + for i in range(self.tensors.shape[0]): + mask = self.mask[i] + maxH = (~mask).sum(0).max() + maxW = (~mask).sum(1).max() + res.append(torch.Tensor([maxH, maxW])) + return res + + def to(self, device): + # type: (Device) -> NestedTensor # noqa + cast_tensor = self.tensors.to(device) + mask = self.mask + if mask is not None: + assert mask is not None + cast_mask = mask.to(device) + else: + cast_mask = None + return NestedTensor(cast_tensor, cast_mask) + + def to_img_list_single(self, tensor, mask): + assert tensor.dim() == 3, "dim of tensor should be 3 but {}".format(tensor.dim()) + maxH = (~mask).sum(0).max() + maxW = (~mask).sum(1).max() + img = tensor[:, :maxH, :maxW] + return img + + def to_img_list(self): + """remove the padding and convert to img list + + Returns: + [type]: [description] + """ + if self.tensors.dim() == 3: + return self.to_img_list_single(self.tensors, self.mask) + else: + res = [] + for i in range(self.tensors.shape[0]): + tensor_i = self.tensors[i] + mask_i = self.mask[i] + res.append(self.to_img_list_single(tensor_i, mask_i)) + return res + + @property + def device(self): + return self.tensors.device + + def decompose(self): + return self.tensors, self.mask + + def __repr__(self): + return str(self.tensors) + + @property + def shape(self): + return {"tensors.shape": self.tensors.shape, "mask.shape": self.mask.shape} + + +def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): + # TODO make this more general + if tensor_list[0].ndim == 3: + if torchvision._is_tracing(): + # nested_tensor_from_tensor_list() does not export well to ONNX + # call _onnx_nested_tensor_from_tensor_list() instead + return _onnx_nested_tensor_from_tensor_list(tensor_list) + + # TODO make it support different-sized images + max_size = _max_by_axis([list(img.shape) for img in tensor_list]) + # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) + batch_shape = [len(tensor_list)] + max_size + b, c, h, w = batch_shape + dtype = tensor_list[0].dtype + device = tensor_list[0].device + tensor = torch.zeros(batch_shape, dtype=dtype, device=device) + mask = torch.ones((b, h, w), dtype=torch.bool, device=device) + for img, pad_img, m in zip(tensor_list, tensor, mask): + pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + m[: img.shape[1], : img.shape[2]] = False + else: + raise ValueError("not supported") + return NestedTensor(tensor, mask) + + +# _onnx_nested_tensor_from_tensor_list() is an implementation of +# nested_tensor_from_tensor_list() that is supported by ONNX tracing. +@torch.jit.unused +def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor: + max_size = [] + for i in range(tensor_list[0].dim()): + max_size_i = torch.max( + torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32) + ).to(torch.int64) + max_size.append(max_size_i) + max_size = tuple(max_size) + + # work around for + # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + # m[: img.shape[1], :img.shape[2]] = False + # which is not yet supported in onnx + padded_imgs = [] + padded_masks = [] + for img in tensor_list: + padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))] + padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0])) + padded_imgs.append(padded_img) + + m = torch.zeros_like(img[0], dtype=torch.int, device=img.device) + padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1) + padded_masks.append(padded_mask.to(torch.bool)) + + tensor = torch.stack(padded_imgs) + mask = torch.stack(padded_masks) + + return NestedTensor(tensor, mask=mask) + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(*args, **kwargs): + if is_main_process(): + torch.save(*args, **kwargs) + + +def init_distributed_mode(args): + if "WORLD_SIZE" in os.environ and os.environ["WORLD_SIZE"] != "": # 'RANK' in os.environ and + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = args.local_rank = int(os.environ["LOCAL_RANK"]) + + # launch by torch.distributed.launch + # Single node + # python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 1 --rank 0 ... + # Multi nodes + # python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 2 --rank 0 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' ... + # python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 2 --rank 1 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' ... + # args.rank = int(os.environ.get('OMPI_COMM_WORLD_RANK')) + # local_world_size = int(os.environ['GPU_PER_NODE_COUNT']) + # args.world_size = args.world_size * local_world_size + # args.gpu = args.local_rank = int(os.environ['LOCAL_RANK']) + # args.rank = args.rank * local_world_size + args.local_rank + print( + "world size: {}, rank: {}, local rank: {}".format( + args.world_size, args.rank, args.local_rank + ) + ) + print(json.dumps(dict(os.environ), indent=2)) + elif "SLURM_PROCID" in os.environ: + args.rank = int(os.environ["SLURM_PROCID"]) + args.gpu = args.local_rank = int(os.environ["SLURM_LOCALID"]) + args.world_size = int(os.environ["SLURM_NPROCS"]) + + print( + "world size: {}, world rank: {}, local rank: {}, device_count: {}".format( + args.world_size, args.rank, args.local_rank, torch.cuda.device_count() + ) + ) + else: + print("Not using distributed mode") + args.distributed = False + args.world_size = 1 + args.rank = 0 + args.local_rank = 0 + return + + print("world_size:{} rank:{} local_rank:{}".format(args.world_size, args.rank, args.local_rank)) + args.distributed = True + torch.cuda.set_device(args.local_rank) + args.dist_backend = "nccl" + print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True) + + torch.distributed.init_process_group( + backend=args.dist_backend, + world_size=args.world_size, + rank=args.rank, + init_method=args.dist_url, + ) + + print("Before torch.distributed.barrier()") + torch.distributed.barrier() + print("End torch.distributed.barrier()") + setup_for_distributed(args.rank == 0) + + +@torch.no_grad() +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + if target.numel() == 0: + return [torch.zeros([], device=output.device)] + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +@torch.no_grad() +def accuracy_onehot(pred, gt): + """_summary_ + + Args: + pred (_type_): n, c + gt (_type_): n, c + """ + tp = ((pred - gt).abs().sum(-1) < 1e-4).float().sum() + acc = tp / gt.shape[0] * 100 + return acc + + +def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None): + # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor + """ + Equivalent to nn.functional.interpolate, but with support for empty batch sizes. + This will eventually be supported natively by PyTorch, and this + class can go away. + """ + if __torchvision_need_compat_flag < 0.7: + if input.numel() > 0: + return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners) + + output_shape = _output_size(2, input, size, scale_factor) + output_shape = list(input.shape[:-2]) + list(output_shape) + return _new_empty_tensor(input, output_shape) + else: + return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners) + + +class color_sys: + def __init__(self, num_colors) -> None: + self.num_colors = num_colors + colors = [] + for i in np.arange(0.0, 360.0, 360.0 / num_colors): + hue = i / 360.0 + lightness = (50 + np.random.rand() * 10) / 100.0 + saturation = (90 + np.random.rand() * 10) / 100.0 + colors.append( + tuple([int(j * 255) for j in colorsys.hls_to_rgb(hue, lightness, saturation)]) + ) + self.colors = colors + + def __call__(self, idx): + return self.colors[idx] + + +def inverse_sigmoid(x, eps=1e-3): + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + + +def clean_state_dict(state_dict): + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + if k[:7] == "module.": + k = k[7:] # remove `module.` + new_state_dict[k] = v + return new_state_dict diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/slconfig.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/slconfig.py new file mode 100644 index 0000000000000000000000000000000000000000..3f293e3aff215a3c7c2f7d21d27853493b6ebfbc --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/slconfig.py @@ -0,0 +1,427 @@ +# ========================================================== +# Modified from mmcv +# ========================================================== +import ast +import os.path as osp +import shutil +import sys +import tempfile +from argparse import Action +from importlib import import_module +import platform + +from addict import Dict +from yapf.yapflib.yapf_api import FormatCode + +BASE_KEY = "_base_" +DELETE_KEY = "_delete_" +RESERVED_KEYS = ["filename", "text", "pretty_text", "get", "dump", "merge_from_dict"] + + +def check_file_exist(filename, msg_tmpl='file "{}" does not exist'): + if not osp.isfile(filename): + raise FileNotFoundError(msg_tmpl.format(filename)) + + +class ConfigDict(Dict): + def __missing__(self, name): + raise KeyError(name) + + def __getattr__(self, name): + try: + value = super(ConfigDict, self).__getattr__(name) + except KeyError: + ex = AttributeError(f"'{self.__class__.__name__}' object has no " f"attribute '{name}'") + except Exception as e: + ex = e + else: + return value + raise ex + + +class SLConfig(object): + """ + config files. + only support .py file as config now. + + ref: mmcv.utils.config + + Example: + >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1]))) + >>> cfg.a + 1 + >>> cfg.b + {'b1': [0, 1]} + >>> cfg.b.b1 + [0, 1] + >>> cfg = Config.fromfile('tests/data/config/a.py') + >>> cfg.filename + "/home/kchen/projects/mmcv/tests/data/config/a.py" + >>> cfg.item4 + 'test' + >>> cfg + "Config [path: /home/kchen/projects/mmcv/tests/data/config/a.py]: " + "{'item1': [1, 2], 'item2': {'a': 0}, 'item3': True, 'item4': 'test'}" + """ + + @staticmethod + def _validate_py_syntax(filename): + with open(filename) as f: + content = f.read() + try: + ast.parse(content) + except SyntaxError: + raise SyntaxError("There are syntax errors in config " f"file {filename}") + + @staticmethod + def _file2dict(filename): + filename = osp.abspath(osp.expanduser(filename)) + check_file_exist(filename) + if filename.lower().endswith(".py"): + with tempfile.TemporaryDirectory() as temp_config_dir: + temp_config_file = tempfile.NamedTemporaryFile(dir=temp_config_dir, suffix=".py") + temp_config_name = osp.basename(temp_config_file.name) + if platform.system() == 'Windows': + temp_config_file.close() + shutil.copyfile(filename, osp.join(temp_config_dir, temp_config_name)) + temp_module_name = osp.splitext(temp_config_name)[0] + sys.path.insert(0, temp_config_dir) + SLConfig._validate_py_syntax(filename) + mod = import_module(temp_module_name) + sys.path.pop(0) + cfg_dict = { + name: value for name, value in mod.__dict__.items() if not name.startswith("__") + } + # delete imported module + del sys.modules[temp_module_name] + # close temp file + temp_config_file.close() + elif filename.lower().endswith((".yml", ".yaml", ".json")): + from .slio import slload + + cfg_dict = slload(filename) + else: + raise IOError("Only py/yml/yaml/json type are supported now!") + + cfg_text = filename + "\n" + with open(filename, "r") as f: + cfg_text += f.read() + + # parse the base file + if BASE_KEY in cfg_dict: + cfg_dir = osp.dirname(filename) + base_filename = cfg_dict.pop(BASE_KEY) + base_filename = base_filename if isinstance(base_filename, list) else [base_filename] + + cfg_dict_list = list() + cfg_text_list = list() + for f in base_filename: + _cfg_dict, _cfg_text = SLConfig._file2dict(osp.join(cfg_dir, f)) + cfg_dict_list.append(_cfg_dict) + cfg_text_list.append(_cfg_text) + + base_cfg_dict = dict() + for c in cfg_dict_list: + if len(base_cfg_dict.keys() & c.keys()) > 0: + raise KeyError("Duplicate key is not allowed among bases") + # TODO Allow the duplicate key while warnning user + base_cfg_dict.update(c) + + base_cfg_dict = SLConfig._merge_a_into_b(cfg_dict, base_cfg_dict) + cfg_dict = base_cfg_dict + + # merge cfg_text + cfg_text_list.append(cfg_text) + cfg_text = "\n".join(cfg_text_list) + + return cfg_dict, cfg_text + + @staticmethod + def _merge_a_into_b(a, b): + """merge dict `a` into dict `b` (non-inplace). + values in `a` will overwrite `b`. + copy first to avoid inplace modification + + Args: + a ([type]): [description] + b ([type]): [description] + + Returns: + [dict]: [description] + """ + # import ipdb; ipdb.set_trace() + if not isinstance(a, dict): + return a + + b = b.copy() + for k, v in a.items(): + if isinstance(v, dict) and k in b and not v.pop(DELETE_KEY, False): + + if not isinstance(b[k], dict) and not isinstance(b[k], list): + # if : + # import ipdb; ipdb.set_trace() + raise TypeError( + f"{k}={v} in child config cannot inherit from base " + f"because {k} is a dict in the child config but is of " + f"type {type(b[k])} in base config. You may set " + f"`{DELETE_KEY}=True` to ignore the base config" + ) + b[k] = SLConfig._merge_a_into_b(v, b[k]) + elif isinstance(b, list): + try: + _ = int(k) + except: + raise TypeError( + f"b is a list, " f"index {k} should be an int when input but {type(k)}" + ) + b[int(k)] = SLConfig._merge_a_into_b(v, b[int(k)]) + else: + b[k] = v + + return b + + @staticmethod + def fromfile(filename): + cfg_dict, cfg_text = SLConfig._file2dict(filename) + return SLConfig(cfg_dict, cfg_text=cfg_text, filename=filename) + + def __init__(self, cfg_dict=None, cfg_text=None, filename=None): + if cfg_dict is None: + cfg_dict = dict() + elif not isinstance(cfg_dict, dict): + raise TypeError("cfg_dict must be a dict, but " f"got {type(cfg_dict)}") + for key in cfg_dict: + if key in RESERVED_KEYS: + raise KeyError(f"{key} is reserved for config file") + + super(SLConfig, self).__setattr__("_cfg_dict", ConfigDict(cfg_dict)) + super(SLConfig, self).__setattr__("_filename", filename) + if cfg_text: + text = cfg_text + elif filename: + with open(filename, "r") as f: + text = f.read() + else: + text = "" + super(SLConfig, self).__setattr__("_text", text) + + @property + def filename(self): + return self._filename + + @property + def text(self): + return self._text + + @property + def pretty_text(self): + + indent = 4 + + def _indent(s_, num_spaces): + s = s_.split("\n") + if len(s) == 1: + return s_ + first = s.pop(0) + s = [(num_spaces * " ") + line for line in s] + s = "\n".join(s) + s = first + "\n" + s + return s + + def _format_basic_types(k, v, use_mapping=False): + if isinstance(v, str): + v_str = f"'{v}'" + else: + v_str = str(v) + + if use_mapping: + k_str = f"'{k}'" if isinstance(k, str) else str(k) + attr_str = f"{k_str}: {v_str}" + else: + attr_str = f"{str(k)}={v_str}" + attr_str = _indent(attr_str, indent) + + return attr_str + + def _format_list(k, v, use_mapping=False): + # check if all items in the list are dict + if all(isinstance(_, dict) for _ in v): + v_str = "[\n" + v_str += "\n".join( + f"dict({_indent(_format_dict(v_), indent)})," for v_ in v + ).rstrip(",") + if use_mapping: + k_str = f"'{k}'" if isinstance(k, str) else str(k) + attr_str = f"{k_str}: {v_str}" + else: + attr_str = f"{str(k)}={v_str}" + attr_str = _indent(attr_str, indent) + "]" + else: + attr_str = _format_basic_types(k, v, use_mapping) + return attr_str + + def _contain_invalid_identifier(dict_str): + contain_invalid_identifier = False + for key_name in dict_str: + contain_invalid_identifier |= not str(key_name).isidentifier() + return contain_invalid_identifier + + def _format_dict(input_dict, outest_level=False): + r = "" + s = [] + + use_mapping = _contain_invalid_identifier(input_dict) + if use_mapping: + r += "{" + for idx, (k, v) in enumerate(input_dict.items()): + is_last = idx >= len(input_dict) - 1 + end = "" if outest_level or is_last else "," + if isinstance(v, dict): + v_str = "\n" + _format_dict(v) + if use_mapping: + k_str = f"'{k}'" if isinstance(k, str) else str(k) + attr_str = f"{k_str}: dict({v_str}" + else: + attr_str = f"{str(k)}=dict({v_str}" + attr_str = _indent(attr_str, indent) + ")" + end + elif isinstance(v, list): + attr_str = _format_list(k, v, use_mapping) + end + else: + attr_str = _format_basic_types(k, v, use_mapping) + end + + s.append(attr_str) + r += "\n".join(s) + if use_mapping: + r += "}" + return r + + cfg_dict = self._cfg_dict.to_dict() + text = _format_dict(cfg_dict, outest_level=True) + # copied from setup.cfg + yapf_style = dict( + based_on_style="pep8", + blank_line_before_nested_class_or_def=True, + split_before_expression_after_opening_paren=True, + ) + text, _ = FormatCode(text, style_config=yapf_style, verify=True) + + return text + + def __repr__(self): + return f"Config (path: {self.filename}): {self._cfg_dict.__repr__()}" + + def __len__(self): + return len(self._cfg_dict) + + def __getattr__(self, name): + # # debug + # print('+'*15) + # print('name=%s' % name) + # print("addr:", id(self)) + # # print('type(self):', type(self)) + # print(self.__dict__) + # print('+'*15) + # if self.__dict__ == {}: + # raise ValueError + + return getattr(self._cfg_dict, name) + + def __getitem__(self, name): + return self._cfg_dict.__getitem__(name) + + def __setattr__(self, name, value): + if isinstance(value, dict): + value = ConfigDict(value) + self._cfg_dict.__setattr__(name, value) + + def __setitem__(self, name, value): + if isinstance(value, dict): + value = ConfigDict(value) + self._cfg_dict.__setitem__(name, value) + + def __iter__(self): + return iter(self._cfg_dict) + + def dump(self, file=None): + # import ipdb; ipdb.set_trace() + if file is None: + return self.pretty_text + else: + with open(file, "w") as f: + f.write(self.pretty_text) + + def merge_from_dict(self, options): + """Merge list into cfg_dict + + Merge the dict parsed by MultipleKVAction into this cfg. + + Examples: + >>> options = {'model.backbone.depth': 50, + ... 'model.backbone.with_cp':True} + >>> cfg = Config(dict(model=dict(backbone=dict(type='ResNet')))) + >>> cfg.merge_from_dict(options) + >>> cfg_dict = super(Config, self).__getattribute__('_cfg_dict') + >>> assert cfg_dict == dict( + ... model=dict(backbone=dict(depth=50, with_cp=True))) + + Args: + options (dict): dict of configs to merge from. + """ + option_cfg_dict = {} + for full_key, v in options.items(): + d = option_cfg_dict + key_list = full_key.split(".") + for subkey in key_list[:-1]: + d.setdefault(subkey, ConfigDict()) + d = d[subkey] + subkey = key_list[-1] + d[subkey] = v + + cfg_dict = super(SLConfig, self).__getattribute__("_cfg_dict") + super(SLConfig, self).__setattr__( + "_cfg_dict", SLConfig._merge_a_into_b(option_cfg_dict, cfg_dict) + ) + + # for multiprocess + def __setstate__(self, state): + self.__init__(state) + + def copy(self): + return SLConfig(self._cfg_dict.copy()) + + def deepcopy(self): + return SLConfig(self._cfg_dict.deepcopy()) + + +class DictAction(Action): + """ + argparse action to split an argument into KEY=VALUE form + on the first = and append to a dictionary. List options should + be passed as comma separated values, i.e KEY=V1,V2,V3 + """ + + @staticmethod + def _parse_int_float_bool(val): + try: + return int(val) + except ValueError: + pass + try: + return float(val) + except ValueError: + pass + if val.lower() in ["true", "false"]: + return True if val.lower() == "true" else False + if val.lower() in ["none", "null"]: + return None + return val + + def __call__(self, parser, namespace, values, option_string=None): + options = {} + for kv in values: + key, val = kv.split("=", maxsplit=1) + val = [self._parse_int_float_bool(v) for v in val.split(",")] + if len(val) == 1: + val = val[0] + options[key] = val + setattr(namespace, self.dest, options) diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/slio.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/slio.py new file mode 100644 index 0000000000000000000000000000000000000000..72c1f0f7b82cdc931d381feef64fe15815ba657e --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/slio.py @@ -0,0 +1,177 @@ +# ========================================================== +# Modified from mmcv +# ========================================================== + +import json +import pickle +from abc import ABCMeta, abstractmethod +from pathlib import Path + +import yaml + +try: + from yaml import CLoader as Loader, CDumper as Dumper +except ImportError: + from yaml import Loader, Dumper + + +# =========================== +# Rigister handler +# =========================== + + +class BaseFileHandler(metaclass=ABCMeta): + @abstractmethod + def load_from_fileobj(self, file, **kwargs): + pass + + @abstractmethod + def dump_to_fileobj(self, obj, file, **kwargs): + pass + + @abstractmethod + def dump_to_str(self, obj, **kwargs): + pass + + def load_from_path(self, filepath, mode="r", **kwargs): + with open(filepath, mode) as f: + return self.load_from_fileobj(f, **kwargs) + + def dump_to_path(self, obj, filepath, mode="w", **kwargs): + with open(filepath, mode) as f: + self.dump_to_fileobj(obj, f, **kwargs) + + +class JsonHandler(BaseFileHandler): + def load_from_fileobj(self, file): + return json.load(file) + + def dump_to_fileobj(self, obj, file, **kwargs): + json.dump(obj, file, **kwargs) + + def dump_to_str(self, obj, **kwargs): + return json.dumps(obj, **kwargs) + + +class PickleHandler(BaseFileHandler): + def load_from_fileobj(self, file, **kwargs): + return pickle.load(file, **kwargs) + + def load_from_path(self, filepath, **kwargs): + return super(PickleHandler, self).load_from_path(filepath, mode="rb", **kwargs) + + def dump_to_str(self, obj, **kwargs): + kwargs.setdefault("protocol", 2) + return pickle.dumps(obj, **kwargs) + + def dump_to_fileobj(self, obj, file, **kwargs): + kwargs.setdefault("protocol", 2) + pickle.dump(obj, file, **kwargs) + + def dump_to_path(self, obj, filepath, **kwargs): + super(PickleHandler, self).dump_to_path(obj, filepath, mode="wb", **kwargs) + + +class YamlHandler(BaseFileHandler): + def load_from_fileobj(self, file, **kwargs): + kwargs.setdefault("Loader", Loader) + return yaml.load(file, **kwargs) + + def dump_to_fileobj(self, obj, file, **kwargs): + kwargs.setdefault("Dumper", Dumper) + yaml.dump(obj, file, **kwargs) + + def dump_to_str(self, obj, **kwargs): + kwargs.setdefault("Dumper", Dumper) + return yaml.dump(obj, **kwargs) + + +file_handlers = { + "json": JsonHandler(), + "yaml": YamlHandler(), + "yml": YamlHandler(), + "pickle": PickleHandler(), + "pkl": PickleHandler(), +} + +# =========================== +# load and dump +# =========================== + + +def is_str(x): + """Whether the input is an string instance. + + Note: This method is deprecated since python 2 is no longer supported. + """ + return isinstance(x, str) + + +def slload(file, file_format=None, **kwargs): + """Load data from json/yaml/pickle files. + + This method provides a unified api for loading data from serialized files. + + Args: + file (str or :obj:`Path` or file-like object): Filename or a file-like + object. + file_format (str, optional): If not specified, the file format will be + inferred from the file extension, otherwise use the specified one. + Currently supported formats include "json", "yaml/yml" and + "pickle/pkl". + + Returns: + The content from the file. + """ + if isinstance(file, Path): + file = str(file) + if file_format is None and is_str(file): + file_format = file.split(".")[-1] + if file_format not in file_handlers: + raise TypeError(f"Unsupported format: {file_format}") + + handler = file_handlers[file_format] + if is_str(file): + obj = handler.load_from_path(file, **kwargs) + elif hasattr(file, "read"): + obj = handler.load_from_fileobj(file, **kwargs) + else: + raise TypeError('"file" must be a filepath str or a file-object') + return obj + + +def sldump(obj, file=None, file_format=None, **kwargs): + """Dump data to json/yaml/pickle strings or files. + + This method provides a unified api for dumping data as strings or to files, + and also supports custom arguments for each file format. + + Args: + obj (any): The python object to be dumped. + file (str or :obj:`Path` or file-like object, optional): If not + specified, then the object is dump to a str, otherwise to a file + specified by the filename or file-like object. + file_format (str, optional): Same as :func:`load`. + + Returns: + bool: True for success, False otherwise. + """ + if isinstance(file, Path): + file = str(file) + if file_format is None: + if is_str(file): + file_format = file.split(".")[-1] + elif file is None: + raise ValueError("file_format must be specified since file is None") + if file_format not in file_handlers: + raise TypeError(f"Unsupported format: {file_format}") + + handler = file_handlers[file_format] + if file is None: + return handler.dump_to_str(obj, **kwargs) + elif is_str(file): + handler.dump_to_path(obj, file, **kwargs) + elif hasattr(file, "write"): + handler.dump_to_fileobj(obj, file, **kwargs) + else: + raise TypeError('"file" must be a filename str or a file-object') diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/time_counter.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/time_counter.py new file mode 100644 index 0000000000000000000000000000000000000000..0aedb2e4d61bfbe7571dca9d50053f0fedaa1359 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/time_counter.py @@ -0,0 +1,62 @@ +import json +import time + + +class TimeCounter: + def __init__(self) -> None: + pass + + def clear(self): + self.timedict = {} + self.basetime = time.perf_counter() + + def timeit(self, name): + nowtime = time.perf_counter() - self.basetime + self.timedict[name] = nowtime + self.basetime = time.perf_counter() + + +class TimeHolder: + def __init__(self) -> None: + self.timedict = {} + + def update(self, _timedict: dict): + for k, v in _timedict.items(): + if k not in self.timedict: + self.timedict[k] = AverageMeter(name=k, val_only=True) + self.timedict[k].update(val=v) + + def final_res(self): + return {k: v.avg for k, v in self.timedict.items()} + + def __str__(self): + return json.dumps(self.final_res(), indent=2) + + +class AverageMeter(object): + """Computes and stores the average and current value""" + + def __init__(self, name, fmt=":f", val_only=False): + self.name = name + self.fmt = fmt + self.val_only = val_only + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + def __str__(self): + if self.val_only: + fmtstr = "{name} {val" + self.fmt + "}" + else: + fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})" + return fmtstr.format(**self.__dict__) diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/utils.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e9f0318e306fa04bff0ada70486b41aaa69b07c8 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/utils.py @@ -0,0 +1,608 @@ +import argparse +import json +import warnings +from collections import OrderedDict +from copy import deepcopy +from typing import Any, Dict, List + +import numpy as np +import torch +from transformers import AutoTokenizer + +from groundingdino.util.slconfig import SLConfig + + +def slprint(x, name="x"): + if isinstance(x, (torch.Tensor, np.ndarray)): + print(f"{name}.shape:", x.shape) + elif isinstance(x, (tuple, list)): + print("type x:", type(x)) + for i in range(min(10, len(x))): + slprint(x[i], f"{name}[{i}]") + elif isinstance(x, dict): + for k, v in x.items(): + slprint(v, f"{name}[{k}]") + else: + print(f"{name}.type:", type(x)) + + +def clean_state_dict(state_dict): + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + if k[:7] == "module.": + k = k[7:] # remove `module.` + new_state_dict[k] = v + return new_state_dict + + +def renorm( + img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] +) -> torch.FloatTensor: + # img: tensor(3,H,W) or tensor(B,3,H,W) + # return: same as img + assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim() + if img.dim() == 3: + assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % ( + img.size(0), + str(img.size()), + ) + img_perm = img.permute(1, 2, 0) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(2, 0, 1) + else: # img.dim() == 4 + assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % ( + img.size(1), + str(img.size()), + ) + img_perm = img.permute(0, 2, 3, 1) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(0, 3, 1, 2) + + +class CocoClassMapper: + def __init__(self) -> None: + self.category_map_str = { + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "13": 12, + "14": 13, + "15": 14, + "16": 15, + "17": 16, + "18": 17, + "19": 18, + "20": 19, + "21": 20, + "22": 21, + "23": 22, + "24": 23, + "25": 24, + "27": 25, + "28": 26, + "31": 27, + "32": 28, + "33": 29, + "34": 30, + "35": 31, + "36": 32, + "37": 33, + "38": 34, + "39": 35, + "40": 36, + "41": 37, + "42": 38, + "43": 39, + "44": 40, + "46": 41, + "47": 42, + "48": 43, + "49": 44, + "50": 45, + "51": 46, + "52": 47, + "53": 48, + "54": 49, + "55": 50, + "56": 51, + "57": 52, + "58": 53, + "59": 54, + "60": 55, + "61": 56, + "62": 57, + "63": 58, + "64": 59, + "65": 60, + "67": 61, + "70": 62, + "72": 63, + "73": 64, + "74": 65, + "75": 66, + "76": 67, + "77": 68, + "78": 69, + "79": 70, + "80": 71, + "81": 72, + "82": 73, + "84": 74, + "85": 75, + "86": 76, + "87": 77, + "88": 78, + "89": 79, + "90": 80, + } + self.origin2compact_mapper = {int(k): v - 1 for k, v in self.category_map_str.items()} + self.compact2origin_mapper = {int(v - 1): int(k) for k, v in self.category_map_str.items()} + + def origin2compact(self, idx): + return self.origin2compact_mapper[int(idx)] + + def compact2origin(self, idx): + return self.compact2origin_mapper[int(idx)] + + +def to_device(item, device): + if isinstance(item, torch.Tensor): + return item.to(device) + elif isinstance(item, list): + return [to_device(i, device) for i in item] + elif isinstance(item, dict): + return {k: to_device(v, device) for k, v in item.items()} + else: + raise NotImplementedError( + "Call Shilong if you use other containers! type: {}".format(type(item)) + ) + + +# +def get_gaussian_mean(x, axis, other_axis, softmax=True): + """ + + Args: + x (float): Input images(BxCxHxW) + axis (int): The index for weighted mean + other_axis (int): The other index + + Returns: weighted index for axis, BxC + + """ + mat2line = torch.sum(x, axis=other_axis) + # mat2line = mat2line / mat2line.mean() * 10 + if softmax: + u = torch.softmax(mat2line, axis=2) + else: + u = mat2line / (mat2line.sum(2, keepdim=True) + 1e-6) + size = x.shape[axis] + ind = torch.linspace(0, 1, size).to(x.device) + batch = x.shape[0] + channel = x.shape[1] + index = ind.repeat([batch, channel, 1]) + mean_position = torch.sum(index * u, dim=2) + return mean_position + + +def get_expected_points_from_map(hm, softmax=True): + """get_gaussian_map_from_points + B,C,H,W -> B,N,2 float(0, 1) float(0, 1) + softargmax function + + Args: + hm (float): Input images(BxCxHxW) + + Returns: + weighted index for axis, BxCx2. float between 0 and 1. + + """ + # hm = 10*hm + B, C, H, W = hm.shape + y_mean = get_gaussian_mean(hm, 2, 3, softmax=softmax) # B,C + x_mean = get_gaussian_mean(hm, 3, 2, softmax=softmax) # B,C + # return torch.cat((x_mean.unsqueeze(-1), y_mean.unsqueeze(-1)), 2) + return torch.stack([x_mean, y_mean], dim=2) + + +# Positional encoding (section 5.1) +# borrow from nerf +class Embedder: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.create_embedding_fn() + + def create_embedding_fn(self): + embed_fns = [] + d = self.kwargs["input_dims"] + out_dim = 0 + if self.kwargs["include_input"]: + embed_fns.append(lambda x: x) + out_dim += d + + max_freq = self.kwargs["max_freq_log2"] + N_freqs = self.kwargs["num_freqs"] + + if self.kwargs["log_sampling"]: + freq_bands = 2.0 ** torch.linspace(0.0, max_freq, steps=N_freqs) + else: + freq_bands = torch.linspace(2.0**0.0, 2.0**max_freq, steps=N_freqs) + + for freq in freq_bands: + for p_fn in self.kwargs["periodic_fns"]: + embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq)) + out_dim += d + + self.embed_fns = embed_fns + self.out_dim = out_dim + + def embed(self, inputs): + return torch.cat([fn(inputs) for fn in self.embed_fns], -1) + + +def get_embedder(multires, i=0): + import torch.nn as nn + + if i == -1: + return nn.Identity(), 3 + + embed_kwargs = { + "include_input": True, + "input_dims": 3, + "max_freq_log2": multires - 1, + "num_freqs": multires, + "log_sampling": True, + "periodic_fns": [torch.sin, torch.cos], + } + + embedder_obj = Embedder(**embed_kwargs) + embed = lambda x, eo=embedder_obj: eo.embed(x) + return embed, embedder_obj.out_dim + + +class APOPMeter: + def __init__(self) -> None: + self.tp = 0 + self.fp = 0 + self.tn = 0 + self.fn = 0 + + def update(self, pred, gt): + """ + Input: + pred, gt: Tensor() + """ + assert pred.shape == gt.shape + self.tp += torch.logical_and(pred == 1, gt == 1).sum().item() + self.fp += torch.logical_and(pred == 1, gt == 0).sum().item() + self.tn += torch.logical_and(pred == 0, gt == 0).sum().item() + self.tn += torch.logical_and(pred == 1, gt == 0).sum().item() + + def update_cm(self, tp, fp, tn, fn): + self.tp += tp + self.fp += fp + self.tn += tn + self.tn += fn + + +def inverse_sigmoid(x, eps=1e-5): + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + + +def get_raw_dict(args): + """ + return the dicf contained in args. + + e.g: + >>> with open(path, 'w') as f: + json.dump(get_raw_dict(args), f, indent=2) + """ + if isinstance(args, argparse.Namespace): + return vars(args) + elif isinstance(args, dict): + return args + elif isinstance(args, SLConfig): + return args._cfg_dict + else: + raise NotImplementedError("Unknown type {}".format(type(args))) + + +def stat_tensors(tensor): + assert tensor.dim() == 1 + tensor_sm = tensor.softmax(0) + entropy = (tensor_sm * torch.log(tensor_sm + 1e-9)).sum() + + return { + "max": tensor.max(), + "min": tensor.min(), + "mean": tensor.mean(), + "var": tensor.var(), + "std": tensor.var() ** 0.5, + "entropy": entropy, + } + + +class NiceRepr: + """Inherit from this class and define ``__nice__`` to "nicely" print your + objects. + + Defines ``__str__`` and ``__repr__`` in terms of ``__nice__`` function + Classes that inherit from :class:`NiceRepr` should redefine ``__nice__``. + If the inheriting class has a ``__len__``, method then the default + ``__nice__`` method will return its length. + + Example: + >>> class Foo(NiceRepr): + ... def __nice__(self): + ... return 'info' + >>> foo = Foo() + >>> assert str(foo) == '' + >>> assert repr(foo).startswith('>> class Bar(NiceRepr): + ... pass + >>> bar = Bar() + >>> import pytest + >>> with pytest.warns(None) as record: + >>> assert 'object at' in str(bar) + >>> assert 'object at' in repr(bar) + + Example: + >>> class Baz(NiceRepr): + ... def __len__(self): + ... return 5 + >>> baz = Baz() + >>> assert str(baz) == '' + """ + + def __nice__(self): + """str: a "nice" summary string describing this module""" + if hasattr(self, "__len__"): + # It is a common pattern for objects to use __len__ in __nice__ + # As a convenience we define a default __nice__ for these objects + return str(len(self)) + else: + # In all other cases force the subclass to overload __nice__ + raise NotImplementedError(f"Define the __nice__ method for {self.__class__!r}") + + def __repr__(self): + """str: the string of the module""" + try: + nice = self.__nice__() + classname = self.__class__.__name__ + return f"<{classname}({nice}) at {hex(id(self))}>" + except NotImplementedError as ex: + warnings.warn(str(ex), category=RuntimeWarning) + return object.__repr__(self) + + def __str__(self): + """str: the string of the module""" + try: + classname = self.__class__.__name__ + nice = self.__nice__() + return f"<{classname}({nice})>" + except NotImplementedError as ex: + warnings.warn(str(ex), category=RuntimeWarning) + return object.__repr__(self) + + +def ensure_rng(rng=None): + """Coerces input into a random number generator. + + If the input is None, then a global random state is returned. + + If the input is a numeric value, then that is used as a seed to construct a + random state. Otherwise the input is returned as-is. + + Adapted from [1]_. + + Args: + rng (int | numpy.random.RandomState | None): + if None, then defaults to the global rng. Otherwise this can be an + integer or a RandomState class + Returns: + (numpy.random.RandomState) : rng - + a numpy random number generator + + References: + .. [1] https://gitlab.kitware.com/computer-vision/kwarray/blob/master/kwarray/util_random.py#L270 # noqa: E501 + """ + + if rng is None: + rng = np.random.mtrand._rand + elif isinstance(rng, int): + rng = np.random.RandomState(rng) + else: + rng = rng + return rng + + +def random_boxes(num=1, scale=1, rng=None): + """Simple version of ``kwimage.Boxes.random`` + + Returns: + Tensor: shape (n, 4) in x1, y1, x2, y2 format. + + References: + https://gitlab.kitware.com/computer-vision/kwimage/blob/master/kwimage/structs/boxes.py#L1390 + + Example: + >>> num = 3 + >>> scale = 512 + >>> rng = 0 + >>> boxes = random_boxes(num, scale, rng) + >>> print(boxes) + tensor([[280.9925, 278.9802, 308.6148, 366.1769], + [216.9113, 330.6978, 224.0446, 456.5878], + [405.3632, 196.3221, 493.3953, 270.7942]]) + """ + rng = ensure_rng(rng) + + tlbr = rng.rand(num, 4).astype(np.float32) + + tl_x = np.minimum(tlbr[:, 0], tlbr[:, 2]) + tl_y = np.minimum(tlbr[:, 1], tlbr[:, 3]) + br_x = np.maximum(tlbr[:, 0], tlbr[:, 2]) + br_y = np.maximum(tlbr[:, 1], tlbr[:, 3]) + + tlbr[:, 0] = tl_x * scale + tlbr[:, 1] = tl_y * scale + tlbr[:, 2] = br_x * scale + tlbr[:, 3] = br_y * scale + + boxes = torch.from_numpy(tlbr) + return boxes + + +class ModelEma(torch.nn.Module): + def __init__(self, model, decay=0.9997, device=None): + super(ModelEma, self).__init__() + # make a copy of the model for accumulating moving average of weights + self.module = deepcopy(model) + self.module.eval() + + # import ipdb; ipdb.set_trace() + + self.decay = decay + self.device = device # perform ema on different device from model if set + if self.device is not None: + self.module.to(device=device) + + def _update(self, model, update_fn): + with torch.no_grad(): + for ema_v, model_v in zip( + self.module.state_dict().values(), model.state_dict().values() + ): + if self.device is not None: + model_v = model_v.to(device=self.device) + ema_v.copy_(update_fn(ema_v, model_v)) + + def update(self, model): + self._update(model, update_fn=lambda e, m: self.decay * e + (1.0 - self.decay) * m) + + def set(self, model): + self._update(model, update_fn=lambda e, m: m) + + +class BestMetricSingle: + def __init__(self, init_res=0.0, better="large") -> None: + self.init_res = init_res + self.best_res = init_res + self.best_ep = -1 + + self.better = better + assert better in ["large", "small"] + + def isbetter(self, new_res, old_res): + if self.better == "large": + return new_res > old_res + if self.better == "small": + return new_res < old_res + + def update(self, new_res, ep): + if self.isbetter(new_res, self.best_res): + self.best_res = new_res + self.best_ep = ep + return True + return False + + def __str__(self) -> str: + return "best_res: {}\t best_ep: {}".format(self.best_res, self.best_ep) + + def __repr__(self) -> str: + return self.__str__() + + def summary(self) -> dict: + return { + "best_res": self.best_res, + "best_ep": self.best_ep, + } + + +class BestMetricHolder: + def __init__(self, init_res=0.0, better="large", use_ema=False) -> None: + self.best_all = BestMetricSingle(init_res, better) + self.use_ema = use_ema + if use_ema: + self.best_ema = BestMetricSingle(init_res, better) + self.best_regular = BestMetricSingle(init_res, better) + + def update(self, new_res, epoch, is_ema=False): + """ + return if the results is the best. + """ + if not self.use_ema: + return self.best_all.update(new_res, epoch) + else: + if is_ema: + self.best_ema.update(new_res, epoch) + return self.best_all.update(new_res, epoch) + else: + self.best_regular.update(new_res, epoch) + return self.best_all.update(new_res, epoch) + + def summary(self): + if not self.use_ema: + return self.best_all.summary() + + res = {} + res.update({f"all_{k}": v for k, v in self.best_all.summary().items()}) + res.update({f"regular_{k}": v for k, v in self.best_regular.summary().items()}) + res.update({f"ema_{k}": v for k, v in self.best_ema.summary().items()}) + return res + + def __repr__(self) -> str: + return json.dumps(self.summary(), indent=2) + + def __str__(self) -> str: + return self.__repr__() + + +def targets_to(targets: List[Dict[str, Any]], device): + """Moves the target dicts to the given device.""" + excluded_keys = [ + "questionId", + "tokens_positive", + "strings_positive", + "tokens", + "dataset_name", + "sentence_id", + "original_img_id", + "nb_eval", + "task_id", + "original_id", + "token_span", + "caption", + "dataset_type", + ] + return [ + {k: v.to(device) if k not in excluded_keys else v for k, v in t.items()} for t in targets + ] + + +def get_phrases_from_posmap( + posmap: torch.BoolTensor, tokenized: Dict, tokenizer: AutoTokenizer +): + assert isinstance(posmap, torch.Tensor), "posmap must be torch.Tensor" + if posmap.dim() == 1: + non_zero_idx = posmap.nonzero(as_tuple=True)[0].tolist() + token_ids = [tokenized["input_ids"][i] for i in non_zero_idx] + return tokenizer.decode(token_ids) + else: + raise NotImplementedError("posmap must be 1-dim") diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/visualizer.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7a1b7b101e9b73f75f9136bc67f2063c7c1cf1c1 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/visualizer.py @@ -0,0 +1,318 @@ +# -*- coding: utf-8 -*- +""" +@File : visualizer.py +@Time : 2022/04/05 11:39:33 +@Author : Shilong Liu +@Contact : slongliu86@gmail.com +""" + +import datetime +import os + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import torch +from matplotlib import transforms +from matplotlib.collections import PatchCollection +from matplotlib.patches import Polygon +from pycocotools import mask as maskUtils + + +def renorm( + img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] +) -> torch.FloatTensor: + # img: tensor(3,H,W) or tensor(B,3,H,W) + # return: same as img + assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim() + if img.dim() == 3: + assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % ( + img.size(0), + str(img.size()), + ) + img_perm = img.permute(1, 2, 0) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(2, 0, 1) + else: # img.dim() == 4 + assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % ( + img.size(1), + str(img.size()), + ) + img_perm = img.permute(0, 2, 3, 1) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(0, 3, 1, 2) + + +class ColorMap: + def __init__(self, basergb=[255, 255, 0]): + self.basergb = np.array(basergb) + + def __call__(self, attnmap): + # attnmap: h, w. np.uint8. + # return: h, w, 4. np.uint8. + assert attnmap.dtype == np.uint8 + h, w = attnmap.shape + res = self.basergb.copy() + res = res[None][None].repeat(h, 0).repeat(w, 1) # h, w, 3 + attn1 = attnmap.copy()[..., None] # h, w, 1 + res = np.concatenate((res, attn1), axis=-1).astype(np.uint8) + return res + + +def rainbow_text(x, y, ls, lc, **kw): + """ + Take a list of strings ``ls`` and colors ``lc`` and place them next to each + other, with text ls[i] being shown in color lc[i]. + + This example shows how to do both vertical and horizontal text, and will + pass all keyword arguments to plt.text, so you can set the font size, + family, etc. + """ + t = plt.gca().transData + fig = plt.gcf() + plt.show() + + # horizontal version + for s, c in zip(ls, lc): + text = plt.text(x, y, " " + s + " ", color=c, transform=t, **kw) + text.draw(fig.canvas.get_renderer()) + ex = text.get_window_extent() + t = transforms.offset_copy(text._transform, x=ex.width, units="dots") + + # #vertical version + # for s,c in zip(ls,lc): + # text = plt.text(x,y," "+s+" ",color=c, transform=t, + # rotation=90,va='bottom',ha='center',**kw) + # text.draw(fig.canvas.get_renderer()) + # ex = text.get_window_extent() + # t = transforms.offset_copy(text._transform, y=ex.height, units='dots') + + +class COCOVisualizer: + def __init__(self, coco=None, tokenlizer=None) -> None: + self.coco = coco + + def visualize(self, img, tgt, caption=None, dpi=180, savedir="vis"): + """ + img: tensor(3, H, W) + tgt: make sure they are all on cpu. + must have items: 'image_id', 'boxes', 'size' + """ + plt.figure(dpi=dpi) + plt.rcParams["font.size"] = "5" + ax = plt.gca() + img = renorm(img).permute(1, 2, 0) + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + ax.imshow(img) + + self.addtgt(tgt) + + if tgt is None: + image_id = 0 + elif "image_id" not in tgt: + image_id = 0 + else: + image_id = tgt["image_id"] + + if caption is None: + savename = "{}/{}-{}.png".format( + savedir, int(image_id), str(datetime.datetime.now()).replace(" ", "-") + ) + else: + savename = "{}/{}-{}-{}.png".format( + savedir, caption, int(image_id), str(datetime.datetime.now()).replace(" ", "-") + ) + print("savename: {}".format(savename)) + os.makedirs(os.path.dirname(savename), exist_ok=True) + plt.savefig(savename) + plt.close() + + def addtgt(self, tgt): + """ """ + if tgt is None or not "boxes" in tgt: + ax = plt.gca() + + if "caption" in tgt: + ax.set_title(tgt["caption"], wrap=True) + + ax.set_axis_off() + return + + ax = plt.gca() + H, W = tgt["size"] + numbox = tgt["boxes"].shape[0] + + color = [] + polygons = [] + boxes = [] + for box in tgt["boxes"].cpu(): + unnormbbox = box * torch.Tensor([W, H, W, H]) + unnormbbox[:2] -= unnormbbox[2:] / 2 + [bbox_x, bbox_y, bbox_w, bbox_h] = unnormbbox.tolist() + boxes.append([bbox_x, bbox_y, bbox_w, bbox_h]) + poly = [ + [bbox_x, bbox_y], + [bbox_x, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y], + ] + np_poly = np.array(poly).reshape((4, 2)) + polygons.append(Polygon(np_poly)) + c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] + color.append(c) + + p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.1) + ax.add_collection(p) + p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2) + ax.add_collection(p) + + if "strings_positive" in tgt and len(tgt["strings_positive"]) > 0: + assert ( + len(tgt["strings_positive"]) == numbox + ), f"{len(tgt['strings_positive'])} = {numbox}, " + for idx, strlist in enumerate(tgt["strings_positive"]): + cate_id = int(tgt["labels"][idx]) + _string = str(cate_id) + ":" + " ".join(strlist) + bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx] + # ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1}) + ax.text( + bbox_x, + bbox_y, + _string, + color="black", + bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1}, + ) + + if "box_label" in tgt: + assert len(tgt["box_label"]) == numbox, f"{len(tgt['box_label'])} = {numbox}, " + for idx, bl in enumerate(tgt["box_label"]): + _string = str(bl) + bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx] + # ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1}) + ax.text( + bbox_x, + bbox_y, + _string, + color="black", + bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1}, + ) + + if "caption" in tgt: + ax.set_title(tgt["caption"], wrap=True) + # plt.figure() + # rainbow_text(0.0,0.0,"all unicorns poop rainbows ! ! !".split(), + # ['red', 'orange', 'brown', 'green', 'blue', 'purple', 'black']) + + if "attn" in tgt: + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + if isinstance(tgt["attn"], tuple): + tgt["attn"] = [tgt["attn"]] + for item in tgt["attn"]: + attn_map, basergb = item + attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-3) + attn_map = (attn_map * 255).astype(np.uint8) + cm = ColorMap(basergb) + heatmap = cm(attn_map) + ax.imshow(heatmap) + ax.set_axis_off() + + def showAnns(self, anns, draw_bbox=False): + """ + Display the specified annotations. + :param anns (array of object): annotations to display + :return: None + """ + if len(anns) == 0: + return 0 + if "segmentation" in anns[0] or "keypoints" in anns[0]: + datasetType = "instances" + elif "caption" in anns[0]: + datasetType = "captions" + else: + raise Exception("datasetType not supported") + if datasetType == "instances": + ax = plt.gca() + ax.set_autoscale_on(False) + polygons = [] + color = [] + for ann in anns: + c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] + if "segmentation" in ann: + if type(ann["segmentation"]) == list: + # polygon + for seg in ann["segmentation"]: + poly = np.array(seg).reshape((int(len(seg) / 2), 2)) + polygons.append(Polygon(poly)) + color.append(c) + else: + # mask + t = self.imgs[ann["image_id"]] + if type(ann["segmentation"]["counts"]) == list: + rle = maskUtils.frPyObjects( + [ann["segmentation"]], t["height"], t["width"] + ) + else: + rle = [ann["segmentation"]] + m = maskUtils.decode(rle) + img = np.ones((m.shape[0], m.shape[1], 3)) + if ann["iscrowd"] == 1: + color_mask = np.array([2.0, 166.0, 101.0]) / 255 + if ann["iscrowd"] == 0: + color_mask = np.random.random((1, 3)).tolist()[0] + for i in range(3): + img[:, :, i] = color_mask[i] + ax.imshow(np.dstack((img, m * 0.5))) + if "keypoints" in ann and type(ann["keypoints"]) == list: + # turn skeleton into zero-based index + sks = np.array(self.loadCats(ann["category_id"])[0]["skeleton"]) - 1 + kp = np.array(ann["keypoints"]) + x = kp[0::3] + y = kp[1::3] + v = kp[2::3] + for sk in sks: + if np.all(v[sk] > 0): + plt.plot(x[sk], y[sk], linewidth=3, color=c) + plt.plot( + x[v > 0], + y[v > 0], + "o", + markersize=8, + markerfacecolor=c, + markeredgecolor="k", + markeredgewidth=2, + ) + plt.plot( + x[v > 1], + y[v > 1], + "o", + markersize=8, + markerfacecolor=c, + markeredgecolor=c, + markeredgewidth=2, + ) + + if draw_bbox: + [bbox_x, bbox_y, bbox_w, bbox_h] = ann["bbox"] + poly = [ + [bbox_x, bbox_y], + [bbox_x, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y], + ] + np_poly = np.array(poly).reshape((4, 2)) + polygons.append(Polygon(np_poly)) + color.append(c) + + # p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4) + # ax.add_collection(p) + p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2) + ax.add_collection(p) + elif datasetType == "captions": + for ann in anns: + print(ann["caption"]) diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/vl_utils.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/vl_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c91bb02f584398f08a28e6b7719e2b99f6e28616 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/util/vl_utils.py @@ -0,0 +1,100 @@ +import os +import random +from typing import List + +import torch + + +def create_positive_map_from_span(tokenized, token_span, max_text_len=256): + """construct a map such that positive_map[i,j] = True iff box i is associated to token j + Input: + - tokenized: + - input_ids: Tensor[1, ntokens] + - attention_mask: Tensor[1, ntokens] + - token_span: list with length num_boxes. + - each item: [start_idx, end_idx] + """ + positive_map = torch.zeros((len(token_span), max_text_len), dtype=torch.float) + for j, tok_list in enumerate(token_span): + for (beg, end) in tok_list: + beg_pos = tokenized.char_to_token(beg) + end_pos = tokenized.char_to_token(end - 1) + if beg_pos is None: + try: + beg_pos = tokenized.char_to_token(beg + 1) + if beg_pos is None: + beg_pos = tokenized.char_to_token(beg + 2) + except: + beg_pos = None + if end_pos is None: + try: + end_pos = tokenized.char_to_token(end - 2) + if end_pos is None: + end_pos = tokenized.char_to_token(end - 3) + except: + end_pos = None + if beg_pos is None or end_pos is None: + continue + + assert beg_pos is not None and end_pos is not None + if os.environ.get("SHILONG_DEBUG_ONLY_ONE_POS", None) == "TRUE": + positive_map[j, beg_pos] = 1 + break + else: + positive_map[j, beg_pos : end_pos + 1].fill_(1) + + return positive_map / (positive_map.sum(-1)[:, None] + 1e-6) + + +def build_captions_and_token_span(cat_list, force_lowercase): + """ + Return: + captions: str + cat2tokenspan: dict + { + 'dog': [[0, 2]], + ... + } + """ + + cat2tokenspan = {} + captions = "" + for catname in cat_list: + class_name = catname + if force_lowercase: + class_name = class_name.lower() + if "/" in class_name: + class_name_list: List = class_name.strip().split("/") + class_name_list.append(class_name) + class_name: str = random.choice(class_name_list) + + tokens_positive_i = [] + subnamelist = [i.strip() for i in class_name.strip().split(" ")] + for subname in subnamelist: + if len(subname) == 0: + continue + if len(captions) > 0: + captions = captions + " " + strat_idx = len(captions) + end_idx = strat_idx + len(subname) + tokens_positive_i.append([strat_idx, end_idx]) + captions = captions + subname + + if len(tokens_positive_i) > 0: + captions = captions + " ." + cat2tokenspan[class_name] = tokens_positive_i + + return captions, cat2tokenspan + + +def build_id2posspan_and_caption(category_dict: dict): + """Build id2pos_span and caption from category_dict + + Args: + category_dict (dict): category_dict + """ + cat_list = [item["name"].lower() for item in category_dict] + id2catname = {item["id"]: item["name"].lower() for item in category_dict} + caption, cat2posspan = build_captions_and_token_span(cat_list, force_lowercase=True) + id2posspan = {catid: cat2posspan[catname] for catid, catname in id2catname.items()} + return id2posspan, caption diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/version.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/version.py new file mode 100644 index 0000000000000000000000000000000000000000..b794fd409a5e3b3b65ad76a43d6a01a318877640 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/groundingdino/version.py @@ -0,0 +1 @@ +__version__ = '0.1.0' diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/pyproject.toml b/ArtiAgent - DefectDiffu/src/GroundingDINO/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..24dcc68d94ea5aaee6bb7a903a0e1638cf14e6b1 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/pyproject.toml @@ -0,0 +1,8 @@ +[build-system] +requires = [ + "setuptools", + "torch", + "wheel", + "torch" +] +build-backend = "setuptools.build_meta" diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/requirements.txt b/ArtiAgent - DefectDiffu/src/GroundingDINO/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..32e2538a959167f9ce248a5c99cf275bc53cc51b --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/requirements.txt @@ -0,0 +1,12 @@ +torch +torchvision +transformers +addict +yapf +timm +numpy +opencv-python +supervision==0.21.0 +pycocotools +lpips +openai \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/GroundingDINO/setup.py b/ArtiAgent - DefectDiffu/src/GroundingDINO/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..a58340d44eca86b09cb69630465dfbdfe8acb742 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/GroundingDINO/setup.py @@ -0,0 +1,216 @@ +# coding=utf-8 +# Copyright 2022 The IDEA Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ------------------------------------------------------------------------------------------------ +# Modified from +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/setup.py +# https://github.com/facebookresearch/detectron2/blob/main/setup.py +# https://github.com/open-mmlab/mmdetection/blob/master/setup.py +# https://github.com/Oneflow-Inc/libai/blob/main/setup.py +# ------------------------------------------------------------------------------------------------ + +import glob +import os +import subprocess + +import torch +from setuptools import find_packages, setup +from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension + +# groundingdino version info +version = "0.1.0" +package_name = "groundingdino" +cwd = os.path.dirname(os.path.abspath(__file__)) + + +sha = "Unknown" +try: + sha = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=cwd).decode("ascii").strip() +except Exception: + pass + + +def write_version_file(): + version_path = os.path.join(cwd, "groundingdino", "version.py") + with open(version_path, "w") as f: + f.write(f"__version__ = '{version}'\n") + # f.write(f"git_version = {repr(sha)}\n") + + +requirements = ["torch", "torchvision"] + +torch_ver = [int(x) for x in torch.__version__.split(".")[:2]] + + +def get_extensions(): + this_dir = os.path.dirname(os.path.abspath(__file__)) + extensions_dir = os.path.join(this_dir, "groundingdino", "models", "GroundingDINO", "csrc") + + main_source = os.path.join(extensions_dir, "vision.cpp") + sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp")) + source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu")) + glob.glob( + os.path.join(extensions_dir, "*.cu") + ) + + sources = [main_source] + sources + + # We need these variables to build with CUDA when we create the Docker image + # It solves https://github.com/IDEA-Research/Grounded-Segment-Anything/issues/53 + # and https://github.com/IDEA-Research/Grounded-Segment-Anything/issues/84 when running + # inside a Docker container. + am_i_docker = os.environ.get('AM_I_DOCKER', '').casefold() in ['true', '1', 't'] + use_cuda = os.environ.get('BUILD_WITH_CUDA', '').casefold() in ['true', '1', 't'] + + extension = CppExtension + + extra_compile_args = {"cxx": []} + define_macros = [] + + if (torch.cuda.is_available() and CUDA_HOME is not None) or \ + (am_i_docker and use_cuda): + print("Compiling with CUDA") + extension = CUDAExtension + sources += source_cuda + define_macros += [("WITH_CUDA", None)] + extra_compile_args["nvcc"] = [ + "-DCUDA_HAS_FP16=1", + "-D__CUDA_NO_HALF_OPERATORS__", + "-D__CUDA_NO_HALF_CONVERSIONS__", + "-D__CUDA_NO_HALF2_OPERATORS__", + ] + else: + print("Compiling without CUDA") + define_macros += [("WITH_HIP", None)] + extra_compile_args["nvcc"] = [] + return None + + sources = [os.path.join(extensions_dir, s) for s in sources] + include_dirs = [extensions_dir] + + ext_modules = [ + extension( + "groundingdino._C", + sources, + include_dirs=include_dirs, + define_macros=define_macros, + extra_compile_args=extra_compile_args, + ) + ] + + return ext_modules + + +def parse_requirements(fname="requirements.txt", with_version=True): + """Parse the package dependencies listed in a requirements file but strips + specific versioning information. + + Args: + fname (str): path to requirements file + with_version (bool, default=False): if True include version specs + + Returns: + List[str]: list of requirements items + + CommandLine: + python -c "import setup; print(setup.parse_requirements())" + """ + import re + import sys + from os.path import exists + + require_fpath = fname + + def parse_line(line): + """Parse information from a line in a requirements text file.""" + if line.startswith("-r "): + # Allow specifying requirements in other files + target = line.split(" ")[1] + for info in parse_require_file(target): + yield info + else: + info = {"line": line} + if line.startswith("-e "): + info["package"] = line.split("#egg=")[1] + elif "@git+" in line: + info["package"] = line + else: + # Remove versioning from the package + pat = "(" + "|".join([">=", "==", ">"]) + ")" + parts = re.split(pat, line, maxsplit=1) + parts = [p.strip() for p in parts] + + info["package"] = parts[0] + if len(parts) > 1: + op, rest = parts[1:] + if ";" in rest: + # Handle platform specific dependencies + # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies + version, platform_deps = map(str.strip, rest.split(";")) + info["platform_deps"] = platform_deps + else: + version = rest # NOQA + info["version"] = (op, version) + yield info + + def parse_require_file(fpath): + with open(fpath, "r") as f: + for line in f.readlines(): + line = line.strip() + if line and not line.startswith("#"): + for info in parse_line(line): + yield info + + def gen_packages_items(): + if exists(require_fpath): + for info in parse_require_file(require_fpath): + parts = [info["package"]] + if with_version and "version" in info: + parts.extend(info["version"]) + if not sys.version.startswith("3.4"): + # apparently package_deps are broken in 3.4 + platform_deps = info.get("platform_deps") + if platform_deps is not None: + parts.append(";" + platform_deps) + item = "".join(parts) + yield item + + packages = list(gen_packages_items()) + return packages + + +if __name__ == "__main__": + print(f"Building wheel {package_name}-{version}") + + with open("LICENSE", "r", encoding="utf-8") as f: + license = f.read() + + write_version_file() + + setup( + name="groundingdino", + version="0.1.0", + author="International Digital Economy Academy, Shilong Liu", + url="https://github.com/IDEA-Research/GroundingDINO", + description="open-set object detector", + license=license, + install_requires=parse_requirements("requirements.txt"), + packages=find_packages( + exclude=( + "configs", + "tests", + ) + ), + ext_modules=get_extensions(), + cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, + ) diff --git a/ArtiAgent - DefectDiffu/src/__pycache__/artiagent_orchestrator.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/__pycache__/artiagent_orchestrator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aaa244be4faf6f2087c37a38c7d3fec2d6f6ffff Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/__pycache__/artiagent_orchestrator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/artiagent_orchestrator.py b/ArtiAgent - DefectDiffu/src/artiagent_orchestrator.py new file mode 100644 index 0000000000000000000000000000000000000000..8fd18c63316874841c557d0c4e2975287e219fea --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/artiagent_orchestrator.py @@ -0,0 +1,777 @@ +""" +ArtiAgent Orchestrator โ€” DefectDiffu Edition + +Adapted from the FLUX-based pipeline to use DefectDiffu (ECCV 2024). + +KEY ARCHITECTURAL CHANGES vs. FLUX version: + 1. DefectDiffu generates NEW images from noise (text-to-image), rather than + editing an existing image via inversion-injection. + 2. Three disentangled text prompts drive generation: + c_p = background/product consistency prompt + c_d = defect consistency prompt + c_f = fusion prompt + 3. Double-free strategy controls defect strength (w_d) and product fidelity (w_p). + 4. Masks are generated automatically from defect-block cross-attention maps. + 5. Patch-based artifact mappings (16x16 FLUX patches) are REMOVED entirely. + 6. The input "clean image" is used for PLANNING and VERIFICATION only. + +Usage: + python artiagent_orchestrator.py \\ + --product-desc "VCSEL laser diode with glass lens cap" \\ + --image ./clean_chip.png \\ + --output-dir ./defect_output \\ + --defectdiffu-ckpt ./defectdiffu_ckpt.pt \\ + --vae-path ./sd-vae-ft-mse \\ + --device cuda +""" + +import os +import sys +import json +import argparse +import uuid +import traceback +from pathlib import Path +from typing import Dict, List, Optional, Tuple +from datetime import datetime + +import numpy as np +import torch +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from pipeline.local_vlm_client import LocalVLMClient +from pipeline.prompts import ( + plan_defects_for_product, + artifact_description, + MoneyManager +) +from pipeline.gsam_detector import GSAMDetector +from pipeline.defectdiffu_generator import DefectDiffuGenerator, DefectDiffuConfig # NEW +from pipeline.instance_processor import InstanceProcessor +from pipeline.defect_rag import get_rag +from pipeline.domain_router import get_router + +import cv2 + +def blend_defect_onto_real_image( + real_image: np.ndarray, # Clean real VCSEL factory photo [H, W, 3] + defect_image: Image.Image, # Pure DefectDiffu generated output (PIL) + defect_mask: np.ndarray, # Binary mask from DefectDiffu [512, 512] + target_bbox: List[int], # [x1, y1, x2, y2] from VLM perception + max_defect_ratio: Optional[float] = None, # None or >= 1.0 means 100% full ROI coverage + mask_shape: str = "free" # "circle", "square", "rectangle", or "free" +) -> Tuple[np.ndarray, np.ndarray]: + """ + Injects a DefectDiffu defect patch onto a real clean factory image at + target_bbox using Poisson Seamless Cloning (Method 2). + """ + defect_np = np.array(defect_image) + mask_uint8 = (defect_mask.astype(np.uint8) * 255) if defect_mask.dtype == bool else defect_mask.astype(np.uint8) + + # 1. Crop tight patch around the generated defect mask + ys, xs = np.where(mask_uint8 > 0) + if len(ys) == 0 or len(xs) == 0: + return real_image.copy(), np.zeros(real_image.shape[:2], dtype=np.uint8) + + y1_d, y2_d = ys.min(), ys.max() + x1_d, x2_d = xs.min(), xs.max() + + defect_patch = defect_np[y1_d:y2_d + 1, x1_d:x2_d + 1] + mask_patch = mask_uint8[y1_d:y2_d + 1, x1_d:x2_d + 1] + + # 2. Extract VLM target ROI dimensions + x1_t, y1_t, x2_t, y2_t = target_bbox + target_w = max(1, x2_t - x1_t) + target_h = max(1, y2_t - y1_t) + + # 3. Dynamic Sizing Logic + # max_defect_ratio = 1.0 + if max_defect_ratio < 0.2: + max_defect_ratio = 0.2 + + if max_defect_ratio is None or max_defect_ratio >= 1.0: + # OPTION A: 100% Full target ROI fit (for smudges, contamination, large scratches) + final_w = target_w + final_h = target_h + else: + # OPTION B: Scaled down relative to target ROI (for bubbles, pinholes, particles) + scale_factor = np.sqrt(max_defect_ratio) + scaled_w = int(target_w * scale_factor) + scaled_h = int(target_h * scale_factor) + + patch_h, patch_w = defect_patch.shape[:2] + aspect_ratio = patch_w / max(1, patch_h) + + if aspect_ratio > 1: + final_w = max(15, scaled_w) + final_h = max(15, int(final_w / aspect_ratio)) + else: + final_h = max(15, scaled_h) + final_w = max(15, int(final_h * aspect_ratio)) + + # 4. Resize patch and mask to fit VLM bounding box + defect_patch_resized = cv2.resize(defect_patch, (final_w, final_h), interpolation=cv2.INTER_AREA) + mask_patch_resized = cv2.resize(mask_patch, (final_w, final_h), interpolation=cv2.INTER_NEAREST) + + # ========================================================================= + # ๐ŸŽจ VLM DYNAMIC MASK SHAPE GENERATION + # ========================================================================= + # shape_type = mask_shape.lower().strip() + shape_type = "free" + + if shape_type == "circle": + # Draw a perfect filled circle + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + center = (final_w // 2, final_h // 2) + radius = max(1, min(final_w, final_h) // 2 - 1) + cv2.circle(geom_mask, center, radius, 255, thickness=-1) + mask_patch_resized = geom_mask + + elif shape_type == "square": + # Draw a centered square + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + side = max(1, min(final_w, final_h) - 2) + top_left_x = (final_w - side) // 2 + top_left_y = (final_h - side) // 2 + cv2.rectangle( + geom_mask, + (top_left_x, top_left_y), + (top_left_x + side, top_left_y + side), + 255, + thickness=-1 + ) + mask_patch_resized = geom_mask + + elif shape_type == "rectangle": + # Fill the entire patch as a solid rectangle + mask_patch_resized = np.full((final_h, final_w), 255, dtype=np.uint8) + + elif shape_type == "free" or shape_type == "irregular": + # Keep original organic AI-generated mask from DefectDiffu + pass + # ========================================================================= + + # 5. Calculate center point for cv2.seamlessClone + center_x = x1_t + target_w // 2 + center_y = y1_t + target_h // 2 + center = (center_x, center_y) + + # 6. Convert RGB -> BGR for OpenCV Poisson Blending + real_bgr = cv2.cvtColor(real_image, cv2.COLOR_RGB2BGR) + patch_bgr = cv2.cvtColor(defect_patch_resized, cv2.COLOR_RGB2BGR) + + # Inside blend_defect_onto_real_image: + patch_mean = np.mean(defect_patch_resized) + + if patch_mean < 30: + # Use NORMAL_CLONE for dark/subtle features to prevent Poisson smoothing from erasing them + clone_mode = cv2.NORMAL_CLONE + else: + clone_mode = cv2.MIXED_CLONE # Preserves underlying substrate structure while injecting defect texture + + blended_bgr = cv2.seamlessClone( + patch_bgr, + real_bgr, + mask_patch_resized, + center, + clone_mode + ) + blended_rgb = cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB) + + # 7. Map binary mask to full real image resolution for segmentation ground-truth + full_mask = np.zeros(real_image.shape[:2], dtype=np.uint8) + top_left_x = max(0, center_x - final_w // 2) + top_left_y = max(0, center_y - final_h // 2) + + h_end = min(real_image.shape[0], top_left_y + final_h) + w_end = min(real_image.shape[1], top_left_x + final_w) + + mask_crop_h = h_end - top_left_y + mask_crop_w = w_end - top_left_x + + if mask_crop_h > 0 and mask_crop_w > 0: + full_mask[top_left_y:h_end, top_left_x:w_end] = ( + mask_patch_resized[:mask_crop_h, :mask_crop_w] > 128 + ).astype(np.uint8) + + return blended_rgb, full_mask + +def create_visual_prompt_image(full_image: np.ndarray, bbox: list) -> np.ndarray: + """Draws a bright neon bounding box on the full image around the target ROI.""" + viz_img = full_image.copy() + x1, y1, x2, y2 = bbox + + # Draw a 2px bright neon green or red rectangle + cv2.rectangle(viz_img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2) + return viz_img + +class ArtiAgentOrchestrator: + """Agentic orchestrator for directed defect generation with DefectDiffu.""" + + def __init__( + self, + device='cuda', + output_dir='./defect_output', + vlm_model='gemma3:12b', + defectdiffu_ckpt: str = "", + vae_path: str = "", + image_size: int = 512, + num_steps: int = 50 + ): + self.device = device + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.vlm_client = LocalVLMClient(model=vlm_model) + self.money_manager = MoneyManager(model="gpt-4o") + + self.gsam_detector = None + self.defectdiffu_generator = None + + # DefectDiffu config + self.defectdiffu_ckpt = defectdiffu_ckpt + self.vae_path = vae_path + self.image_size = image_size + self.num_steps = num_steps + + self.rag = get_rag() + self.router = get_router() + + # ------------------------------------------------------------------ + # Lazy initializers + # ------------------------------------------------------------------ + + def _init_gsam(self): + if self.gsam_detector is None or getattr(self.gsam_detector, 'sam_predictor', None) is None: + print("[Agent] Initializing GSAM detector...") + self.gsam_detector = GSAMDetector( + device=self.device, + openai_client=self.vlm_client + ) + + def _init_defectdiffu(self): + if self.defectdiffu_generator is None: + print("[Agent] Initializing DefectDiffu generator...") + config = DefectDiffuConfig( + ckpt_path=self.defectdiffu_ckpt, + vae_path=self.vae_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + seed=42 + ) + self.defectdiffu_generator = DefectDiffuGenerator(config) + + # ------------------------------------------------------------------ + # Step 1: Planning + # ------------------------------------------------------------------ + + def plan(self, product_description: str, image: np.ndarray, + defect_type: Optional[str] = None, max_defects: int = 3): + """Agent plans defects based on product knowledge.""" + print(f"\\n{'='*60}") + print("[Agent] Step 1: Planning defects from product description...") + print(f"Product: {product_description}") + if defect_type: + print(f"[Agent] Target user defect type requested: {defect_type}") + + plan = plan_defects_for_product( + self.vlm_client, + product_description, + image, + money_manager=self.money_manager, + target_defect_type=defect_type, + max_defects=max_defects + ) + + if plan is None: + raise RuntimeError("Defect planning failed") + + print(f"[Agent] Product type: {plan.product_type}") + print(f"[Agent] Analysis: {plan.analysis}") + print(f"[Agent] Proposed {len(plan.possible_defects)} defects:") + for i, d in enumerate(plan.possible_defects, 1): + print(f" {i}. [{d.defect_type.upper()}] {d.description}") + print(f" Target: {d.target_entity} / {d.target_subentity or '(whole)'}" + f" | Location: {d.location_hint}") + if hasattr(d, 'c_p'): + print(f" c_p: {d.c_p}") + print(f" c_d: {d.c_d}") + print(f" w_d: {d.w_d}") + + return plan + + # ------------------------------------------------------------------ + # Step 2: Perception (optional โ€” for verification bbox only) + # ------------------------------------------------------------------ + + def perceive(self, image: np.ndarray, defect_plan): + """ + Directed perception โ€” detect target entity for verification cropping. + With DefectDiffu this is OPTIONAL; the generator does not need patches. + We keep it to obtain a bbox for the VLM verification step. + """ + print(f"\\n{'='*60}") + print("[Agent] Step 2: Directed perception (verification bbox)...") + + self._init_gsam() + + entity = defect_plan.target_entity + synonym_map = { + "metal can package": ["TO-can", "metal can", "can body", "package body", "metal ring"], + "lens cap": ["glass lens cap", "lens", "glass dome", "optical window"], + "electrode bars": ["vertical bars", "electrodes", "metal lines"], + } + search_terms = [entity] + synonym_map.get(entity.lower(), []) + + predictions = [] + for term in search_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[defect_plan.target_subentity] if defect_plan.target_subentity else [], + entity_subentity_mapping={}, + min_area_ratio=0.005, + max_area_ratio=0.5, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + if term != entity: + print(f"[Agent] Fallback: detected '{term}' instead of '{entity}'") + break + + if not predictions: + print(f"[Agent] Warning: No detections for {entity}; verification will use full image") + h, w = image.shape[:2] + best_pred = { + 'bbox': [0, 0, w, h], + 'pred_mask': torch.ones((h, w), dtype=torch.bool) + } + if predictions: + best_pred = max(predictions, key=lambda p: p.get('area_ratio', 0)) + bbox = best_pred['bbox'] + h, w = image.shape[:2] + + # Check if detected bbox center is too close to border (< 10% margin) + cx = (bbox[0] + bbox[2]) / 2 + cy = (bbox[1] + bbox[3]) / 2 + if cx < 0.1 * w or cx > 0.9 * w or cy < 0.1 * h or cy > 0.9 * h: + print(f"[Agent] Warning: Bbox {bbox} on edge. Falling back to image center.") + best_pred['bbox'] = [int(0.25 * w), int(0.25 * h), int(0.75 * w), int(0.75 * h)] + + return best_pred + + # ------------------------------------------------------------------ + # Step 3: Prepare DefectDiffu generation conditions + # ------------------------------------------------------------------ + + def prepare_generation_conditions( + self, + defect_plan, + product_description: str, + rag_k: int = 3 + ) -> Dict: + """ + Build the three DefectDiffu text prompts + double-free scales. + Retrieves RAG examples to enrich the defect prompt c_d and set w_d. + """ + print(f"\\n{'='*60}") + print("[Agent] Step 3: Preparing DefectDiffu generation conditions...") + + domain = self.router.route(product_description) + + # --- RAG: Retrieve in-context defect examples --- + rag_examples = [] + try: + # DefectRAG expects an object with artifact_type, description, target_entity + rag_query_obj = type('RAGQuery', (), { + 'artifact_type': defect_plan.defect_type, + 'description': defect_plan.description, + 'target_entity': defect_plan.target_entity + })() + rag_examples = self.rag.retrieve( + rag_query_obj, + k=rag_k, + domain_filter=domain if domain != "general" else None, + commercial_only=True + ) + if rag_examples: + print(f"[RAG] Retrieved {len(rag_examples)} example(s) for '{defect_plan.description}'") + for ex in rag_examples: + print(f" โ†’ {ex.get('defect_name', 'unknown')} ({ex.get('domain', 'unknown')})") + else: + print(f"[RAG] No examples found for '{defect_plan.description}'") + except Exception as e: + print(f"[RAG] Retrieval failed: {e}") + + # --- Build prompts --- + # c_p: product / background consistency + c_p = f"A photo of {product_description}" + + # c_d: defect consistency โ€” enrich with RAG if available + clean_defect_plan_description = defect_plan.description.replace("A photo of ", "") + defect_name = clean_defect_plan_description + if rag_examples: + # Use the most similar example's caption to enrich + rag_defect_desc = rag_examples[0].get('caption', '') + if rag_defect_desc: + defect_name = f"{clean_defect_plan_description}, {rag_defect_desc}" + c_d = f"A photo of {defect_name}" + + # c_f: fusion prompt + c_f = f"A photo of {product_description} with {clean_defect_plan_description}" + + # --- Double-free scales --- + severity_to_wd = {"low": 0.6, "minor": 0.6, + "medium": 1.0, "moderate": 1.0, + "high": 1.5, "severe": 1.5} + w_d = severity_to_wd.get(getattr(defect_plan, 'severity', 'medium').lower(), 1.0) + + # If RAG suggests a strength adjustment, apply it + for ex in rag_examples: + meta_wd = ex.get('metadata', {}).get('recommended_wd') + if meta_wd is not None: + w_d = float(meta_wd) + print(f"[RAG] Adjusted w_d to {w_d} from retrieved example") + break + + w_p = 1.0 # Default product consistency; increase if background drifts + + print(f"[Agent] c_p: {c_p}") + print(f"[Agent] c_d: {c_d}") + print(f"[Agent] c_f: {c_f}") + print(f"[Agent] w_d={w_d}, w_p={w_p}") + + return { + 'c_p': c_p, + 'c_d': c_d, + 'c_f': c_f, + 'w_d': w_d, + 'w_p': w_p, + 'rag_examples': rag_examples, + 'rag_domain': domain, + 'defect_plan': defect_plan + } + + # ------------------------------------------------------------------ + # Step 4: Synthesize with DefectDiffu + # ------------------------------------------------------------------ + + def synthesize(self, gen_conditions: Dict) -> Tuple[Image.Image, np.ndarray]: + """Generate defect image + mask with DefectDiffu.""" + print(f"\\n{'='*60}") + print("[Agent] Step 4: Synthesizing defect with DefectDiffu...") + + self._init_defectdiffu() + + img, mask, meta = self.defectdiffu_generator.generate_from_plan( + product_description=gen_conditions['c_p'].replace("A photo of ", ""), + defect_description=gen_conditions['c_d'].replace("A photo of ", ""), + w_d=gen_conditions['w_d'], + w_p=gen_conditions['w_p'], + seed=42 + ) + + print("[Agent] DefectDiffu synthesis complete") + return img, mask + + # ------------------------------------------------------------------ + # Step 5: Verification + # ------------------------------------------------------------------ + + def verify(self, original_image: np.ndarray, generated_image: Image.Image, + defect_mask: np.ndarray, defect_plan) -> Dict: + """ + VLM verification of generated defect. + Since DefectDiffu generates from noise (not editing the original), + we verify that the generated image contains the planned defect + in a plausible location. + """ + print(f"\\n{'='*60}") + print("[Agent] Step 5: Verifying generated defect...") + + # Crop to the mask region (or full image if mask is empty) + mask_bool = defect_mask.astype(bool) + if mask_bool.sum() == 0: + print("[Agent] Warning: Empty mask; verifying full image") + y1, x1 = 0, 0 + y2, x2 = original_image.shape[0], original_image.shape[1] + else: + ys, xs = np.where(mask_bool) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + # Add margin + margin = 32 + h, w = original_image.shape[:2] + y1 = max(0, y1 - margin) + x1 = max(0, x1 - margin) + y2 = min(h, y2 + margin) + x2 = min(w, x2 + margin) + + gen_crop = np.array(generated_image)[y1:y2, x1:x2] + orig_crop = original_image[y1:y2, x1:x2] + + obj_name = f"a {defect_plan.target_subentity or defect_plan.target_entity}" + + result = artifact_description( + self.vlm_client, + original_image, # masked original (full image) + orig_crop, # original crop + gen_crop, # generated crop + obj_name, + defect_plan.defect_type, + self.money_manager + ) + + print(f"[Agent] Verification result: has_artifact={result.has_artifact}") + print(f"[Agent] Explanation: {result.explanation}") + print(f"[Agent] Label: {result.label}") + + return { + 'passed': result.has_artifact, + 'explanation': result.explanation, + 'label': result.label + } + + # ------------------------------------------------------------------ + # Main pipeline + # ------------------------------------------------------------------ + + def run(self, product_description: str, image_path: str, + caption: Optional[str] = None, max_defects: int = 3, + defect_type: Optional[str] = None) -> Dict: + """Run the full agentic pipeline with DefectDiffu.""" + start_time = datetime.now() + exp_id = str(uuid.uuid4())[:8] + + image = np.array(Image.open(image_path).convert('RGB')) + # Resize to DefectDiffu resolution if needed + if image.shape[0] != self.image_size or image.shape[1] != self.image_size: + image_pil = Image.fromarray(image).resize((self.image_size, self.image_size), Image.LANCZOS) + image = np.array(image_pil) + print(f"[Agent] Resized image to {self.image_size}x{self.image_size} for DefectDiffu") + + print(f"[Agent] Loaded image: {image.shape}") + + plan = self.plan(product_description, image, defect_type=defect_type, max_defects=max_defects) + + results = [] + defects_to_process = plan.possible_defects[:max_defects] + print(f"[Agent] Processing top {len(defects_to_process)} of {len(plan.possible_defects)} planned defects") + + for i, defect_plan in enumerate(defects_to_process): + print(f"\\n{'='*80}") + print(f"[Agent] Defect {i+1}/{len(defects_to_process)}: [{defect_plan.defect_type.upper()}] {defect_plan.description}") + print(f"{'='*80}") + + try: + # Step 2: Perceive (VLM ROI bounding box for placement) + prediction = self.perceive(image, defect_plan) + target_bbox = prediction.get('bbox', [0, 0, image.shape[1], image.shape[0]]) + + # Step 3: Prepare conditions + gen_conditions = self.prepare_generation_conditions( + defect_plan, product_description=product_description + ) + + # Step 4: Generate patch from DefectDiffu + generated_image, defect_mask = self.synthesize(gen_conditions) + + # Default ratio presets per defect type + DEFECT_RATIO_PRESETS = { + "bubble": 0.25, # Small localized bubble + "pinhole": 0.15, # Very small point defect + "particle": 0.20, # Dust / particle + "scratch": 0.40, # Medium line scratch + "crack": 0.50, # Medium crack + "smudge": 0.85, # Large surface coverage + "contamination": 1.0, # 100% full ROI coverage + "discoloration": 1.0, # 100% full ROI coverage + "residue": 0.80 # Large area residue + } + + # Get defect type from defect_plan + current_defect_type = getattr(defect_plan, 'defect_type', '').lower() + + # Pick preset ratio (defaults to None / 100% if defect type isn't in presets) + # selected_ratio = DEFECT_RATIO_PRESETS.get(current_defect_type, None) + + # Extract VLM dynamic scale & shape decisions + vlm_ratio = getattr(defect_plan, 'defect_coverage_ratio', None) + vlm_shape = getattr(defect_plan, 'mask_shape', 'free') + + # Fallback ratio if VLM didn't specify + if vlm_ratio is None: + current_defect_type = getattr(defect_plan, 'defect_type', '').lower() + vlm_ratio = DEFECT_RATIO_PRESETS.get(current_defect_type, 1.0) + + # Step 4.5: Blend defect onto REAL image using dynamic ratio + blended_image, full_defect_mask = blend_defect_onto_real_image( + real_image=image, + defect_image=generated_image, + defect_mask=defect_mask, + target_bbox=target_bbox, + max_defect_ratio=vlm_ratio, + mask_shape=vlm_shape + ) + + # Extract the exact pixel bounding box from full_defect_mask + ys, xs = np.where(full_defect_mask > 0) + + if len(ys) > 0 and len(xs) > 0: + # Exact coordinates of the blended defect + mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] + else: + # Fallback to target_bbox if mask is empty + mask_bbox = target_bbox + + # Draw the green verification box around the EXACT mask coordinates + blended_with_green_box = create_visual_prompt_image(blended_image, mask_bbox) + + # Step 5: Verify blended result + verification = self.verify(image, Image.fromarray(blended_image), full_defect_mask, defect_plan) + + # Save outputs + # Dynamically set target folder depending on VLM verification result + if not verification['passed']: + defect_dir = self.output_dir / "failed" / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + else: + defect_dir = self.output_dir / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + + defect_dir.mkdir(parents=True, exist_ok=True) + + Image.fromarray(image).save(defect_dir / "real_clean_image.png") + Image.fromarray(blended_image).save(defect_dir / "blended_factory_defect.png") # Final injected photo + generated_image.save(defect_dir / "raw_defectdiffu_patch.png") + + # Save pixel-exact segmentation mask for model training + mask_img = Image.fromarray(full_defect_mask * 255) + mask_img.save(defect_dir / "defect_mask.png") + + metadata = { + 'experiment_id': exp_id, + 'product_description': product_description, + 'product_type': plan.product_type, + 'defect_plan': defect_plan.dict() if hasattr(defect_plan, 'dict') else vars(defect_plan), + 'generation_conditions': {k: v for k, v in gen_conditions.items() if k != 'defect_plan'}, + 'verification': verification, + 'timestamp': datetime.now().isoformat() + } + with open(defect_dir / "metadata.json", 'w') as f: + json.dump(metadata, f, indent=2, default=str) + + if not verification['passed']: + results.append({ + 'defect_type': defect_plan.defect_type, + 'success': False, + 'verification_passed': False, + 'error': f"VLM verification failed: {verification.get('explanation', 'no explanation')}" + }) + print(f"[Agent] Verification FAILED for {defect_plan.defect_type}.") + else: + results.append({ + 'defect_type': defect_plan.defect_type, + 'success': True, + 'verification_passed': verification['passed'], + 'output_dir': str(defect_dir) + }) + print(f"[Agent] Defect {i+1} complete. Saved to {defect_dir}") + + except Exception as e: + print(f"[Agent] ERROR processing defect {i+1}: {str(e)}") + traceback.print_exc() + results.append({ + 'defect_type': defect_plan.defect_type, + 'success': False, + 'error': str(e) + }) + + # Cleanup caches + if hasattr(self, '_gsam_cache'): + keys_to_remove = [k for k in self._gsam_cache if k.startswith(str(image_path) + "::")] + for k in keys_to_remove: + self._gsam_cache.pop(k, None) + + elapsed = (datetime.now() - start_time).total_seconds() + print(f"\\n{'='*60}") + print(f"[Agent] Pipeline complete in {elapsed:.1f}s") + print(f"[Agent] Results: {sum(1 for r in results if r['success'])}/{len(results)} succeeded") + + return { + 'experiment_id': exp_id, + 'product_type': plan.product_type, + 'results': results, + 'output_dir': str(self.output_dir), + 'elapsed_time': elapsed + } + + def cleanup(self): + """Call once after all batch processing is done.""" + if self.gsam_detector: + self.gsam_detector.cleanup() + self.gsam_detector = None + if self.defectdiffu_generator: + self.defectdiffu_generator.unload_models() + self.defectdiffu_generator = None + print("[Agent] All models cleaned up.") + + +# ============================================================================= +# CLI +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description='ArtiAgent โ€” DefectDiffu Edition') + parser.add_argument('--product-desc', required=True, help='Product description') + parser.add_argument('--image', required=True, help='Path to clean product image (for planning/verification)') + + # DefectDiffu model paths (REQUIRED) + parser.add_argument('--defectdiffu-ckpt', required=True, + help='Path to trained DefectDiffu checkpoint (.pt)') + parser.add_argument('--vae-path', required=True, + help='Path to Stable Diffusion VAE (e.g. stabilityai/sd-vae-ft-mse)') + + # Generation control + parser.add_argument('--defect-type', default=None, + help='Specific defect type to generate (e.g., bubble, scratch)') + parser.add_argument('--output-dir', default='./defect_output', help='Output directory') + parser.add_argument('--caption', default=None, help='Optional image caption') + parser.add_argument('--max-defects', type=int, default=3, + help='Max defects to generate (default: 3)') + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, + help='DefectDiffu generation resolution (default: 512)') + parser.add_argument('--num-steps', type=int, default=50, + help='Denoising steps for DefectDiffu (default: 50)') + + args = parser.parse_args() + + orchestrator = ArtiAgentOrchestrator( + device=args.device, + output_dir=args.output_dir, + vlm_model=args.vlm_model, + defectdiffu_ckpt=args.defectdiffu_ckpt, + vae_path=args.vae_path, + image_size=args.image_size, + num_steps=args.num_steps + ) + + result = orchestrator.run( + product_description=args.product_desc, + image_path=args.image, + caption=args.caption, + max_defects=args.max_defects, + defect_type=args.defect_type + ) + + print(f"\\nFinal output saved to: {result['output_dir']}") + + +if __name__ == "__main__": + main() diff --git a/ArtiAgent - DefectDiffu/src/batch_agent_orchestrator.py b/ArtiAgent - DefectDiffu/src/batch_agent_orchestrator.py new file mode 100644 index 0000000000000000000000000000000000000000..4c2f9b33e5e6fdf223c55462d633530fe8607b4c --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/batch_agent_orchestrator.py @@ -0,0 +1,293 @@ +""" +Batch Agent Orchestrator โ€” DefectDiffu Edition + +Usage: + # Scenario B: Global description for entire folder + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage/vcsel_batch" \ + --product-desc "VCSEL laser diode with emission aperture and surrounding mesa" \ + --output-dir "C:/AgentOutput" \ + --defectdiffu-ckpt "./defectdiffu_ckpt.pt" \ + --vae-path "./sd-vae-ft-mse" \ + --device cuda + + # Scenario A: CSV manifest + python batch_agent_orchestrator.py \\ + --input-dir "C:/TestingImage" \\ + --manifest "C:/products.csv" \\ + --output-dir "C:/AgentOutput" \\ + --defectdiffu-ckpt "./defectdiffu_ckpt.pt" \\ + --vae-path "./sd-vae-ft-mse" + + # Mode 3: Infer from folder names + python batch_agent_orchestrator.py \\ + --input-dir "C:/TestingImage" \\ + --output-dir "C:/AgentOutput" \\ + --defectdiffu-ckpt "./defectdiffu_ckpt.pt" \\ + --vae-path "./sd-vae-ft-mse" +""" + +import os +import sys +import csv +import json +import argparse +import traceback +from pathlib import Path +from datetime import datetime +from typing import Dict, List, Optional +from tqdm import tqdm + +import numpy as np +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from artiagent_orchestrator import ArtiAgentOrchestrator + + +def infer_product_from_path(image_path: Path) -> str: + """Infer product description from folder structure or filename.""" + parent = image_path.parent.name.lower() + if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']: + return parent.replace('_', ' ').replace('-', ' ') + stem = image_path.stem.lower() + for keyword in ['vcsel', 'lens', 'die', 'photodiode', 'sensor', 'chip', 'led', 'laser', 'optical']: + if keyword in stem: + return keyword + return "electronic component" + + +def load_manifest(manifest_path: str) -> Dict[str, str]: + """Load CSV manifest mapping image paths to product descriptions.""" + manifest = {} + with open(manifest_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + img_path = row.get('image_path', row.get('path', row.get('image', ''))).strip() + desc = row.get('product_description', row.get('description', row.get('product', ''))).strip() + if img_path and desc: + manifest[Path(img_path).resolve()] = desc + print(f"[Batch] Loaded manifest with {len(manifest)} entries") + return manifest + + +def discover_images(input_dir: str, extensions=('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')) -> List[Path]: + """Recursively discover all images in input directory.""" + input_path = Path(input_dir) + images = [] + for ext in extensions: + images.extend(input_path.rglob(f"*{ext}")) + images.extend(input_path.rglob(f"*{ext.upper()}")) + unique = sorted(set(images)) + print(f"[Batch] Discovered {len(unique)} images in {input_dir}") + return unique + + +def run_batch( + input_dir: str, + output_dir: str, + defectdiffu_ckpt: str, + vae_path: str, + product_desc: Optional[str] = None, + manifest_path: Optional[str] = None, + defect_type: Optional[str] = None, + max_defects_per_image: int = 3, + device: str = 'cuda', + vlm_model: str = 'gemma3:12b', + image_size: int = 512, + num_steps: int = 50, + resume: bool = False, + save_failed: bool = True +): + """Run agent orchestrator over all images in input directory.""" + + timestamp = datetime.now().strftime("%Y%m%d_%H%M") + output_path = Path(output_dir) / timestamp + output_path.mkdir(parents=True, exist_ok=True) + print(f"[Batch] Output folder: {output_path}") + + manifest = {} + if product_desc and not manifest_path: + print(f"[Batch] Scenario B Active: Global Description = '{product_desc}'") + elif manifest_path and os.path.exists(manifest_path): + manifest = load_manifest(manifest_path) + print(f"[Batch] Scenario A Active: CSV Manifest ({len(manifest)} entries)") + elif product_desc: + print(f"[Batch] Scenario B Active (Fallback): Global Description = '{product_desc}'") + else: + print("[Batch] Scenario C Active: Folder Name Inference (no description provided)") + + images = discover_images(input_dir) + if not images: + print("[Batch] No images found. Exiting.") + return + + progress_file = output_path / "batch_progress.json" + processed_ids = set() + if resume and progress_file.exists(): + with open(progress_file, 'r') as f: + progress = json.load(f) + processed_ids = set(progress.get('processed_paths', [])) + print(f"[Batch] Resuming: {len(processed_ids)} images already processed") + + orchestrator = ArtiAgentOrchestrator( + device=device, + output_dir=str(output_path), + vlm_model=vlm_model, + defectdiffu_ckpt=defectdiffu_ckpt, + vae_path=vae_path, + image_size=image_size, + num_steps=num_steps + ) + + stats = { + 'total': len(images), + 'processed': 0, + 'successful': 0, + 'failed': 0, + 'defects_generated': 0, + 'start_time': datetime.now().isoformat(), + 'processed_paths': [], + 'failed_images': [] + } + + if resume: + images = [img for img in images if str(img.resolve()) not in processed_ids] + + print(f"[Batch] Processing {len(images)} images...") + print(f"[Batch] Max defects per image: {max_defects_per_image}") + print("=" * 70) + + for img_path in tqdm(images, desc="Agent Batch Processing"): + img_key = str(img_path.resolve()) + + try: + if img_key in manifest: + desc = manifest[img_key] + source = "manifest" + elif product_desc: + desc = product_desc + source = "global" + else: + desc = infer_product_from_path(img_path) + source = "inferred" + + print(f"\\n[Batch] Processing: {img_path.name} | desc source: {source}") + if source in ['inferred', 'global']: + print(f"[Batch] Using description: '{desc}'") + + result = orchestrator.run( + product_description=desc, + image_path=str(img_path), + max_defects=max_defects_per_image, + defect_type=defect_type + ) + + successful_defects = sum(1 for r in result['results'] if r['success']) + + stats['processed'] += 1 + stats['successful'] += 1 if successful_defects > 0 else 0 + stats['defects_generated'] += successful_defects + stats['processed_paths'].append(img_key) + + if successful_defects == 0: + stats['failed'] += 1 + stats['failed_images'].append({'path': img_key, 'reason': 'no_defects_generated'}) + + if stats['processed'] % 5 == 0: + with open(progress_file, 'w') as f: + json.dump(stats, f, indent=2) + + except Exception as e: + stats['failed'] += 1 + stats['failed_images'].append({'path': img_key, 'reason': str(e)}) + print(f"[Batch] FAILED: {img_path.name} -> {str(e)}") + if save_failed: + fail_dir = output_path / "_failed" / img_path.stem + fail_dir.mkdir(parents=True, exist_ok=True) + with open(fail_dir / "error.txt", 'w') as f: + f.write(traceback.format_exc()) + + with open(progress_file, 'w') as f: + json.dump(stats, f, indent=2) + + orchestrator.cleanup() + + elapsed = (datetime.now() - datetime.fromisoformat(stats['start_time'])).total_seconds() + hours = int(elapsed // 3600) + minutes = int((elapsed % 3600) // 60) + seconds = int(elapsed % 60) + + print("\\n" + "=" * 70) + print("BATCH ORCHESTRATION COMPLETE") + print("=" * 70) + print(f"Total images: {stats['total']}") + print(f"Processed: {stats['processed']}") + print(f"Successful: {stats['successful']}") + print(f"Failed: {stats['failed']}") + print(f"Defects generated: {stats['defects_generated']}") + print(f"Total time: {hours}h {minutes}m {seconds}s") + print(f"Output directory: {output_path}") + print("=" * 70) + + +def main(): + parser = argparse.ArgumentParser(description='Batch Agent-Driven Defect Generation (DefectDiffu)') + + # Input / Output + parser.add_argument('--input-dir', required=True, help='Directory containing clean product images') + parser.add_argument('--output-dir', required=True, help='Output directory for all defect images') + + # DefectDiffu model paths (REQUIRED) + parser.add_argument('--defectdiffu-ckpt', required=True, + help='Path to trained DefectDiffu checkpoint') + parser.add_argument('--vae-path', required=True, + help='Path to Stable Diffusion VAE (e.g. stabilityai/sd-vae-ft-mse)') + + # Description sources + parser.add_argument('--product-desc', default=None, + help='[Scenario B] ONE global description applied to ALL images in the folder') + parser.add_argument('--manifest', default=None, + help='[Scenario A] CSV manifest with columns: image_path,product_description') + + # Defect control + parser.add_argument('--defect-type', default=None, + help='Specify defect type for batch generation (e.g., bubble, scratch)') + parser.add_argument('--max-defects-per-image', type=int, default=3, + help='Maximum defects to generate per image (default: 3)') + + # Generation control + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, + help='DefectDiffu generation resolution (default: 512)') + parser.add_argument('--num-steps', type=int, default=50, + help='Denoising steps for DefectDiffu (default: 50)') + parser.add_argument('--resume', action='store_true', help='Resume from previous batch run') + parser.add_argument('--no-save-failed', action='store_true', help='Do not save failed case logs') + + args = parser.parse_args() + + run_batch( + input_dir=args.input_dir, + output_dir=args.output_dir, + defectdiffu_ckpt=args.defectdiffu_ckpt, + vae_path=args.vae_path, + product_desc=args.product_desc, + manifest_path=args.manifest, + defect_type=args.defect_type, + max_defects_per_image=args.max_defects_per_image, + device=args.device, + vlm_model=args.vlm_model, + image_size=args.image_size, + num_steps=args.num_steps, + resume=args.resume, + save_failed=not args.no_save_failed + ) + + +if __name__ == "__main__": + main() diff --git a/ArtiAgent - DefectDiffu/src/flux/__init__.py b/ArtiAgent - DefectDiffu/src/flux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..43c365a49d6980e88acba10ef3069f110a59644a --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/__init__.py @@ -0,0 +1,11 @@ +try: + from ._version import version as __version__ # type: ignore + from ._version import version_tuple +except ImportError: + __version__ = "unknown (no version information available)" + version_tuple = (0, 0, "unknown", "noinfo") + +from pathlib import Path + +PACKAGE = __package__.replace("_", "-") +PACKAGE_ROOT = Path(__file__).parent diff --git a/ArtiAgent - DefectDiffu/src/flux/__main__.py b/ArtiAgent - DefectDiffu/src/flux/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..d5cf0fd2444d4cda4053fa74dad3371556b886e5 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/__main__.py @@ -0,0 +1,4 @@ +from .cli import app + +if __name__ == "__main__": + app() diff --git a/ArtiAgent - DefectDiffu/src/flux/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46cf619948bf6c42c3e8f33cb0710cf5f75492c2 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/__pycache__/artifacts_util.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/__pycache__/artifacts_util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb6eccc55fd52291945088cd96b2e2e04ba57333 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/__pycache__/artifacts_util.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/__pycache__/math.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/__pycache__/math.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3752064f736ab300fcf0f6bd9e72a412c3416df0 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/__pycache__/math.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/__pycache__/model.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86e79a513113889165d28e2121f97243b014a50c Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/__pycache__/model.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/__pycache__/sampling.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/__pycache__/sampling.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edf1f9c6090dbb91b2f7f54ccf86f1afeb62ffef Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/__pycache__/sampling.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/__pycache__/util.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/__pycache__/util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3e2b9ff8fdf2cad9c44a1008c1943dac8ab8073 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/__pycache__/util.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/api.py b/ArtiAgent - DefectDiffu/src/flux/api.py new file mode 100644 index 0000000000000000000000000000000000000000..b08202adb35d2ffae320bb9b47f567e538837836 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/api.py @@ -0,0 +1,194 @@ +import io +import os +import time +from pathlib import Path + +import requests +from PIL import Image + +API_ENDPOINT = "https://api.bfl.ml" + + +class ApiException(Exception): + def __init__(self, status_code: int, detail: str | list[dict] | None = None): + super().__init__() + self.detail = detail + self.status_code = status_code + + def __str__(self) -> str: + return self.__repr__() + + def __repr__(self) -> str: + if self.detail is None: + message = None + elif isinstance(self.detail, str): + message = self.detail + else: + message = "[" + ",".join(d["msg"] for d in self.detail) + "]" + return f"ApiException({self.status_code=}, {message=}, detail={self.detail})" + + +class ImageRequest: + def __init__( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + name: str = "flux.1-pro", + num_steps: int = 50, + prompt_upsampling: bool = False, + seed: int | None = None, + validate: bool = True, + launch: bool = True, + api_key: str | None = None, + ): + """ + Manages an image generation request to the API. + + Args: + prompt: Prompt to sample + width: Width of the image in pixel + height: Height of the image in pixel + name: Name of the model + num_steps: Number of network evaluations + prompt_upsampling: Use prompt upsampling + seed: Fix the generation seed + validate: Run input validation + launch: Directly launches request + api_key: Your API key if not provided by the environment + + Raises: + ValueError: For invalid input + ApiException: For errors raised from the API + """ + if validate: + if name not in ["flux.1-pro"]: + raise ValueError(f"Invalid model {name}") + elif width % 32 != 0: + raise ValueError(f"width must be divisible by 32, got {width}") + elif not (256 <= width <= 1440): + raise ValueError(f"width must be between 256 and 1440, got {width}") + elif height % 32 != 0: + raise ValueError(f"height must be divisible by 32, got {height}") + elif not (256 <= height <= 1440): + raise ValueError(f"height must be between 256 and 1440, got {height}") + elif not (1 <= num_steps <= 50): + raise ValueError(f"steps must be between 1 and 50, got {num_steps}") + + self.request_json = { + "prompt": prompt, + "width": width, + "height": height, + "variant": name, + "steps": num_steps, + "prompt_upsampling": prompt_upsampling, + } + if seed is not None: + self.request_json["seed"] = seed + + self.request_id: str | None = None + self.result: dict | None = None + self._image_bytes: bytes | None = None + self._url: str | None = None + if api_key is None: + self.api_key = os.environ.get("BFL_API_KEY") + else: + self.api_key = api_key + + if launch: + self.request() + + def request(self): + """ + Request to generate the image. + """ + if self.request_id is not None: + return + response = requests.post( + f"{API_ENDPOINT}/v1/image", + headers={ + "accept": "application/json", + "x-key": self.api_key, + "Content-Type": "application/json", + }, + json=self.request_json, + ) + result = response.json() + if response.status_code != 200: + raise ApiException(status_code=response.status_code, detail=result.get("detail")) + self.request_id = response.json()["id"] + + def retrieve(self) -> dict: + """ + Wait for the generation to finish and retrieve response. + """ + if self.request_id is None: + self.request() + while self.result is None: + response = requests.get( + f"{API_ENDPOINT}/v1/get_result", + headers={ + "accept": "application/json", + "x-key": self.api_key, + }, + params={ + "id": self.request_id, + }, + ) + result = response.json() + if "status" not in result: + raise ApiException(status_code=response.status_code, detail=result.get("detail")) + elif result["status"] == "Ready": + self.result = result["result"] + elif result["status"] == "Pending": + time.sleep(0.5) + else: + raise ApiException(status_code=200, detail=f"API returned status '{result['status']}'") + return self.result + + @property + def bytes(self) -> bytes: + """ + Generated image as bytes. + """ + if self._image_bytes is None: + response = requests.get(self.url) + if response.status_code == 200: + self._image_bytes = response.content + else: + raise ApiException(status_code=response.status_code) + return self._image_bytes + + @property + def url(self) -> str: + """ + Public url to retrieve the image from + """ + if self._url is None: + result = self.retrieve() + self._url = result["sample"] + return self._url + + @property + def image(self) -> Image.Image: + """ + Load the image as a PIL Image + """ + return Image.open(io.BytesIO(self.bytes)) + + def save(self, path: str): + """ + Save the generated image to a local path + """ + suffix = Path(self.url).suffix + if not path.endswith(suffix): + path = path + suffix + Path(path).resolve().parent.mkdir(parents=True, exist_ok=True) + with open(path, "wb") as file: + file.write(self.bytes) + + +if __name__ == "__main__": + from fire import Fire + + Fire(ImageRequest) diff --git a/ArtiAgent - DefectDiffu/src/flux/artifacts_util.py b/ArtiAgent - DefectDiffu/src/flux/artifacts_util.py new file mode 100644 index 0000000000000000000000000000000000000000..98caf48d4614d0de4eee997a6b3fee88fd35c848 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/artifacts_util.py @@ -0,0 +1,335 @@ +import torch +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as patches +import os + +def patch_coor_to_ind(x, y, w, txt_len): + return y * w + x + txt_len + +def patch_ind_to_coor(ind, w, txt_len, return_shape=False): + ind -= txt_len + y = ind // w + x = ind % w + return [y, x] + +def patch_indices_to_coords(indices, w, txt_len=512): + """ + Convert patch indices back to patch coordinates. + + Args: + indices: List of patch indices + w: Patch width + txt_len: Text length offset + + Returns: + List of (y, x) patch coordinates + """ + return [patch_ind_to_coor(ind, w, txt_len) for ind in indices] + +def bbox_to_patch_indices(bbox_coordinates, h, w, patch_size=16, txt_len=512): + xmin, xmax, ymin, ymax = bbox_coordinates + patch_xmin = xmin // patch_size + patch_xmax = (xmax - 1) // patch_size + patch_ymin = ymin // patch_size + patch_ymax = (ymax - 1) // patch_size + indices = [ + patch_coor_to_ind(px, py, w, txt_len) + for py in range(patch_ymin, patch_ymax + 1) + for px in range(patch_xmin, patch_xmax + 1) + ] + return indices + +def bbox_to_patch_coords(bbox_coordinates, patch_size=16, return_shape=False): + xmin, xmax, ymin, ymax = bbox_coordinates + patch_xmin = xmin // patch_size + patch_xmax = (xmax - 1) // patch_size + patch_ymin = ymin // patch_size + patch_ymax = (ymax - 1) // patch_size + coords = [ (py, px) + for py in range(patch_ymin, patch_ymax + 1) + for px in range(patch_xmin, patch_xmax + 1) + ] + if return_shape: + return coords, patch_ymax-patch_ymin+1, patch_xmax-patch_xmin+1 + return coords + +# New shape-based functions +def mask_to_patch_indices(mask, patch_size=16, txt_len=512): + """ + Convert a binary mask to patch indices. + + Args: + mask: Binary mask (numpy array) where 1 indicates the region of interest + patch_size: Size of each patch (default 16 for Flux) + txt_len: Text length offset (default 512) + + Returns: + List of patch indices + """ + h, w = mask.shape + patch_h, patch_w = h // patch_size, w // patch_size + + # Downsample mask to patch resolution + patch_mask = np.zeros((patch_h, patch_w), dtype=bool) + + for py in range(patch_h): + for px in range(patch_w): + # Get the patch region in the original mask + y_start, y_end = py * patch_size, (py + 1) * patch_size + x_start, x_end = px * patch_size, (px + 1) * patch_size + + # If any part of the patch overlaps with the mask, include it + patch_region = mask[y_start:y_end, x_start:x_end] + if np.any(patch_region): + patch_mask[py, px] = True + + # Convert patch coordinates to indices + indices = [] + for py in range(patch_h): + for px in range(patch_w): + if patch_mask[py, px]: + indices.append(patch_coor_to_ind(px, py, patch_w, txt_len)) + + return indices + +def mask_to_patch_coords(mask, patch_size=16): + """ + Convert a binary mask to patch coordinates. + + Args: + mask: Binary mask (numpy array) where 1 indicates the region of interest + patch_size: Size of each patch (default 16 for Flux) + + Returns: + List of (py, px) patch coordinates + """ + h, w = mask.shape + patch_h, patch_w = h // patch_size, w // patch_size + + # Downsample mask to patch resolution + patch_mask = np.zeros((patch_h, patch_w), dtype=bool) + + for py in range(patch_h): + for px in range(patch_w): + # Get the patch region in the original mask + y_start, y_end = py * patch_size, (py + 1) * patch_size + x_start, x_end = px * patch_size, (px + 1) * patch_size + + # If any part of the patch overlaps with the mask, include it + patch_region = mask[y_start:y_end, x_start:x_end] + if np.any(patch_region): + patch_mask[py, px] = True + + # Convert patch coordinates to list + coords = [] + for py in range(patch_h): + for px in range(patch_w): + if patch_mask[py, px]: + coords.append((py, px)) + + return coords + +def get_closest_patch_ind(h, w, bbox_coordinates, patch_size=16, txt_len=512): + # Get only valid coordinates (value == 1) + + large_array = np.ones((h,w), dtype=int) + small_grid_coords, bbox_h, bbox_w = bbox_to_patch_coords(bbox_coordinates, patch_size=patch_size, return_shape=True) + + for y, x in small_grid_coords: + large_array[y,x] = 0 + + valid_coords = np.argwhere(large_array == 1) + + result = np.empty((bbox_h, bbox_w), dtype=object) + min_h, min_w = small_grid_coords[0] + + for idx, coord in enumerate(small_grid_coords): + y, x = coord + distances = np.abs(valid_coords[:, 0] - y) + np.abs(valid_coords[:, 1] - x) + min_idx = np.argmin(distances) + closest_coord = tuple(valid_coords[min_idx]) + result[y-min_h,x-min_w] = closest_coord + + return [patch_coor_to_ind(x,y,w,txt_len) for y, x in result.flatten()] + +def get_neighbors_patch_ind(h, w, bbox_coordinates, img_ids, patch_size=16, txt_len=512): + # Get the patch coordinates and shape of the bbox + small_grid_coords, bbox_h, bbox_w = bbox_to_patch_coords(bbox_coordinates, patch_size=patch_size, return_shape=True) + small_grid_coords = np.array(small_grid_coords) + + indices_array = img_ids.cpu().numpy().squeeze().copy() + indices_array = indices_array.reshape((h*w,-1)) + + # indices_array = indices_array.reshape((h*w,-1)) + + # Compute bbox center in patch coordinates + min_h, min_w = np.min(small_grid_coords, axis=0) + max_h, max_w = np.max(small_grid_coords, axis=0) + center_y = (min_h + max_h) // 2 + center_x = (min_w + max_w) // 2 + + min_h_p, min_w_p = max(min_h-2, 0), max(min_w-2, 0) + max_h_p, max_w_p = min(h, max_h+2), min(w, max_w+2) + + # Compute shortest distance from center to bbox edge + radius_y = min(center_y - min_h, max_h - center_y) + radius_x = min(center_x - min_w, max_w - center_x) + radius = min(radius_y, radius_x) + + # Get all valid coordinates in the grid + yy, xx = np.meshgrid(np.arange(h), np.arange(w), indexing='ij') + all_coords = np.stack([yy.ravel(), xx.ravel()], axis=1) + + # Exclude bbox coordinates + bbox_set = set(map(tuple, small_grid_coords)) + filtered_coords = [tuple(coord) for coord in all_coords if tuple(coord) not in bbox_set] + + # result_coords=[] + # for center_y, center_x in small_grid_coords: + # neighbors = [] + # for dy in range(-radius-1, radius + 2): + # for dx in range(-radius-1, radius + 2): + # if abs(dy) + abs(dx) <= radius+2: + # ny, nx = center_y + dy, center_x + dx + # if 0 <= ny < h and 0 <= nx < w: + # if (ny, nx) not in bbox_set: + # neighbors.append((ny, nx)) + # result_coords.append(neighbors) + + # return small_grid_coords.tolist(), result_coords + + result=[] + for center_y, center_x in small_grid_coords: + neighbors = [] + for dy in range(-radius-1, radius + 2): + for dx in range(-radius-1, radius + 2): + if abs(dy) + abs(dx) <= radius+1: + ny, nx = center_y + dy, center_x + dx + if min_h_p <= ny < max_h_p and min_w_p <= nx < max_w_p: + if (ny, nx) not in bbox_set: + neighbors.append((ny, nx)) + # if len(neighbors) == 0: + # import pdb;pdb.set_trace() + neighbors_ind = indices_array[[patch_coor_to_ind(x,y,w,0) for (y,x) in neighbors],:] + result.append(neighbors_ind.mean(0)) + + return torch.from_numpy(np.array(result)).unsqueeze(0) + +def perturb_pe(h, w, bbox_coordinates, img_ids, patch_size=16, txt_len=512): + patch_ids = bbox_to_patch_indices(bbox_coordinates, h, w, patch_size, txt_len=0) + indices_array = torch.from_numpy(img_ids.cpu().numpy().squeeze().copy()[patch_ids, :]) + noise = torch.randn_like(indices_array) + # Mask for non-zero elements + nonzero_mask = indices_array != 0 + # Clone indices_array to preserve original + perturbed_indices_array = indices_array.clone() + # Apply noise only to non-zero elements + perturbed_indices_array[nonzero_mask] += noise[nonzero_mask] + # perturbed_indices_array = indices_array + torch.randn_like(indices_array) * 0.3 + return perturbed_indices_array.unsqueeze(0) + +def shuffle_pe(h, w, patch_ids, patch_size=16, txt_len=512, intensity=3): + bbox_coords = [patch_ind_to_coor(ind, w, txt_len) for ind in patch_ids] + # bbox_coords = bbox_to_patch_coords(bbox_coordinates, patch_size=patch_size) + + shuffled_coords = [] + + for y, x in bbox_coords: + dx = torch.randint(-intensity, intensity + 1, (1,)).item() + dy = torch.randint(-intensity, intensity + 1, (1,)).item() + # dx = 0 + # dy = 0 + + new_x = max(0, min(x + dx, w - 1)) + new_y = max(0, min(y + dy, h - 1)) + + shuffled_coords.append((new_y, new_x)) + + return [patch_coor_to_ind(x,y,w,txt_len) for y, x in shuffled_coords] + +def sample_closest_patch_ind(h, w, patch_indices, reference_patch_indices, patch_size=16, txt_len=512): + """ + Sample closest patches for arbitrary shape with randomization. + + Args: + h, w: Patch grid dimensions + patch_indices: List of patch indices defining the shape + reference_patch_indices: List of patch indices to use as reference/candidates + patch_size: Size of each patch + txt_len: Text length offset + + Returns: + List of sampled closest patch indices + """ + shape_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in patch_indices]) + + # Convert reference patch indices to coordinates + reference_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in reference_patch_indices]) + + result = [] + for y, x in shape_coords: + if 0 <= y < h and 0 <= x < w: + # Calculate distances only to reference patch coordinates + distances = np.abs(reference_coords[:, 0] - y) + np.abs(reference_coords[:, 1] - x) + inv_d = 1.0 / (distances + 1e-8) + inv_d = np.pow(inv_d, 2) + p_weight = inv_d / np.sum(inv_d) + idx = np.random.choice(len(distances), p=p_weight) + closest_coord = tuple(reference_coords[idx]) + result.append(patch_coor_to_ind(closest_coord[1], closest_coord[0], w, txt_len)) + + return result + +def get_closest_patch_coords(target_coords, reference_coords): + """ + Map each target coordinate to its closest reference coordinate. + + Args: + target_coords: List of (y, x) coordinates that need to be mapped + reference_coords: List of (y, x) coordinates to use as reference/candidates + + Returns: + List of closest reference coordinates for each target coordinate + """ + target_coords = np.array(target_coords) + reference_coords = np.array(reference_coords) + + result = [] + for ty, tx in target_coords: + # Calculate Manhattan distances to all reference coordinates + distances = np.abs(reference_coords[:, 0] - ty) + np.abs(reference_coords[:, 1] - tx) + min_idx = np.argmin(distances) + closest_coord = tuple(reference_coords[min_idx]) + result.append(closest_coord) + + return result + + +def get_closest_patch_inds(h, w, target_patch_indices, reference_patch_indices, txt_len=512): + """ + Map each target patch index to the closest reference patch index using Manhattan distance. + + Args: + h, w: Patch grid dimensions (not used directly but kept for API symmetry) + target_patch_indices: List of patch indices to map + reference_patch_indices: List of candidate reference patch indices + txt_len: Text length offset used in index<->coord conversions + + Returns: + List of closest reference patch indices corresponding to each target patch index + """ + if len(target_patch_indices) == 0 or len(reference_patch_indices) == 0: + return [] + + # Convert indices to (y, x) coordinates + target_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in target_patch_indices]) + reference_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in reference_patch_indices]) + + result = [] + for ty, tx in target_coords: + distances = np.abs(reference_coords[:, 0] - ty) + np.abs(reference_coords[:, 1] - tx) + min_idx = int(np.argmin(distances)) + result.append(reference_patch_indices[min_idx]) + + return result diff --git a/ArtiAgent - DefectDiffu/src/flux/math.py b/ArtiAgent - DefectDiffu/src/flux/math.py new file mode 100644 index 0000000000000000000000000000000000000000..a1f0a62c0d63a5425551e5bc411ea86fc47c53ee --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/math.py @@ -0,0 +1,54 @@ +import math +import torch +from einops import rearrange +from torch import Tensor + + +def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor: + q, k = apply_rope(q, k, pe) + + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "B H L D -> B L (H D)") + + return x + +def attention_masked(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, patch_ids: list[int], mask: Tensor, return_weight:bool=False) -> Tensor: + q, k = apply_rope(q, k, pe) + # x = torch.nn.functional.scaled_dot_product_attention(q,k,v,attn_mask=attn_mask) + if return_weight: + x, m = scaled_dot_product_attention_masked(q, k, v, patch_ids, mask, return_weight) + x = rearrange(x, "B H L D -> B L (H D)") + return x, m + else: + x = scaled_dot_product_attention_masked(q, k, v, patch_ids, mask, return_weight) + x = rearrange(x, "B H L D -> B L (H D)") + return x + + # return x, m + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + assert dim % 2 == 0 + scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = torch.einsum("...n,d->...nd", pos, omega) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2) + return out.float() + + +def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) + + +def scaled_dot_product_attention_masked(query, key, value, patch_ids, attn_mask, return_weight=False): + scale_factor = 1 / math.sqrt(query.size(-1)) + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_mask + attn_weight = torch.softmax(attn_weight, dim=-1) + if return_weight: + return attn_weight @ value, attn_weight + return attn_weight @ value diff --git a/ArtiAgent - DefectDiffu/src/flux/model.py b/ArtiAgent - DefectDiffu/src/flux/model.py new file mode 100644 index 0000000000000000000000000000000000000000..582990668e435b82472d4eb6c6f0fab541d116fa --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/model.py @@ -0,0 +1,249 @@ +from dataclasses import dataclass + +import torch +from torch import Tensor, nn +import numpy as np + +from flux.modules.layers import (DoubleStreamBlock, EmbedND, LastLayer, + MLPEmbedder, SingleStreamBlock, + timestep_embedding) + + +@dataclass +class FluxParams: + in_channels: int + out_channels: int + vec_in_dim: int + context_in_dim: int + hidden_size: int + mlp_ratio: float + num_heads: int + depth: int + depth_single_blocks: int + axes_dim: list[int] + theta: int + qkv_bias: bool + guidance_embed: bool + + +class Flux(nn.Module): + """ + Transformer model for flow matching on sequences. + """ + + def __init__(self, params: FluxParams): + super().__init__() + + self.params = params + self.in_channels = params.in_channels + self.out_channels = params.out_channels + if params.hidden_size % params.num_heads != 0: + raise ValueError( + f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}" + ) + pe_dim = params.hidden_size // params.num_heads + if sum(params.axes_dim) != pe_dim: + raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}") + self.hidden_size = params.hidden_size + self.num_heads = params.num_heads + self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim) + self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) + self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size) + self.guidance_in = ( + MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + ) + self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size) + + self.double_blocks = nn.ModuleList( + [ + DoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + ) + for _ in range(params.depth) + ] + ) + + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio) + for _ in range(params.depth_single_blocks) + ] + ) + + self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels) + self._sequential_offload = False + + def enable_sequential_cpu_offload(self): + self._sequential_offload = True + + def forward( + self, + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor | None = None, + info = None, + ref_img: Tensor | None = None, # โ† NEW + ref_img_ids: Tensor | None = None, # โ† NEW + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + # Ensure inputs match the model's dtype (NF4 can silently upcast to float32) + target_dtype = self.img_in.weight.dtype + img = img.to(target_dtype) + txt = txt.to(target_dtype) + if y.dtype != target_dtype: + y = y.to(target_dtype) + + # running on sequences img + img = self.img_in(img) + original_img_seq_len = img.shape[1] # โ† REMEMBER: how many original img tokens + + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + if guidance is None: + raise ValueError("Didn't get guidance strength for guidance distilled model.") + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # NEW: Concatenate reference tokens into image stream + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + if ref_img is not None and ref_img_ids is not None: + ref = self.img_in(ref_img) # project reference patches same way + img = torch.cat([img, ref], dim=1) + img_ids = torch.cat([img_ids, ref_img_ids], dim=1) + + if ref_img is not None: + print(f"[Flux.forward] Attending to {ref_img.shape[1]} ref tokens + {original_img_seq_len} img tokens") + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + inject_pe = pe.clone() + if not info['inverse']: + # Defensive clamp: GSAM indices may exceed seq_len for non-16-divisible images + seq_len = pe.shape[2] + + # Initialize accumulated lists for tracking all processed IDs + accumulated_target_ids = [] + accumulated_ref_ids = [] + for artifact_data in info['artifact_data']: + if artifact_data['artifact_type'] == 'addition' and info['addition']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + + # Accumulate IDs + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + elif artifact_data['artifact_type'] == 'removal' and info['removal']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + + # ref_ids = get_closest_patch_inds(info['patch_h'], info['patch_w'], target_ids, ref_ids) + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + # Accumulate IDs (after target_ids modification) + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + elif artifact_data['artifact_type'] == 'distortion' and info['distortion']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + + if len(ref_ids) == 0: + # For distortion with no reference patches, shuffle target patches + ref_ids = target_ids.copy() + np.random.shuffle(ref_ids) + # Ensure target_ids and ref_ids are different for distortion + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + # Accumulate IDs (after any ref_ids modification) + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + elif artifact_data['artifact_type'] == 'fusion' and info['fusion']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + + # np.random.shuffle(ref_ids) + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + + # Accumulate IDs + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + info['patch_ids'] = accumulated_target_ids + info['patch_ref_ids'] = accumulated_ref_ids + info['timesteps'] = timesteps + + + if self._sequential_offload: + for block in self.double_blocks: + block = block.to(img.device) + img, txt = block(img=img, txt=txt, vec=vec, pe=inject_pe, info=info) + block = block.cpu() + torch.cuda.empty_cache() + else: + for block in self.double_blocks: + img, txt = block(img=img, txt=txt, vec=vec, pe=inject_pe, info=info) + + cnt = 0 + img = torch.cat((txt, img), 1) + info['type'] = 'single' + if self._sequential_offload: + for block in self.single_blocks: + block = block.to(img.device) + info['id'] = cnt + if cnt < 19: + img, info = block(img, vec=vec, pe=inject_pe, info=info) + else: + img, info = block(img, vec=vec, pe=pe, info=info) + block = block.cpu() + torch.cuda.empty_cache() + cnt += 1 + else: + for block in self.single_blocks: + info['id'] = cnt + if cnt < 19: + img, info = block(img, vec=vec, pe=inject_pe, info=info) + else: + img, info = block(img, vec=vec, pe=pe, info=info) + cnt += 1 + + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # MODIFIED: Extract only ORIGINAL img tokens + # Before: img = img[:, txt.shape[1] :, ...] (gets img + ref) + # After: img = img[:, txt.shape[1] : txt.shape[1] + original_img_seq_len, ...] + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + img = img[:, txt.shape[1] : txt.shape[1] + original_img_seq_len, ...] + + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + return img, info \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/autoencoder.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/autoencoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c340ba212c90e5afb1a6e4ca2bcb5bbd35b4eba Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/autoencoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/conditioner.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/conditioner.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..658ee8c803d824ca7c62beef87d4bb4264b04c73 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/conditioner.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/layers.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2a6b4c42c2cfffa108a5eeae740e00fb71fb517 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/flux/modules/__pycache__/layers.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/flux/modules/autoencoder.py b/ArtiAgent - DefectDiffu/src/flux/modules/autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..86bdec01bd09c872721fe267fe1bd83d32d5fdec --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/modules/autoencoder.py @@ -0,0 +1,313 @@ +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn + + +@dataclass +class AutoEncoderParams: + resolution: int + in_channels: int + ch: int + out_ch: int + ch_mult: list[int] + num_res_blocks: int + z_channels: int + scale_factor: float + shift_factor: float + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, h, w = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + + return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + # no asymmetric padding in torch conv, must do it ourselves + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: Tensor): + pad = (0, 1, 0, 1) + x = nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + return x + + +class Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor): + x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + x = self.conv(x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + ch: int, + ch_mult: list[int], + num_res_blocks: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + # downsampling + self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) + + curr_res = resolution + in_ch_mult = (1,) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + block_in = self.ch + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + ch: int, + out_ch: int, + ch_mult: list[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.ffactor = 2 ** (self.num_resolutions - 1) + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = ch * ch_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + + # z to block_in + self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: Tensor) -> Tensor: + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class DiagonalGaussian(nn.Module): + def __init__(self, sample: bool = True, chunk_dim: int = 1): + super().__init__() + self.sample = sample + self.chunk_dim = chunk_dim + + def forward(self, z: Tensor) -> Tensor: + mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim) + # import pdb;pdb.set_trace() + if self.sample: + std = torch.exp(0.5 * logvar) + return mean #+ std * torch.randn_like(mean) + else: + return mean + + +class AutoEncoder(nn.Module): + def __init__(self, params: AutoEncoderParams): + super().__init__() + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian() + + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: Tensor) -> Tensor: + z = self.reg(self.encoder(x)) + z = self.scale_factor * (z - self.shift_factor) + return z + + def decode(self, z: Tensor) -> Tensor: + z = z / self.scale_factor + self.shift_factor + return self.decoder(z) + + def forward(self, x: Tensor) -> Tensor: + return self.decode(self.encode(x)) diff --git a/ArtiAgent - DefectDiffu/src/flux/modules/conditioner.py b/ArtiAgent - DefectDiffu/src/flux/modules/conditioner.py new file mode 100644 index 0000000000000000000000000000000000000000..98dbffd132a61cb74f1cc50ad3405d2cd58f3268 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/modules/conditioner.py @@ -0,0 +1,64 @@ +import torch +from torch import Tensor, nn +from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer, BitsAndBytesConfig) + + +class HFEmbedder(nn.Module): + def __init__(self, version: str, max_length: int, is_clip, **hf_kwargs): + super().__init__() + self.is_clip = is_clip + self.max_length = max_length + self.output_key = "pooler_output" if self.is_clip else "last_hidden_state" + + # Safely remove 'load_in_8bit' and 'device_map' from hf_kwargs so they don't get passed to __init__ + self.is_8bit = hf_kwargs.pop("load_in_8bit", False) + device_map = hf_kwargs.pop("device_map", "cuda") + + if self.is_clip: + self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs) + else: + self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length) + if self.is_8bit: + # Use BitsAndBytesConfig for modern transformers + # Remove torch_dtype conflict if present in kwargs + hf_kwargs.pop("torch_dtype", None) + q_config = BitsAndBytesConfig(load_in_8bit=True) + # Remove torch_dtype conflict if present in hf_kwargs + hf_kwargs.pop("torch_dtype", None) + + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained( + version, + quantization_config=q_config, + device_map=hf_kwargs.pop("device_map", "cuda"), + **hf_kwargs + ) + else: + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs) + + self.hf_module = self.hf_module.eval().requires_grad_(False) + + def to(self, *args, **kwargs): + # If loaded in 8-bit, bitsandbytes handles device placement automatically. + # Calling .to() on an 8-bit model will crash, so we skip it. + if self.is_8bit: + return self + return super().to(*args, **kwargs) + + def forward(self, text: list[str]) -> Tensor: + batch_encoding = self.tokenizer( + text, + truncation=True, + max_length=self.max_length, + return_length=False, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + + outputs = self.hf_module( + input_ids=batch_encoding["input_ids"].to(self.hf_module.device), + attention_mask=None, + output_hidden_states=False, + ) + return outputs[self.output_key].to(torch.bfloat16) diff --git a/ArtiAgent - DefectDiffu/src/flux/modules/layers.py b/ArtiAgent - DefectDiffu/src/flux/modules/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..3ecc8b7600e0e02eb74d163965cee5f0349acbf5 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/modules/layers.py @@ -0,0 +1,288 @@ +import math +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn +import random + +from flux.math import attention, attention_masked, rope + +import os + +class EmbedND(nn.Module): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + n_axes = ids.shape[-1] + emb = torch.cat( + [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], + dim=-3, + ) + + return emb.unsqueeze(1) + + +def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + t = time_factor * t + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + t.device + ) + + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(t) + return embedding + + +class MLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class RMSNorm(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor): + x_dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=x_dtype) * self.scale.to(dtype=x_dtype) + + +class QKNorm(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = RMSNorm(dim) + self.key_norm = RMSNorm(dim) + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: + q = self.query_norm(q) + k = self.key_norm(k) + return q.to(v), k.to(v) + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.norm = QKNorm(head_dim) + self.proj = nn.Linear(dim, dim) + + def forward(self, x: Tensor, pe: Tensor, patch_ids=None) -> Tensor: + qkv = self.qkv(x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + if patch_ids is None: + x = attention(q, k, v, pe=pe) + else: + x = attention_masked(q, k, v, pe=pe, patch_ids=patch_ids) + + x = self.proj(x) + return x + + +@dataclass +class ModulationOut: + shift: Tensor + scale: Tensor + gate: Tensor + + +class Modulation(nn.Module): + def __init__(self, dim: int, double: bool): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) + + def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1) + + return ( + ModulationOut(*out[:3]), + ModulationOut(*out[3:]) if self.is_double else None, + ) + + +class DoubleStreamBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False): + super().__init__() + + mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.num_heads = num_heads + self.hidden_size = hidden_size + self.img_mod = Modulation(hidden_size, double=True) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + self.txt_mod = Modulation(hidden_size, double=True) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, info) -> tuple[Tensor, Tensor]: + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + # prepare image for attention + img_modulated = self.img_norm1(img) + img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift + img_qkv = self.img_attn.qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) + + # prepare txt for attention + txt_modulated = self.txt_norm1(txt) + txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift + txt_qkv = self.txt_attn.qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) + + # run actual attention + q = torch.cat((txt_q, img_q), dim=2) #[8, 24, 512, 128] + [8, 24, 900, 128] -> [8, 24, 1412, 128] + k = torch.cat((txt_k, img_k), dim=2) + v = torch.cat((txt_v, img_v), dim=2) + + attn = attention(q, k, v, pe=pe) + + txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + + # calculate the txt bloks + txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) + return img, txt + + +class SingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + qk_scale: float | None = None, + ): + super().__init__() + self.hidden_dim = hidden_size + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.scale = qk_scale or head_dim**-0.5 + + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + # qkv and mlp_in + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) + # proj and mlp_out + self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) + + self.norm = QKNorm(head_dim) + + self.hidden_size = hidden_size + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + self.mlp_act = nn.GELU(approximate="tanh") + self.modulation = Modulation(hidden_size, double=False) + + def forward(self, x: Tensor, vec: Tensor, pe: Tensor, info) -> Tensor: + mod, _ = self.modulation(vec) + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + + # Note: If the memory of your device is not enough, you may consider uncommenting the following code. + # if info['inject'] and info['id'] > 19: + # store_path = os.path.join(info['feature_path'], str(info['t']) + '_' + str(info['second_order']) + '_' + str(info['id']) + '_' + info['type'] + '_' + 'V' + '.pth') + # if info['inverse']: + # torch.save(v, store_path) + # if not info['inverse']: + # v = torch.load(store_path, weights_only=True) + + # Save the features in the memory + + if info['inject'] and info['id'] > 19: + feature_name = str(info['t']) + '_' + str(info['second_order']) + '_' + str(info['id']) + '_' + info['type'] + '_' + 'V' + if info['inverse']: + info['feature'][feature_name] = v.clone() + else: + # Try to load feature, but continue if it doesn't exist + if feature_name in info['feature']: + v = info['feature'][feature_name] + + num_patches = v.size(2) + mask = torch.ones(num_patches, dtype=torch.bool, device=v.device) + mask[info['patch_ids']] = False + mask[:512] = False + keep_indices = mask.nonzero(as_tuple=True)[0] + if feature_name in info['feature']: + backup = info['feature'][feature_name] + v[:, :, keep_indices, :] = backup[:, :, keep_indices, :] + if len(info['patch_ref_ids']) != 0: + v[:, :, info['patch_ids'], :] = v[:, :, info['patch_ref_ids'], :] + + attn = attention(q, k, v, pe=pe) + # compute activation in mlp stream, cat again and run second linear layer + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output, info + + +class LastLayer(nn.Module): + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x: Tensor, vec: Tensor) -> Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + x = self.linear(x) + return x diff --git a/ArtiAgent - DefectDiffu/src/flux/sampling.py b/ArtiAgent - DefectDiffu/src/flux/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..f6bb29287c2bf83d19bd268d398bbd4e6b65ecac --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/sampling.py @@ -0,0 +1,390 @@ +import math +from typing import Callable + +import torch +from einops import rearrange, repeat +from torch import Tensor + +from .model import Flux +from .modules.conditioner import HFEmbedder + +def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str], + info=None) -> dict[str, Tensor]: + """ + Prepare inputs for the flux model with support for patch indices. + + Args: + t5, clip: Text encoders + img: Input image tensor + prompt: Text prompt(s) + info: Additional information dictionary, must contain 'artifact_data'. + + Returns: + Dictionary containing prepared inputs + """ + bs, c, h, w = img.shape + if bs == 1 and not isinstance(prompt, str): + bs = len(prompt) + img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + if img.shape[0] == 1 and bs > 1: + img = repeat(img, "1 ... -> bs ...", bs=bs) + + img_ids = torch.zeros(h // 2, w // 2, 3) + img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None] + img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :] + img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) + if isinstance(prompt, str): + prompt = [prompt] + txt = t5(prompt) + if txt.shape[0] == 1 and bs > 1: + txt = repeat(txt, "1 ... -> bs ...", bs=bs) + txt_ids = torch.zeros(bs, txt.shape[1], 3) + + vec = clip(prompt) + if vec.shape[0] == 1 and bs > 1: + vec = repeat(vec, "1 ... -> bs ...", bs=bs) + + patch_h, patch_w = h // 2, w // 2 + + # Add patch dimensions to info for model to use + info['patch_h'] = patch_h + info['patch_w'] = patch_w + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # NEW: Patchify reference latents for RAG visual conditioning + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + if info is not None and info.get('reference_latents'): + ref_list = info['reference_latents'] # list of [B, 16, H, W] tensors + ref_tokens_all = [] + ref_ids_all = [] + + for ref_lat in ref_list: + ref_bs, ref_c, ref_h, ref_w = ref_lat.shape + # Patchify same way as img: [B, 16, H, W] -> [B, (H/2)*(W/2), 64] + ref_p = rearrange(ref_lat, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + if ref_p.shape[0] == 1 and bs > 1: + ref_p = repeat(ref_p, "1 ... -> bs ...", bs=bs) + + # Position IDs matching img_ids pattern + ref_ids = torch.zeros(ref_h // 2, ref_w // 2, 3) + ref_ids[..., 1] = ref_ids[..., 1] + torch.arange(ref_h // 2)[:, None] + ref_ids[..., 2] = ref_ids[..., 2] + torch.arange(ref_w // 2)[None, :] + ref_ids = repeat(ref_ids, "h w c -> b (h w) c", b=ref_p.shape[0]) + + ref_tokens_all.append(ref_p) + ref_ids_all.append(ref_ids) + + # Store in info dict (will be added to return dict below) + info['ref_img'] = torch.cat(ref_tokens_all, dim=1).to(img.device, dtype=torch.bfloat16) + info['ref_img_ids'] = torch.cat(ref_ids_all, dim=1).to(img.device) + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + result = { + "img": img, + "img_ids": img_ids.to(img.device), + "txt": txt.to(device=img.device, dtype=torch.bfloat16), # <--- Cast to bfloat16 + "txt_ids": txt_ids.to(img.device), + "vec": vec.to(device=img.device, dtype=torch.bfloat16), # <--- Cast to bfloat16 + } + + # Add reference tensors if they were computed above + if info is not None and 'ref_img' in info: + result["ref_img"] = info['ref_img'] + result["ref_img_ids"] = info['ref_img_ids'] + + if "ref_img" in result: + print(f"[FLUX prepare] ref_img shape: {result['ref_img'].shape}") + + return result, (patch_h, patch_w) + + +def time_shift(mu: float, sigma: float, t: Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + +def get_lin_function( + x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15 +) -> Callable[[float], float]: + m = (y2 - y1) / (x2 - x1) + b = y1 - m * x1 + return lambda x: m * x + b + + +def get_schedule( + num_steps: int, + image_seq_len: int, + base_shift: float = 0.5, + max_shift: float = 1.15, + shift: bool = True, +) -> list[float]: + # extra step for zero + timesteps = torch.linspace(1, 0, num_steps + 1) + + # shifting the schedule to favor high timesteps for higher signal images + if shift: + # estimate mu based on linear estimation between two points + mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len) + timesteps = time_shift(mu, 1.0, timesteps) + + return timesteps.tolist() + +def denoise_first_order( + model: Flux, + # model input + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + vec: Tensor, + # sampling parameters + timesteps: list[float], + inverse, + info, + percentage_of_steps = 1.0, + guidance: float = 5.0, + ref_img: Tensor | None = None, # โ† ADD + ref_img_ids: Tensor | None = None, # โ† ADD +): + # this is ignored for schnell + inject_list = [True] * int(info['inject_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['inject_step'])) + attn_mask_list = [True] * int(info['attn_mask_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['attn_mask_step'])) + + # PE step lists for each artifact type + pe_step_addition_list = [True] * int(info['pe_step_addition']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_addition'])) + pe_step_removal_list = [True] * int(info['pe_step_removal']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_removal'])) + pe_step_distortion_list = [True] * int(info['pe_step_distortion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_distortion'])) + pe_step_fusion_list = [True] * int(info['pe_step_fusion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_fusion'])) + if inverse: + timesteps = timesteps[::-1] + inject_list = inject_list[::-1] + if percentage_of_steps != 1: + end_timestep_idx = int(len(timesteps) * percentage_of_steps) + if inverse: + timesteps = timesteps[:end_timestep_idx] + # inject_list = inject_list[:end_timestep_idx - 1] + else: + timesteps = timesteps[len(timesteps) - end_timestep_idx:] + # inject_list = inject_list[len(inject_list) - end_timestep_idx + 1:] + + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) + for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) + info['t'] = t_prev if inverse else t_curr + info['inverse'] = inverse + info['second_order'] = False + info['inject'] = inject_list[i] + info['attn_mask'] = attn_mask_list[i] + info['addition'] = pe_step_addition_list[i] + info['removal'] = pe_step_removal_list[i] + info['distortion'] = pe_step_distortion_list[i] + info['fusion'] = pe_step_fusion_list[i] + + pred, info = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec, + guidance=guidance_vec, + info=info, + # โ•โ•โ• ADD THESE TWO LINES โ•โ•โ• + ref_img=ref_img, + ref_img_ids=ref_img_ids, + ) + + img = img + (t_prev - t_curr) * pred + return img, info + +def denoise_fireflow( + model: Flux, + # model input + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + vec: Tensor, + # sampling parameters + timesteps: list[float], + inverse, + info, + percentage_of_steps = 1.0, + guidance: float = 5.0, + ref_img: Tensor | None = None, # โ† ADD + ref_img_ids: Tensor | None = None, # โ† ADD + ): + # this is ignored for schnell + inject_list = [True] * int(info['inject_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['inject_step'])) + attn_mask_list = [True] * int(info['attn_mask_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['attn_mask_step'])) + + # PE step lists for each artifact type + pe_step_addition_list = [True] * int(info['pe_step_addition']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_addition'])) + pe_step_removal_list = [True] * int(info['pe_step_removal']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_removal'])) + pe_step_distortion_list = [True] * int(info['pe_step_distortion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_distortion'])) + pe_step_fusion_list = [True] * int(info['pe_step_fusion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_fusion'])) + + if inverse: + timesteps = timesteps[::-1] + inject_list = inject_list[::-1] + if percentage_of_steps != 1: + end_timestep_idx = int(len(timesteps) * percentage_of_steps) + if inverse: + timesteps = timesteps[:end_timestep_idx] + # inject_list = inject_list[:end_timestep_idx - 1] + else: + timesteps = timesteps[len(timesteps) - end_timestep_idx:] + # inject_list = inject_list[len(inject_list) - end_timestep_idx + 1:] + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) + + step_list = [] + next_step_velocity = None + for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) + info['t'] = t_prev if inverse else t_curr + info['inverse'] = inverse + info['second_order'] = False + info['inject'] = inject_list[i] + info['attn_mask'] = attn_mask_list[i] + info['addition'] = pe_step_addition_list[i] + info['removal'] = pe_step_removal_list[i] + info['distortion'] = pe_step_distortion_list[i] + info['fusion'] = pe_step_fusion_list[i] + + if next_step_velocity is None: + pred, info = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec, + guidance=guidance_vec, + info=info, + # โ•โ•โ• ADD THESE TWO LINES โ•โ•โ• + ref_img=ref_img, + ref_img_ids=ref_img_ids, + ) + else: + pred = next_step_velocity + + img_mid = img + (t_prev - t_curr) / 2 * pred + + t_vec_mid = torch.full((img.shape[0],), t_curr + (t_prev - t_curr) / 2, dtype=img.dtype, device=img.device) + info['second_order'] = True + pred_mid, info = model( + img=img_mid, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec_mid, + guidance=guidance_vec, + info=info, + ref_img=ref_img, + ref_img_ids=ref_img_ids + ) + next_step_velocity = pred_mid + + img = img + (t_prev - t_curr) * pred_mid + + return img, info + + +def denoise( + model: Flux, + # model input + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + vec: Tensor, + # sampling parameters + timesteps: list[float], + inverse, + info, + percentage_of_steps = 1.0, + guidance: float = 4.0, + ref_img: Tensor | None = None, # โ† ADD + ref_img_ids: Tensor | None = None, # โ† ADD +): + # this is ignored for schnell + inject_list = [True] * int(info['inject_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['inject_step'])) + attn_mask_list = [True] * int(info['attn_mask_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['attn_mask_step'])) + + # PE step lists for each artifact type + pe_step_addition_list = [True] * int(info['pe_step_addition']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_addition'])) + pe_step_removal_list = [True] * int(info['pe_step_removal']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_removal'])) + pe_step_distortion_list = [True] * int(info['pe_step_distortion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_distortion'])) + pe_step_fusion_list = [True] * int(info['pe_step_fusion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_fusion'])) + + if inverse: + timesteps = timesteps[::-1] + inject_list = inject_list[::-1] + + if percentage_of_steps != 1: + end_timestep_idx = int(len(timesteps) * percentage_of_steps) + if inverse: + timesteps = timesteps[:end_timestep_idx] + # inject_list = inject_list[:end_timestep_idx - 1] + else: + timesteps = timesteps[len(timesteps) - end_timestep_idx:] + # inject_list = inject_list[len(inject_list) - end_timestep_idx + 1:] + + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) + for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) + info['t'] = t_prev if inverse else t_curr + info['inverse'] = inverse + info['second_order'] = False + info['inject'] = inject_list[i] + info['attn_mask'] = attn_mask_list[i] + info['addition'] = pe_step_addition_list[i] + info['removal'] = pe_step_removal_list[i] + info['distortion'] = pe_step_distortion_list[i] + info['fusion'] = pe_step_fusion_list[i] + + pred, info = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec, + guidance=guidance_vec, + info=info, + # โ•โ•โ• ADD THESE TWO LINES โ•โ•โ• + ref_img=ref_img, + ref_img_ids=ref_img_ids + ) + + + img_mid = img + (t_prev - t_curr) / 2 * pred + + t_vec_mid = torch.full((img.shape[0],), (t_curr + (t_prev - t_curr) / 2), dtype=img.dtype, device=img.device) + info['second_order'] = True + pred_mid, info = model( + img=img_mid, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec_mid, + guidance=guidance_vec, + info=info + ) + + first_order = (pred_mid - pred) / ((t_prev - t_curr) / 2) + img = img + (t_prev - t_curr) * pred + 0.5 * (t_prev - t_curr) ** 2 * first_order + + return img, info + + +def unpack(x: Tensor, height: int, width: int) -> Tensor: + return rearrange( + x, + "b (h w) (c ph pw) -> b c (h ph) (w pw)", + h=math.ceil(height / 16), + w=math.ceil(width / 16), + ph=2, + pw=2, + ) diff --git a/ArtiAgent - DefectDiffu/src/flux/util.py b/ArtiAgent - DefectDiffu/src/flux/util.py new file mode 100644 index 0000000000000000000000000000000000000000..c3f8e8092462a72bcb75af58e716f1b87a6751d8 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/flux/util.py @@ -0,0 +1,346 @@ +import os +from dataclasses import dataclass + +import torch +from einops import rearrange +from huggingface_hub import hf_hub_download +# from imwatermark import WatermarkEncoder +from safetensors.torch import load_file as load_sft + +from flux.model import Flux, FluxParams +from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams +from flux.modules.conditioner import HFEmbedder +from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel, + T5Tokenizer, BitsAndBytesConfig) # <--- Added BitsAndBytesConfig + + +@dataclass +class ModelSpec: + params: FluxParams + ae_params: AutoEncoderParams + ckpt_path: str | None + ae_path: str | None + repo_id: str | None + repo_flow: str | None + repo_ae: str | None + +configs = { + "flux-dev": ModelSpec( + repo_id="black-forest-labs/FLUX.1-dev", + repo_flow="flux1-dev.safetensors", + repo_ae=None, + ckpt_path=os.getenv("FLUX_DEV"), + params=FluxParams( + in_channels=64, + out_channels=64, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=True, + ), + ae_path=os.getenv("AE"), + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), + "flux-fill-dev": ModelSpec( + repo_id="black-forest-labs/FLUX.1-Fill-dev", + repo_flow="flux1-fill-dev.safetensors", + repo_ae="ae.safetensors", + ckpt_path=os.getenv("FLUX_FILL_DEV"), + params=FluxParams( + in_channels=64, + out_channels=384, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=True, + ), + ae_path=os.getenv("AE"), + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), + "flux-kontext-dev": ModelSpec( + repo_id="black-forest-labs/FLUX.1-Kontext-dev", + repo_flow="flux1-kontext-dev.safetensors", + repo_ae="ae.safetensors", + ckpt_path=os.getenv("FLUX_FILL_DEV"), + params=FluxParams( + in_channels=64, + out_channels=64, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=True, + ), + ae_path=os.getenv("AE"), + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), + "flux-schnell": ModelSpec( + repo_id="black-forest-labs/FLUX.1-schnell", + repo_flow="flux1-schnell.safetensors", + repo_ae="black-forest-labs/FLUX.1-schnell", + ckpt_path=os.getenv("FLUX_SCHNELL"), + params=FluxParams( + in_channels=64, # ArtiAgent custom input dimension logic handled in load_flow_model + out_channels=64, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10000.0, + qkv_bias=True, + guidance_embed=False, + ), + ae_path="ae.safetensors", + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), +} + + +def print_load_warning(missing: list[str], unexpected: list[str]) -> None: + if len(missing) > 0 and len(unexpected) > 0: + print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing)) + print("\n" + "-" * 79 + "\n") + print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected)) + elif len(missing) > 0: + print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing)) + elif len(unexpected) > 0: + print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected)) + +def _replace_linear_with_4bit(module, compute_dtype=torch.bfloat16): + """Recursively replace all nn.Linear with bitsandbytes 4-bit layers""" + import bitsandbytes as bnb + for name, child in module.named_children(): + if name == "img_in": + continue # Skip img_in to preserve ArtiAgent's custom shape handling + if isinstance(child, torch.nn.Linear): + has_bias = child.bias is not None + new_layer = bnb.nn.Linear4bit( + child.in_features, + child.out_features, + bias=has_bias, + compute_dtype=compute_dtype, + compress_statistics=True, + quant_type="nf4", + ) + new_layer.weight = bnb.nn.Params4bit( + child.weight.data, + requires_grad=False, + quant_type="nf4", + ) + if has_bias: + new_layer.bias = torch.nn.Parameter(child.bias.data) + setattr(module, name, new_layer) + else: + _replace_linear_with_4bit(child, compute_dtype) + +def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True): + # Loading Flux + print("Init model") + + ckpt_path = configs[name].ckpt_path + if ( + ckpt_path is None + and configs[name].repo_id is not None + and configs[name].repo_flow is not None + and hf_download + ): + ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow) + + # Initialize model directly on CPU or target device (avoids meta-tensor shape replacement) + target_device = torch.device(device) + model = Flux(configs[name].params).to(dtype=torch.bfloat16) + + if ckpt_path is not None: + print("Loading checkpoint") + # load_sft doesn't support torch.device + sd = load_sft(ckpt_path, device="cpu") + + # --- ADD THIS LINE TO STRIP FP8 / COMFYUI KEY PREFIXES --- + sd = {k.replace("model.diffusion_model.", ""): v for k, v in sd.items()} + # --------------------------------------------------------- + + # --- FIX: HANDLE EXPANDED IMG_IN (384 channels vs 64 channels) --- + img_in_weight = sd.pop("img_in.weight", None) + img_in_bias = sd.pop("img_in.bias", None) + + # Load all standard layers safely + missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) + print_load_warning(missing, unexpected) + + # Copy base 64 channels into ArtiAgent's expanded 384-channel input layer + # In src/flux/util.py inside load_flow_model(): + + if img_in_weight is not None: + with torch.no_grad(): + w = img_in_weight.to(device=device, dtype=torch.bfloat16) + + # Check if model.img_in weight expects 384 channels while checkpoint has 64 + if model.img_in.weight.shape[1] != w.shape[1]: + # Slice model.img_in.weight to match the 64-channel input tensor + model.img_in.weight = torch.nn.Parameter(model.img_in.weight[:, :w.shape[1]]) + + model.img_in.weight.copy_(w) + + if img_in_bias is not None and getattr(model.img_in, "bias", None) is not None: + with torch.no_grad(): + b = img_in_bias.to(device=device, dtype=torch.bfloat16) + model.img_in.bias.copy_(b) + + # Quantize all Linear layers to NF4 on CPU before moving to GPU + print("Quantizing model to NF4 (this may take a minute)...") + _replace_linear_with_4bit(model, compute_dtype=torch.bfloat16) + print("NF4 quantization complete.") + + # Move model to target CUDA device + model = model.to(target_device) + return model + + +def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder: + # Force T5 onto CPU; sampling.py already moves the encoded txt tensor to GPU + return HFEmbedder( + "google/t5-v1_1-xxl", + max_length=max_length, + is_clip=False, + torch_dtype=torch.bfloat16, + device_map="cpu" + ) + + +def load_clip(device: str | torch.device = "cuda") -> HFEmbedder: + # Keep on CPU; sampling.py moves vec to GPU after encoding + return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, is_clip=True, torch_dtype=torch.bfloat16) + + +def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder: + ckpt_path = configs[name].ae_path + + # If ckpt_path is just a filename and doesn't exist locally, download it + if ckpt_path is not None and not os.path.exists(ckpt_path) and hf_download: + repo_id = configs[name].repo_ae or configs[name].repo_id + ckpt_path = hf_hub_download(repo_id, ckpt_path) + elif ckpt_path is None and configs[name].repo_id is not None and hf_download: + repo_id = configs[name].repo_ae or configs[name].repo_id + ckpt_path = hf_hub_download(repo_id, "ae.safetensors") + + # Loading the autoencoder + print("Init AE") + + # Initialize directly on CPU to avoid meta-tensor initialization issues + ae = AutoEncoder(configs[name].ae_params) + + if ckpt_path is not None: + sd = load_sft(ckpt_path, device=str(device)) + missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True) + print_load_warning(missing, unexpected) + + ae = ae.to(device) + return ae + + +# class WatermarkEmbedder: +# def __init__(self, watermark): +# self.watermark = watermark +# self.num_bits = len(WATERMARK_BITS) +# self.encoder = WatermarkEncoder() +# self.encoder.set_watermark("bits", self.watermark) + +# def __call__(self, image: torch.Tensor) -> torch.Tensor: +# """ +# Adds a predefined watermark to the input image + +# Args: +# image: ([N,] B, RGB, H, W) in range [-1, 1] + +# Returns: +# same as input but watermarked +# """ +# image = 0.5 * image + 0.5 +# squeeze = len(image.shape) == 4 +# if squeeze: +# image = image[None, ...] +# n = image.shape[0] +# image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1] +# # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255] +# # watermarking libary expects input as cv2 BGR format +# for k in range(image_np.shape[0]): +# image_np[k] = self.encoder.encode(image_np[k], "dwtDct") +# image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to( +# image.device +# ) +# image = torch.clamp(image / 255, min=0.0, max=1.0) +# if squeeze: +# image = image[0] +# image = 2 * image - 1 +# return image + + +# # A fixed 48-bit message that was chosen at random +# WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110 +# # bin(x)[2:] gives bits of x as str, use int to convert them to 0/1 +# WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]] +# embed_watermark = WatermarkEmbedder(WATERMARK_BITS) diff --git a/ArtiAgent - DefectDiffu/src/loop-file.py b/ArtiAgent - DefectDiffu/src/loop-file.py new file mode 100644 index 0000000000000000000000000000000000000000..efd848f76e60a08cbe40e368f809cc34d1a263ec --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/loop-file.py @@ -0,0 +1,44 @@ +import subprocess +import sys + +# Common base arguments across all runs +BASE_CMD = [ + sys.executable, # Uses the currently active Python interpreter + "batch_agent_orchestrator.py", + "--input-dir", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\TestingImage\01_bubble_sensor", + "--output-dir", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\ArtiAgent - Defect\agent_output", + "--defectdiffu-ckpt", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\ArtiAgent - Defect\engine\DefectDiffu\checkpoint_old-4\model_300.pth", + "--vae-path", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\ArtiAgent - Defect\engine\DefectDiffu\checkpoints\sd-vae-ft-mse", + "--product-desc", "VCSEL laser diode with emission aperture and surrounding mesa", + "--max-defects-per-image", "3", + "--device", "cuda" +] + +# Variations for each step in a single loop +TASKS = [ + ["--defect-type", "tiger-strip"], + ["--defect-type", "bubble"], + [] # Default/No defect-type specified +] + +TOTAL_LOOPS = 16 + +if __name__ == "__main__": + for loop_num in range(1, TOTAL_LOOPS + 1): + print(f"\n" + "=" * 50) + print(f"๐Ÿš€ STARTING LOOP {loop_num} OF {TOTAL_LOOPS}") + print("=" * 50) + + for step_idx, task_args in enumerate(TASKS, start=1): + full_command = BASE_CMD + task_args + print(f"\n--> Running Task {step_idx}/3 for Loop {loop_num}...") + + # Run command and block until it completes + result = subprocess.run(full_command) + + # Check if command failed + if result.returncode != 0: + print(f"โŒ Error encountered on Loop {loop_num}, Task {step_idx}. Execution stopped.") + sys.exit(result.returncode) + + print("\n๐ŸŽ‰ Success! Completed all 16 loops.") \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/README.md b/ArtiAgent - DefectDiffu/src/pipeline/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ae88fed6ed971ca929c7cf2eb66d0ebf04e606b9 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/README.md @@ -0,0 +1,319 @@ +# Image Artifacts Pipeline + +A modular pipeline for generating synthetic image artifacts using GSAM (Grounded Segment Anything Model) for part detection and FLUX diffusion model for artifact generation. + +## Overview + +This pipeline provides clean, modular Python components for the two-stage artifact generation workflow: + +``` +๐Ÿ“ฆ pipeline/ +โ”œโ”€โ”€ ๐Ÿ“„ __init__.py # Package initialization +โ”œโ”€โ”€ ๐Ÿ“„ data_loader.py # Dataset handling (COCO, ImageNet, Custom) +โ”œโ”€โ”€ ๐Ÿ“„ gsam_detector.py # GSAM model integration (GroundingDINO + SAM) +โ”œโ”€โ”€ ๐Ÿ“„ instance_processor.py # Instance filtering and bbox operations +โ”œโ”€โ”€ ๐Ÿ“„ flux_generator.py # FLUX model operations +โ”œโ”€โ”€ ๐Ÿ“„ visualization.py # Image visualization utilities +โ”œโ”€โ”€ ๐Ÿ“„ prompts.py # OpenAI API utilities for vocabulary generation +โ””โ”€โ”€ ๐Ÿ“„ README.md # This file +``` + +## Architecture + +``` +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ COCO/ImageNet โ”‚ +โ”‚ Custom Dataset โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ Data Loader โ”‚ +โ”‚ (data_loader) โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ GSAM Detector โ”‚โ—„โ”€โ”€โ”€โ”€โ”ค OpenAI API โ”‚ +โ”‚ (GroundingDINO โ”‚ โ”‚ (Vocabulary)โ”‚ +โ”‚ + SAM/SAM-HQ) โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚Instance Processorโ”‚ +โ”‚ (Filter, Sample,โ”‚ +โ”‚ Create Patches)โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ FLUX Generator โ”‚ +โ”‚ (Diffusion with โ”‚ +โ”‚ Patch Guidance) โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ Visualizer โ”‚ +โ”‚ (Results/Masks) โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +``` + +## Components + +### 1. Data Loaders + +Handles dataset loading and image sampling for COCO, ImageNet, and custom directories. + +```python +from pipeline import COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader + +# COCO Dataset +coco_loader = COCODataLoader( + dataset_path="/path/to/coco/annotations", + image_path="/path/to/coco/images" +) +cat_ids = coco_loader.get_category_ids(['person']) +img_info, img_array, caption = coco_loader.sample_image_by_category(cat_ids) + +# ImageNet Dataset +imagenet_loader = ImageNetDataLoader( + dataset_path="/path/to/imagenet", + split="train" +) + +# Custom Directory +custom_loader = CustomDirectoryDataLoader( + directory_path="/path/to/images" +) +``` + +### 2. GSAMDetector + +Integrates GroundingDINO and SAM/SAM-HQ for part detection with OpenAI-generated vocabulary. + +```python +from pipeline import GSAMDetector + +detector = GSAMDetector( + grounding_config="GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py", + grounding_checkpoint="weight/groundingdino_swint_ogc.pth", + sam_checkpoint="weight/sam_vit_h_4b8939.pth", + sam_version="vit_h", + use_sam_hq=False, + box_threshold=0.3, + text_threshold=0.25, + device='cuda' +) + +# Detect parts in image +entity_predictions, subentity_predictions, vis_output = detector.detect_parts( + image=img_array, + entities=['person'], + subentities=['person head', 'person arm', 'person leg'], + entity_subentity_mapping={'person': ['person head', 'person arm', 'person leg']}, + min_area_ratio=0.005, + max_area_ratio=0.5 +) +``` + +### 3. InstanceProcessor + +Handles instance filtering, sampling, bounding box operations, and patch annotation creation. + +```python +from pipeline import InstanceProcessor + +# Sample instance by confidence score +sampled_instance = InstanceProcessor.sample_instance_by_score( + predictions, + min_area_ratio=0.01, + max_area_ratio=0.5 +) + +# Generate bbox suggestions for addition artifacts +suggested_bbox = InstanceProcessor.generate_bbox_suggestion( + predictions=predictions, + reference_bbox=reference_bbox, + class_name=class_name, + vocab=vocab, + max_ref_overlap=0.3, + min_entity_overlap=0.1 +) + +# Create annotation with patch indices +annotation_data = InstanceProcessor.create_annotation_dict( + instance=sampled_instance, + img_shape=image.shape, + artifact_type='distortion', + patch_size=16 +) +``` + +### 4. FluxGenerator + +Manages FLUX diffusion model operations for artifact generation with patch-based guidance. + +```python +from pipeline import FluxGenerator, FluxConfig + +# Configure FLUX model +config = FluxConfig( + name='flux-dev', + guidance=5.0, + num_steps=25, + pe_step=0.5, # Position encoding step + seed=42 +) + +# Artifact-type-specific PE steps +config_advanced = FluxConfig( + name='flux-dev', + guidance=5.0, + num_steps=25, + pe_step={ + 'addition': 0.3, + 'removal': 0.3, + 'distortion': 0.5 + }, + seed=42 +) + +generator = FluxGenerator(device='cuda', config=config_advanced) + +# Generate artifact image +generated_image = generator.generate_with_artifacts( + source_prompt="a photo of a person", + target_prompt="a photo of a person", + bbox=target_bbox, + bbox_ref=reference_bbox, + artifact_type='distortion', + source_img=image +) +``` + +### 5. ImageVisualizer + +Provides visualization utilities for debugging and quality assurance. + +```python +from pipeline import ImageVisualizer + +visualizer = ImageVisualizer() + +# Show single image with caption +visualizer.show_image(image, caption, title="Original", base_dir="output/") + +# Show comparison +visualizer.show_comparison( + original_image, + generated_image, + artifact_data, + caption="Distortion Artifact", + base_dir="output/" +) + +# Show bounding box overlay +visualizer.show_bbox_overlay( + image, + target_bbox, + base_dir="output/", + filename="bbox_overlay.png" +) + +# Show patch masks +visualizer.show_patch_masks( + image, + reference_patches, + target_patches, + base_dir="output/" +) +``` + +## Artifact Types + +The pipeline supports three types of image artifacts: + +### 1. Distortion +Modifies the appearance of existing parts while keeping them in place. +- Uses reference patches to guide where distortion is applied +- Configurable distortion kernels (jitter, swirl, voronoi) + +### 2. Removal +Removes detected parts from images naturally. +- Uses reference patches to identify removal areas +- FLUX inpainting fills removed areas contextually + +### 3. Addition +Adds new instances of detected parts in suitable locations. +- Uses reference patches as source templates +- Generates target patches using IoU-based intelligent placement +- Maintains visual consistency with surrounding context + +## Configuration + +### Detection Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `box_threshold` | `0.3` | Detection confidence threshold | +| `text_threshold` | `0.25` | Text-image matching threshold | +| `min_area_ratio` | `0.005` | Minimum part size (0.5% of image) | +| `max_area_ratio` | `0.5` | Maximum part size (50% of image) | +| `nms_threshold` | `0.5` | Non-maximum suppression threshold | + +### Generation Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `guidance` | `5.0` | Guidance scale for FLUX | +| `num_steps` | `25` | Number of diffusion steps | +| `pe_step` | `0.5` | Position encoding step size | +| `inject` | `25` | Injection step in diffusion | +| `seed` | `42` | Random seed for reproducibility | + +## Dependencies + +- `torch` - PyTorch for deep learning +- `openai` - For vocabulary generation +- `pycocotools` - For COCO dataset handling +- `supervision` - For detection utilities +- `groundingdino` - For grounded detection +- `segment_anything` - For segmentation +- `matplotlib` - For visualization +- `PIL` - For image processing +- `numpy` - For numerical operations + +## Installation + +1. Install GroundingDINO and SAM following their respective installation guides +2. Download model weights and place them in `src/weight/` directory +3. Set up OpenAI API key: `export OPENAI_API_KEY='your-key'` +4. Install required dependencies + +See the main README for detailed installation instructions. + +## Usage + +See `batch_gsam_segmentation.py` and `batch_flux_generation.py` for complete batch processing examples. + +## Performance Tips + +1. **GPU Memory**: Use SAM `vit_b` for limited GPU memory, `vit_h` for best quality +2. **Filtering**: Adjust area ratios to balance quality vs. quantity +3. **Speed**: Lower `num_steps` (15-20) for faster generation +4. **Quality**: Higher `num_steps` (25-35) for better results + +## Troubleshooting + +### Common Issues + +1. **GSAM setup issues**: Ensure GroundingDINO and SAM weights are downloaded +2. **OpenAI API errors**: Check API key and rate limits +3. **COCO dataset errors**: Verify dataset paths and structure +4. **GPU memory issues**: Use smaller SAM model or reduce batch size + +## License + +This pipeline integrates multiple open-source components, each with their own licenses. See the main repository LICENSE and model_licenses/ directory for details. diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__init__.py b/ArtiAgent - DefectDiffu/src/pipeline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5fe885cb6ab84b2cbc29196b9ff9137fe157ff35 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/__init__.py @@ -0,0 +1,38 @@ +""" +Image Artifacts Pipeline + +This module provides detection and processing pipeline components for generating +image artifacts using various segmentation models. +""" + +# Import GSAMDetector if gsam dependencies exist +try: + from .gsam_detector import GSAMDetector +except ImportError: + GSAMDetector = None + +# Import Flux components if flux dependencies exist +try: + from .flux_generator import FluxGenerator, FluxConfig +except ImportError: + FluxGenerator = None + FluxConfig = None + + +try: + from .data_loader import COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader +except ImportError: + COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader = None, None, None + +from .instance_processor import InstanceProcessor + + +__all__ = [ + 'GSAMDetector', + 'FluxGenerator', + 'FluxConfig', + 'COCODataLoader', + 'ImageNetDataLoader', + 'CustomDirectoryDataLoader', + 'InstanceProcessor', +] \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d78504750b49843a5f2db8319acdbcf7ad3cae23 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/__init__.cpython-311.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4d9bf0093a7c1cc2d9b9876e45aa4ed59a4b452 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/__init__.cpython-311.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/data_loader.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/data_loader.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a5ac44e13dfd72bfa22f1e97874e64304100073e Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/data_loader.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/data_loader.cpython-311.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/data_loader.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..473fe9fd556d225594c2b8a68d5aedc5dc10b6c7 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/data_loader.cpython-311.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/defect_rag.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/defect_rag.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80f51cd7cdc94258049b55efd3857eea9b917825 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/defect_rag.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/defectdiffu_generator.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/defectdiffu_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9ab2fd99095fa93719a614d764ac0152cada728 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/defectdiffu_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/domain_router.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/domain_router.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7bc285e60d99a509d2e9eee3fa41da4ee63f7a5e Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/domain_router.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/flux_generator.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/flux_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a73586f58dace3605d0fda0e13ee2730641ae625 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/flux_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/flux_generator.cpython-311.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/flux_generator.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8624a69492ad38e34059020ab04204f90011c20 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/flux_generator.cpython-311.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/gsam_detector.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/gsam_detector.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..655c50f390a59efa14a82e6167e4e738b163908d Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/gsam_detector.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/gsam_detector.cpython-311.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/gsam_detector.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad080bb5978f041e8611c69460caa534674ce6bf Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/gsam_detector.cpython-311.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/instance_processor.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/instance_processor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e1e3cc6986fd5de2a1b77ff3dc672c3d40485cf Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/instance_processor.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/instance_processor.cpython-311.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/instance_processor.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9486a6a6b5c41bffe9d70fe5389a2134c35e88b Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/instance_processor.cpython-311.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/local_vlm_client.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/local_vlm_client.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c207a38c827deb328374f2e63c14b40e3fe9c1a Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/local_vlm_client.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/prompts.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/prompts.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c653701eb8a95240628db59971b80b69d6c89f76 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/pipeline/__pycache__/prompts.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/pipeline/data_loader.py b/ArtiAgent - DefectDiffu/src/pipeline/data_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..05ad8e165ecb1c753161f67452cf91bb56c43402 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/data_loader.py @@ -0,0 +1,732 @@ +import os +import numpy as np +from pycocotools.coco import COCO +from typing import List, Dict, Tuple, Optional, Union +import pathlib +import json +import glob +from PIL import Image +import logging +from typing import Any + + +def preprocess_image_for_flux(image_path_or_pil: Union[str, Image.Image]) -> np.ndarray: + """ + Shared image preprocessing function for flux model compatibility + + Args: + image_path_or_pil: Either a file path to image or PIL Image object + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + # Load image with PIL if path provided + if isinstance(image_path_or_pil, str): + img = Image.open(image_path_or_pil) + else: + img = image_path_or_pil + + if img.mode != 'RGB': + img = img.convert('RGB') + + # Rescale if shortest side is less than 480 + width, height = img.size + if min(width, height) < 480: + scale_factor = 480 / min(width, height) + new_width = int(width * scale_factor) + new_height = int(height * scale_factor) + img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + img_array = np.array(img) + + # Ensure dimensions are divisible by 16 for flux model compatibility + shape = img_array.shape + new_h = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16 + new_w = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16 + + # Crop image to new dimensions + img_array = img_array[:new_h, :new_w, :] + + return img_array + + +class COCODataLoader: + """Handler for COCO dataset loading and image sampling""" + + def __init__(self, dataset_path: str, image_path: str): + """ + Initialize COCO data loader + + Args: + dataset_path: Path to COCO annotations directory + image_path: Path to COCO images directory + """ + self.dataset_path = dataset_path + self.image_path = image_path + + # Load COCO annotations + self.caption_file = os.path.join(dataset_path, "captions_train2017.json") + self.class_file = os.path.join(dataset_path, "instances_train2017.json") + + self.coco_cap = COCO(self.caption_file) + self.coco_class = COCO(self.class_file) + + # Get all image IDs + self.image_ids = self.coco_cap.getImgIds() + + def get_category_ids(self, super_categories: List[str]) -> List[int]: + """ + Get category IDs for given super categories + + Args: + super_categories: List of super category names (e.g., ['person', 'animal']) + + Returns: + List of category IDs + """ + cat_ids = self.coco_class.getCatIds(supNms=super_categories) + return cat_ids + + def get_category_names(self, cat_ids: List[int]) -> List[str]: + """Get category names from category IDs""" + cats = self.coco_class.loadCats(cat_ids) + return [cat['name'] for cat in cats] + + def sample_image_by_category(self, cat_ids: List[int]) -> Tuple[Dict, np.ndarray, str]: + """ + Sample a random image containing objects from specified categories + + Args: + cat_ids: List of category IDs to sample from + + Returns: + Tuple of (image_info, image_array, caption) with image dimensions adjusted to be divisible by 16 + """ + # Get images containing specified categories + img_ids = self.coco_class.getImgIds(catIds=cat_ids[0]) # Use first category for sampling + + # Sample random image + sampled_id = img_ids[np.random.randint(0, len(img_ids))] + + # Load image info and array + img_info = self.coco_class.loadImgs(sampled_id)[0] + + # Load and preprocess image + img_path = os.path.join(self.image_path, img_info['file_name']) + img_array = preprocess_image_for_flux(img_path) + + # Get caption + ann_ids = self.coco_cap.getAnnIds(imgIds=img_info['id']) + anns = self.coco_cap.loadAnns(ann_ids) + caption = anns[0]['caption'] if anns else "" + + return img_info, img_array, caption + + def get_image_categories(self, img_info: Dict) -> List[str]: + """ + Get all category names present in an image + + Args: + img_info: Image information dictionary + + Returns: + List of unique category names in the image + """ + # Get category information for the image + ann_ids_class = self.coco_class.getAnnIds(imgIds=img_info['id']) + anns_class = self.coco_class.loadAnns(ann_ids_class) + + # Extract category IDs from annotations + cat_ids_in_image = [ann['category_id'] for ann in anns_class] + + # Get category names + categories_in_image = [] + for cat_id in cat_ids_in_image: + cat_info = self.coco_class.loadCats([cat_id])[0] + categories_in_image.append(cat_info['name']) + + # Remove duplicates and return + return list(set(categories_in_image)) + + def load_image_by_info(self, img_info: Dict) -> np.ndarray: + """ + Load image array from image info dictionary + + Args: + img_info: COCO image info dictionary + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + img_path = os.path.join(self.image_path, img_info['file_name']) + return preprocess_image_for_flux(img_path) + + def get_image_caption(self, img_info: Dict) -> str: + """ + Get caption for a specific image + + Args: + img_info: COCO image info dictionary + + Returns: + Image caption string + """ + ann_ids = self.coco_cap.getAnnIds(imgIds=img_info['id']) + anns = self.coco_cap.loadAnns(ann_ids) + caption = anns[0]['caption'] if anns else "" + return caption + + def create_category_directories(self, category_names: List[str], base_path: str = 'data/coco_2017_extracted'): + """Create directories for each category""" + for category in category_names: + pathlib.Path(f'{base_path}/{category}').mkdir(parents=True, exist_ok=True) + + +class ImageNetDataLoader: + """Handler for ImageNet dataset loading and image sampling""" + + def __init__(self, dataset_path: str, split: str = 'train'): + """ + Initialize ImageNet data loader + + Args: + dataset_path: Path to ImageNet dataset directory + split: Dataset split ('train' or 'val') + """ + self.dataset_path = dataset_path + self.split = split + self.split_path = os.path.join(dataset_path, split) + + # Load class mapping if available + self.class_mapping = self._load_class_mapping() + + # Get all synset directories + self.synsets = [d for d in os.listdir(self.split_path) + if os.path.isdir(os.path.join(self.split_path, d))] + + # Build image index + self._build_image_index() + + def _load_class_mapping(self) -> Dict[str, str]: + """ + Load class mapping from synset IDs to human-readable names + + Returns: + Dictionary mapping synset IDs to class names + """ + mapping_files = [ + os.path.join(self.dataset_path, 'imagenet_class_index.json'), + os.path.join(self.dataset_path, 'synset_words.txt'), + os.path.join(self.dataset_path, 'LOC_synset_mapping.txt') + ] + + class_mapping = {} + + # Try loading from JSON format first + for mapping_file in mapping_files: + if os.path.exists(mapping_file): + if mapping_file.endswith('.json'): + with open(mapping_file, 'r') as f: + data = json.load(f) + for idx, (synset, name) in data.items(): + class_mapping[synset] = name + break + elif mapping_file.endswith('.txt'): + with open(mapping_file, 'r') as f: + for line in f: + parts = line.strip().split('\t') + if len(parts) >= 2: + synset = parts[0] + name = parts[1] + class_mapping[synset] = name + break + + return class_mapping + + def _build_image_index(self): + """Build index of all images in the dataset""" + self.image_index = {} + + for synset in self.synsets: + synset_path = os.path.join(self.split_path, synset) + image_files = [] + + # Support common image formats + for ext in ['*.JPEG', '*.jpg', '*.jpeg', '*.png', '*.bmp']: + image_files.extend(glob.glob(os.path.join(synset_path, ext))) + + self.image_index[synset] = image_files + + def get_class_names(self) -> List[str]: + """ + Get all available class names + + Returns: + List of class names (human-readable if mapping available, else synset IDs) + """ + if self.class_mapping: + return [self.class_mapping.get(synset, synset) for synset in self.synsets] + else: + return self.synsets + + def get_synsets(self) -> List[str]: + """Get all available synset IDs""" + return self.synsets + + def sample_image_by_class(self, class_names: List[str] = None, synsets: List[str] = None) -> Tuple[Dict, np.ndarray, str]: + """ + Sample a random image from specified classes or synsets + + Args: + class_names: List of human-readable class names to sample from + synsets: List of synset IDs to sample from (takes precedence over class_names) + + Returns: + Tuple of (image_info, image_array, class_name) + """ + # Determine synsets to sample from + if synsets: + target_synsets = [s for s in synsets if s in self.synsets] + elif class_names: + # Convert class names to synsets + target_synsets = [] + for class_name in class_names: + for synset, mapped_name in self.class_mapping.items(): + if mapped_name.lower() == class_name.lower() and synset in self.synsets: + target_synsets.append(synset) + else: + # Sample from all available synsets + target_synsets = self.synsets + + if not target_synsets: + raise ValueError("No matching synsets found for the specified classes") + + # Sample random synset + sampled_synset = np.random.choice(target_synsets) + + # Sample random image from the synset + if not self.image_index[sampled_synset]: + raise ValueError(f"No images found for synset {sampled_synset}") + + sampled_image_path = np.random.choice(self.image_index[sampled_synset]) + + # Load image + img_array = self._load_and_preprocess_image(sampled_image_path) + + # Create image info + img_info = { + 'file_name': os.path.basename(sampled_image_path), + 'file_path': sampled_image_path, + 'synset': sampled_synset, + 'class_name': self.class_mapping.get(sampled_synset, sampled_synset), + 'height': img_array.shape[0], + 'width': img_array.shape[1] + } + + class_name = self.class_mapping.get(sampled_synset, sampled_synset) + + return img_info, img_array, class_name + + def load_image_by_path(self, image_path: str) -> np.ndarray: + """ + Load image from file path with preprocessing + + Args: + image_path: Path to image file + + Returns: + Preprocessed image array + """ + return self._load_and_preprocess_image(image_path) + + def _load_and_preprocess_image(self, image_path: str) -> np.ndarray: + """ + Load and preprocess image for flux model compatibility + + Args: + image_path: Path to image file + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + return preprocess_image_for_flux(image_path) + + def get_images_by_synset(self, synset: str) -> List[str]: + """ + Get all image paths for a specific synset + + Args: + synset: Synset ID + + Returns: + List of image paths + """ + return self.image_index.get(synset, []) + + def get_synset_stats(self) -> Dict[str, int]: + """ + Get statistics about number of images per synset + + Returns: + Dictionary mapping synset IDs to image counts + """ + return {synset: len(images) for synset, images in self.image_index.items()} + + def create_class_directories(self, class_names: List[str], base_path: str = 'data/imagenet_extracted'): + """ + Create directories for each class + + Args: + class_names: List of class names or synsets + base_path: Base directory to create class folders in + """ + for class_name in class_names: + # Use synset as folder name if it exists, otherwise use class name + if class_name in self.synsets: + folder_name = class_name + else: + # Find synset for class name + folder_name = class_name + for synset, mapped_name in self.class_mapping.items(): + if mapped_name.lower() == class_name.lower(): + folder_name = synset + break + + pathlib.Path(f'{base_path}/{folder_name}').mkdir(parents=True, exist_ok=True) + + +class CustomDirectoryDataLoader: + """Handler for custom directory structure with images directly in a single directory""" + + def __init__(self, dataset_path: str): + """ + Initialize custom directory data loader + + Args: + dataset_path: Path to directory containing images directly + Expected structure: dataset_path/*.jpg, dataset_path/*.png, etc. + """ + self.dataset_path = dataset_path + + if not os.path.exists(dataset_path): + raise ValueError(f"Dataset path does not exist: {dataset_path}") + + # Build image index from directory + self._build_image_index() + + if not self.image_paths: + raise ValueError(f"No images found in {dataset_path}") + + def _build_image_index(self): + """Build index of all images in the directory""" + self.image_paths = [] + + # Support common image formats + for ext in ['*.jpg', '*.jpeg', '*.JPG', '*.JPEG', '*.png', '*.PNG', + '*.bmp', '*.BMP', '*.tiff', '*.TIFF', '*.tif', '*.TIF']: + self.image_paths.extend(glob.glob(os.path.join(self.dataset_path, ext))) + + self.image_paths.sort() # Sort for consistent ordering + + def get_image_count(self) -> int: + """ + Get total number of images in the directory + + Returns: + Number of images + """ + return len(self.image_paths) + + def get_all_image_paths(self) -> List[str]: + """ + Get all image paths in the directory + + Returns: + List of image paths + """ + return self.image_paths.copy() + + def sample_random_image(self) -> Tuple[Dict, np.ndarray]: + """ + Sample a random image from the directory + + Returns: + Tuple of (image_info, image_array) + """ + if not self.image_paths: + raise ValueError("No images available to sample") + + # Sample random image path + sampled_image_path = np.random.choice(self.image_paths) + + # Load and preprocess image + img_array = self._load_and_preprocess_image(sampled_image_path) + + # Create image info + img_info = { + 'file_name': os.path.basename(sampled_image_path), + 'file_path': sampled_image_path, + 'height': img_array.shape[0], + 'width': img_array.shape[1] + } + + return img_info, img_array + + def sample_multiple_images(self, num_samples: int = 1) -> List[Tuple[Dict, np.ndarray]]: + """ + Sample multiple images from the directory + + Args: + num_samples: Number of images to sample + + Returns: + List of tuples (image_info, image_array) + """ + if num_samples > len(self.image_paths): + raise ValueError(f"Requested {num_samples} samples but only {len(self.image_paths)} images available") + + # Sample without replacement + sampled_paths = np.random.choice(self.image_paths, size=num_samples, replace=False) + + results = [] + for image_path in sampled_paths: + img_array = self._load_and_preprocess_image(image_path) + img_info = { + 'file_name': os.path.basename(image_path), + 'file_path': image_path, + 'height': img_array.shape[0], + 'width': img_array.shape[1] + } + results.append((img_info, img_array)) + + return results + + def load_image_by_path(self, image_path: str) -> np.ndarray: + """ + Load image from file path with preprocessing + + Args: + image_path: Path to image file + + Returns: + Preprocessed image array + """ + return self._load_and_preprocess_image(image_path) + + def _load_and_preprocess_image(self, image_path: str) -> np.ndarray: + """ + Load and preprocess image for flux model compatibility + + Args: + image_path: Path to image file + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + return preprocess_image_for_flux(image_path) + + def load_image_by_info(self, img_info: Dict) -> np.ndarray: + """ + Load image by image info dictionary + + Args: + img_info: Dictionary containing 'file_path' key + + Returns: + Preprocessed image array + """ + image_path = img_info.get('file_path') + if not image_path: + raise ValueError("Image info must contain 'file_path' key") + return self._load_and_preprocess_image(image_path) + + +def _get_coco_image_list( + data_loader: COCODataLoader, + categories: List[str], + max_images: Optional[int] = None, + max_instances_per_image: Optional[int] = None +) -> List[Dict[str, Any]]: + """ + Get image list for COCO dataset with optional filtering. + + Args: + data_loader: COCO data loader instance + categories: List of categories to process + max_images: Maximum number of images to process + max_instances_per_image: Maximum instances per image for filtering + + Returns: + List of image information dictionaries + """ + cat_ids = data_loader.get_category_ids(categories) + image_list = [] + image_ids_seen = set() + + # Count instances per image if filtering is requested + instance_counts = {} + if max_instances_per_image is not None: + print("Counting instances per image...") + from collections import defaultdict + instance_counts = defaultdict(int) + for ann in data_loader.coco_class.dataset['annotations']: + image_id = ann['image_id'] + instance_counts[image_id] += 1 + + for cat_id in cat_ids: + img_ids = data_loader.coco_class.getImgIds(catIds=[cat_id]) + for img_id in img_ids: + if img_id not in image_ids_seen: + # Filter by instance count if specified + if max_instances_per_image is not None: + if instance_counts[img_id] >= max_instances_per_image: + continue + + img_info = data_loader.coco_class.loadImgs([img_id])[0] + image_list.append(img_info) + image_ids_seen.add(img_id) + + print("number of images", len(image_list)) + return image_list + + +def _get_imagenet_image_list( + data_loader: ImageNetDataLoader, + categories: List[str], + max_images: Optional[int] = None +) -> List[Dict[str, Any]]: + """ + Get image list for ImageNet dataset. + + Args: + data_loader: ImageNet data loader instance + categories: List of categories to process + max_images: Maximum number of images to process + + Returns: + List of image information dictionaries + """ + # Determine target synsets + target_synsets = [] + for class_name in categories: + for synset, mapped_name in data_loader.class_mapping.items(): + if mapped_name.lower() == class_name.lower() and synset in data_loader.synsets: + target_synsets.append(synset) + + if not target_synsets: + target_synsets = data_loader.synsets + + image_list = [] + for synset in target_synsets: + image_paths = data_loader.get_images_by_synset(synset) + for img_path in image_paths: + img_info = { + 'id': hash(img_path) % 1000000, # Generate unique ID + 'file_name': os.path.basename(img_path), + 'file_path': img_path, + 'synset': synset, + 'class_name': data_loader.class_mapping.get(synset, synset) + } + image_list.append(img_info) + + if max_images and len(image_list) >= max_images: + break + if max_images and len(image_list) >= max_images: + break + + return image_list + + +def _get_custom_image_list( + data_loader: CustomDirectoryDataLoader, + categories: List[str], + max_images: Optional[int] = None, + logger: logging.Logger = None +) -> List[Dict[str, Any]]: + """ + Get image list for custom dataset. + + Args: + data_loader: Custom directory data loader instance + categories: List of categories (ignored for flat directory structure) + max_images: Maximum number of images to process + logger: Logger instance + + Returns: + List of image information dictionaries + """ + # Get all available image paths from the directory + all_image_paths = data_loader.get_all_image_paths() + if logger: + logger.info(f"Found {len(all_image_paths)} images in custom dataset directory") + + # Limit images if max_images is specified + if max_images and max_images < len(all_image_paths): + all_image_paths = all_image_paths[:max_images] + if logger: + logger.info(f"Limited to first {max_images} images") + + # Create image info list + image_list = [] + for img_path in all_image_paths: + img_info = { + 'id': hash(img_path) % 1000000, # Generate unique ID + 'file_name': os.path.basename(img_path), + 'file_path': img_path + } + image_list.append(img_info) + + return image_list + + +def _get_image_list( + dataset_type: str, + data_loader: Any, + categories: List[str], + max_images: Optional[int] = None, + max_instances_per_image: Optional[int] = None, + logger: logging.Logger = None +) -> List[Dict[str, Any]]: + """ + Get image list based on dataset type. + + Args: + dataset_type: Type of dataset + data_loader: Data loader instance + categories: List of categories to process + max_images: Maximum number of images to process + max_instances_per_image: Maximum number of instances per image + logger: Logger instance + + Returns: + List of image information dictionaries + """ + if dataset_type == "coco": + return _get_coco_image_list(data_loader, categories, max_images, max_instances_per_image) + elif dataset_type == "imagenet": + return _get_imagenet_image_list(data_loader, categories, max_images) + elif dataset_type == "custom": + return _get_custom_image_list(data_loader, categories, max_images, logger) + else: + raise ValueError(f"Unsupported dataset type: {dataset_type}") + + +def _initialize_data_loader(dataset_type: str, config: Dict[str, Any]) -> Any: + """ + Initialize the appropriate data loader based on dataset type. + + Args: + dataset_type: Type of dataset ('coco', 'imagenet', 'custom') + config: Configuration dictionary + + Returns: + Initialized data loader instance + """ + if dataset_type == "coco": + return COCODataLoader(config['dataset_path'], config['image_path']) + elif dataset_type == "imagenet": + return ImageNetDataLoader(config['dataset_path'], config['imagenet_split']) + elif dataset_type == "custom": + return CustomDirectoryDataLoader(config['dataset_path']) + else: + raise ValueError(f"Unsupported dataset type: {dataset_type}") \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/defect_rag.py b/ArtiAgent - DefectDiffu/src/pipeline/defect_rag.py new file mode 100644 index 0000000000000000000000000000000000000000..1d080b7445ffa37c5cbcf2d04d5d5c00bfe6aecc --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/defect_rag.py @@ -0,0 +1,181 @@ +# ============================================ +# DefectRAG โ€” DefectDiffu Edition +# ============================================ +# Retrieves in-context defect examples from ChromaDB and composes +# DefectDiffu text prompts (c_d) from retrieved knowledge. + +import chromadb +from sentence_transformers import SentenceTransformer +import json +from pathlib import Path +from typing import List, Dict, Optional + + +class DefectRAG: + """ + Retrieval-Augmented Generation for defect examples. + Now includes helper to build DefectDiffu defect prompts (c_d) from RAG results. + """ + + def __init__( + self, + db_path: str = "data/defect_db", + collection: str = "defect_patches", + model: str = "all-MiniLM-L6-v2" + ): + self.db_path = db_path + self.collection_name = collection + self._client = None + self._collection = None + self._encoder = None + + @property + def client(self): + if self._client is None: + self._client = chromadb.PersistentClient(path=self.db_path) + return self._client + + @property + def collection(self): + if self._collection is None: + self._collection = self.client.get_collection(self.collection_name) + return self._collection + + @property + def encoder(self): + if self._encoder is None: + self._encoder = SentenceTransformer('all-MiniLM-L6-v2') + return self._encoder + + def _build_where_filter( + self, + commercial_only: bool = True, + domain_filter: Optional[str] = None + ) -> Optional[Dict]: + conditions = [] + if commercial_only: + conditions.append({"commercial_ok": True}) + if domain_filter and domain_filter != "general": + conditions.append({"domain": domain_filter}) + if len(conditions) == 0: + return None + elif len(conditions) == 1: + return conditions[0] + else: + return {"$and": conditions} + + def retrieve( + self, + defect_plan, + k: int = 3, + domain_filter: Optional[str] = None, + commercial_only: bool = True + ) -> List[Dict]: + """Retrieve top-k matching defect examples.""" + query_parts = [ + getattr(defect_plan, 'artifact_type', defect_plan.defect_type), + defect_plan.description, + "on", + defect_plan.target_entity + ] + query = " ".join(query_parts) + query_emb = self.encoder.encode(query) + where_filter = self._build_where_filter(commercial_only, domain_filter) + + results = self.collection.query( + query_embeddings=[query_emb.tolist()], + n_results=k, + where=where_filter + ) + + examples = [] + for i, meta in enumerate(results['metadatas'][0]): + paths = json.loads(meta['paths']) if isinstance(meta.get('paths'), str) else meta.get('paths', {}) + examples.append({ + 'paths': paths, + 'caption': meta.get('caption', ''), + 'domain': meta.get('domain', 'unknown'), + 'license': meta.get('license', 'unknown'), + 'source': meta.get('source', 'unknown'), + 'defect_name': meta.get('defect_name', 'unknown'), + 'score': results['distances'][0][i] if results.get('distances') else None, + 'metadata': {k: v for k, v in meta.items() if k not in {'paths', 'caption', 'domain', 'license', 'source', 'defect_name'}} + }) + return examples + + def retrieve_by_text( + self, + text: str, + k: int = 3, + domain_filter: Optional[str] = None + ) -> List[Dict]: + """Direct text search (for debugging/testing).""" + query_emb = self.encoder.encode(text) + where_filter = self._build_where_filter(True, domain_filter) + results = self.collection.query( + query_embeddings=[query_emb.tolist()], + n_results=k, + where=where_filter + ) + examples = [] + for i, meta in enumerate(results['metadatas'][0]): + paths = json.loads(meta['paths']) if isinstance(meta.get('paths'), str) else meta.get('paths', {}) + examples.append({ + 'paths': paths, + 'caption': meta.get('caption', ''), + 'domain': meta.get('domain', 'unknown'), + 'score': results['distances'][0][i] if results.get('distances') else None + }) + return examples + + def compose_defect_prompt( + self, + base_description: str, + examples: List[Dict], + max_examples: int = 1 + ) -> str: + """ + Build a DefectDiffu defect prompt (c_d) by enriching the base description + with captions from retrieved RAG examples. + + Example output: + "A photo of a small transparent bubble trapped under glass, similar to + a spherical air pocket with dark meniscus ring" + """ + if not examples: + return f"A photo of {base_description}" + + captions = [ex.get('caption', '') for ex in examples[:max_examples] if ex.get('caption')] + if captions: + enriched = f"{base_description}, similar to {captions[0]}" + return f"A photo of {enriched}" + return f"A photo of {base_description}" + + def get_stats(self) -> Dict: + """Get DB statistics.""" + count = self.collection.count() + results = self.collection.get() + domains = {} + licenses = {} + sources = {} + for meta in results["metadatas"]: + domains[meta.get("domain", "unknown")] = domains.get(meta.get("domain"), 0) + 1 + licenses[meta.get("license", "unknown")] = licenses.get(meta.get("license"), 0) + 1 + sources[meta.get("source", "unknown")] = sources.get(meta.get("source"), 0) + 1 + return { + "total_entries": count, + "domains": domains, + "licenses": licenses, + "sources": sources + } + + +# Singleton instance +_rag_instance = None + +def get_rag(db_path="data/defect_db") -> DefectRAG: + """Get or create singleton RAG instance.""" + global _rag_instance + if _rag_instance is None: + _rag_instance = DefectRAG(db_path=db_path) + return _rag_instance diff --git a/ArtiAgent - DefectDiffu/src/pipeline/defectdiffu_generator.py b/ArtiAgent - DefectDiffu/src/pipeline/defectdiffu_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..5b034e7d8934fbdbeb53bf4d1e78ca4973530710 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/defectdiffu_generator.py @@ -0,0 +1,361 @@ +""" +DefectDiffu Generator Wrapper + +Replaces FLUX/RF-Solver-Edit with DefectDiffu's text-guided disentangled +architecture for manufacturing defect generation. + +DefectDiffu (ECCV 2024) uses: + - Three disentangled text prompts: c_p (product/bg), c_d (defect), c_f (fusion) + - Double-free strategy with perturbation scales w_p and w_d + - Automatic mask extraction from defect-block cross-attention maps + - DiT backbone + Stable Diffusion VAE + +INTEGRATION NOTE: + This file contains the INTERFACE. You must plug in the actual DefectDiffu + model forward pass from the official repo at the marked TODO sections. + Repo: https://github.com/FFDD-diffusion/DefectDiffu +""" + +import os +import torch +import torch.nn.functional as F +import numpy as np +from typing import Dict, Tuple, Optional, List +from dataclasses import dataclass +from PIL import Image +import warnings + +# === DefectDiffu actual imports (must be in PYTHONPATH) === +from diffusers.models import AutoencoderKL + +import sys +from pathlib import Path + +# Automatically locate and add the DefectDiffu engine directory to sys.path +CURRENT_DIR = Path(__file__).resolve().parent +DEFECTDIFFU_DIR = CURRENT_DIR.parent.parent / "engine" / "DefectDiffu" + +if DEFECTDIFFU_DIR.exists() and str(DEFECTDIFFU_DIR) not in sys.path: + sys.path.insert(0, str(DEFECTDIFFU_DIR)) +import clip.clip as clip + +from models_add_cross_concate import DiT +from diffusion import create_diffusion + + +# ========================================================================= +# Mask binarization helpers (copied from test.py) +# ========================================================================= + +def rgb_to_gray(tensor): + r, g, b = tensor[:, 0], tensor[:, 1], tensor[:, 2] + gray = 0.299 * r + 0.587 * g + 0.114 * b + return gray + + +def iterative_thresholding_batch(gray_tensor): + gray_np = gray_tensor.detach().cpu().numpy() + binarized = np.zeros_like(gray_np, dtype=np.uint8) + + for i in range(gray_np.shape[0]): + img = gray_np[i] + T = img.mean() + prev_T = -1 + + while abs(T - prev_T) > 1e-4: + prev_T = T + G1 = img[img >= T] + G2 = img[img < T] + m1 = G1.mean() if G1.size > 0 else 0 + m2 = G2.mean() if G2.size > 0 else 0 + T = (m1 + m2) / 2 + + binarized[i] = (img >= T).astype(np.uint8) + + return torch.from_numpy(binarized).to(gray_tensor.device) + + +def binarize_tensor_iterative(x): + gray = rgb_to_gray(x) + binary = iterative_thresholding_batch(gray) + return binary.unsqueeze(1) + +@dataclass +class DefectDiffuConfig: + """Configuration for DefectDiffu inference.""" + ckpt_path: str # Path to trained DefectDiffu checkpoint + vae_path: str # Path to SD VAE (stabilityai/sd-vae-ft-mse) + dit_model: str = "DiT-XL/2" # DiT variant (DiT-XL/2, DiT-L/2, etc.) + image_size: int = 512 # Must match training resolution + num_steps: int = 50 # DDPM/DDIM inference steps + cfg_scale: float = 1.0 # Classifier-free guidance (if used) + device: str = "cuda" + offload: bool = False # CPU offload for low-VRAM GPUs + seed: int = 42 + + +class DefectDiffuGenerator: + """ + Wrapper around DefectDiffu for the agentic pipeline. + + Unlike FLUX (which edits an existing image via inversion-injection), + DefectDiffu generates a NEW image from noise conditioned on three text + prompts. The input "clean image" is used only for planning/verification, + not as a pixel-level source for editing. + """ + + def __init__(self, config: DefectDiffuConfig): + self.config = config + self.device = torch.device(config.device) + self._models_loaded = False + + # Placeholders โ€” populated in _load_models() + self.dit = None + self.vae = None + self.text_encoder = None + self.tokenizer = None + self.scheduler = None + + self._load_models() + + # ------------------------------------------------------------------ + # TODO: Replace the methods below with actual DefectDiffu code + # ------------------------------------------------------------------ + + def _load_models(self): + """Load DiT, VAE, text encoder, and scheduler.""" + print(f"[DefectDiffu] Loading checkpoint: {self.config.ckpt_path}") + print(f"[DefectDiffu] VAE: {self.config.vae_path}") + + # 1. CLIP RN50 (must match training) + self.model_clip, _ = clip.load('RN50', self.device) + self.model_clip.eval() + + # 2. DiT architecture (must match train.py exactly) + latent_size = self.config.image_size // 8 + self.dit = DiT( + depth=28, + hidden_size=1152, + patch_size=2, + num_heads=16, + input_size=latent_size, + num_classes=1000 + ).to(self.device) + + print(f"[DefectDiffu] Loading DiT weights from: {self.config.ckpt_path}") + checkpoint = torch.load(self.config.ckpt_path, map_location=self.device) + if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: + self.dit.load_state_dict(checkpoint['model_state_dict']) + else: + self.dit.load_state_dict(checkpoint) + self.dit.eval() + + # 3. Stable Diffusion VAE + self.vae = AutoencoderKL.from_pretrained(self.config.vae_path).to(self.device) + self.vae.eval() + + # 4. Diffusion sampler (respacing = num_steps) + self.diffusion = create_diffusion(timestep_respacing=str(self.config.num_steps)) + + self._models_loaded = True + print("[DefectDiffu] All models loaded successfully.") + + def _encode_text(self, prompt: str) -> torch.Tensor: + """Encode a text prompt into CLIP RN50 text embeddings.""" + with torch.no_grad(): + tokens = clip.tokenize([prompt]).to(self.device) + emb = self.model_clip.encode_text(tokens) + emb = emb / emb.norm(dim=-1, keepdim=True) + emb = emb.float() + return emb + + def _extract_mask_from_attention( + self, + mask_latent: torch.Tensor + ) -> np.ndarray: + """ + Decode mask latent through VAE and binarize using iterative thresholding. + Matches test.py post-processing. + """ + with torch.no_grad(): + mask_decoded = self.vae.decode(mask_latent / 0.18215).sample # [1, 3, H, W] + + # Binarize with iterative thresholding (Otsu-like) + mask_binary = binarize_tensor_iterative(mask_decoded) # [1, 1, H, W] + mask_bool = mask_binary[0, 0].cpu().numpy() > 0 + + return mask_bool + + def _denoise_with_double_free( + self, + z: torch.Tensor, + emb_p: torch.Tensor, + emb_d: torch.Tensor, + emb_f: torch.Tensor, + emb_good: torch.Tensor, + emb_null_good: torch.Tensor, + w_d: float, + w_p: float + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Run DefectDiffu inference via p_sample_loop with dual-branch CFG. + Matches test.py exactly. + + Returns: + (img_latent, mask_latent) in VAE latent space, shape [1, 4, H, W] + """ + # Build paired conditioning: defect_class vs good_class + y_defect_class = [emb_d, emb_p, emb_f] + y_good_class = [emb_good, emb_p, emb_null_good] + y = [y_defect_class, y_good_class] + + # Duplicate latent for CFG (concatenated batch) + z_cfg = torch.cat([z, z], dim=0) + + model_kwargs = dict(y=y, cfg_scale=float(w_d)) + + with torch.no_grad(): + samples, cross = self.diffusion.p_sample_loop( + self.dit.forward_with_cfg_2, + z_cfg.shape, + z_cfg, + clip_denoised=False, + model_kwargs=model_kwargs, + progress=False, + device=self.device + ) + + # Unchunk: first half is the defect-conditioned output + img_latent, _ = samples.chunk(2, dim=0) + mask_latent, _ = cross.chunk(2, dim=0) + + return img_latent, mask_latent + + # ------------------------------------------------------------------ + # Public API โ€” used by the orchestrator + # ------------------------------------------------------------------ + + @torch.no_grad() + def generate( + self, + c_p: str, + c_d: str, + c_f: str, + w_d: float = 1.0, + w_p: float = 1.0, + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray]: + """ + Generate a synthetic defect image and its binary mask. + + Args: + c_p: Background/product prompt (e.g. "A photo of metal nut") + c_d: Defect prompt (e.g. "A photo of scratch") + c_f: Fusion prompt (e.g. "A photo of metal nut with scratch") + w_d: Defect strength perturbation scale (0.0 = no defect, 2.0 = severe) + w_p: Product consistency scale (usually 1.0, increase for stronger product fidelity) + seed: Random seed + + Returns: + (pil_image, binary_mask) where mask is bool array [H, W] + """ + if seed is None: + seed = self.config.seed + torch.manual_seed(seed) + np.random.seed(seed) + + print(f"[DefectDiffu] Generating: w_d={w_d}, w_p={w_p}") + print(f"[DefectDiffu] c_p: {c_p}") + print(f"[DefectDiffu] c_d: {c_d}") + print(f"[DefectDiffu] c_f: {c_f}") + + # 1. Parse product name from c_p for the null-good prompt + product_name = c_p.replace("A photo of ", "").strip() + + # 2. Encode all five text conditions (must match training format) + emb_d = self._encode_text(c_d) # "a photo of scratch" + emb_p = self._encode_text(c_p) # "a photo of vcsel" + emb_f = self._encode_text(c_f) # "a photo of scratch vcsel" + emb_good = self._encode_text("a photo of good") + emb_null_good = self._encode_text(f"a photo of good {product_name}") + + # 3. Initialize latent noise + latent_h = self.config.image_size // 8 + latent_w = self.config.image_size // 8 + z = torch.randn(1, 4, latent_h, latent_w, device=self.device) + + # 4. DefectDiffu double-free denoising + img_latent, mask_latent = self._denoise_with_double_free( + z, emb_p, emb_d, emb_f, emb_good, emb_null_good, w_d, w_p + ) + + # 5. Decode image latent โ†’ RGB + with torch.no_grad(): + img_tensor = self.vae.decode(img_latent / 0.18215).sample # [1, 3, H, W] + img_tensor = (img_tensor + 1) / 2 # [-1, 1] โ†’ [0, 1] + img_tensor = img_tensor.clamp(0, 1) + + img_np = (img_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) + pil_image = Image.fromarray(img_np) + + # 6. Extract defect mask from mask latent + binary_mask = self._extract_mask_from_attention(mask_latent) + + print(f"[DefectDiffu] Generation complete. Mask coverage: {binary_mask.mean():.3f}") + return pil_image, binary_mask + + @torch.no_grad() + def generate_from_plan( + self, + product_description: str, + defect_description: str, + severity: str = "medium", + w_d: Optional[float] = None, + w_p: float = 1.0, + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray, Dict]: + """ + Convenience wrapper that builds the three DefectDiffu prompts from + product/defect descriptions and maps severity to w_d. + """ + # Map severity to defect strength + severity_to_wd = {"low": 0.6, "minor": 0.6, + "medium": 1.0, "moderate": 1.0, + "high": 1.5, "severe": 1.5} + if w_d is None: + w_d = severity_to_wd.get(severity.lower(), 1.0) + + c_p = f"A photo of {product_description}" + c_d = f"A photo of {defect_description}" + c_f = f"A photo of {product_description} with {defect_description}" + + img, mask = self.generate(c_p, c_d, c_f, w_d=w_d, w_p=w_p, seed=seed) + + meta = { + "c_p": c_p, "c_d": c_d, "c_f": c_f, + "w_d": w_d, "w_p": w_p, "seed": seed or self.config.seed + } + return img, mask, meta + + def unload_models(self): + """Free GPU memory.""" + self.dit = None + self.vae = None + self.model_clip = None + self.diffusion = None + if self.device.type == "cuda": + torch.cuda.empty_cache() + print("[DefectDiffu] Models unloaded.") + + +def get_defectdiffu_generator( + ckpt_path: str, + vae_path: str, + device: str = "cuda", + **kwargs +) -> DefectDiffuGenerator: + """Factory function for easy instantiation.""" + config = DefectDiffuConfig(ckpt_path=ckpt_path, vae_path=vae_path, + device=device, **kwargs) + return DefectDiffuGenerator(config) + diff --git a/ArtiAgent - DefectDiffu/src/pipeline/domain_router.py b/ArtiAgent - DefectDiffu/src/pipeline/domain_router.py new file mode 100644 index 0000000000000000000000000000000000000000..65a2c95169e997d27bbcb33207150bad3c049010 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/domain_router.py @@ -0,0 +1,97 @@ +# src/pipeline/domain_router.py +# ============================================ +# Routes queries to correct collection based on product description +# ============================================ + +from typing import Optional + + +class DomainRouter: + """ + Routes defect queries to domain-specific DB collections or filters. + + Usage: + router = DomainRouter() + domain = router.route("VCSEL laser diode in TO-can package...") + # Returns: "glass/vcsel" + """ + + # Keyword โ†’ domain mapping + DOMAIN_MAP = { + # VCSEL / Optoelectronics + "vcsel": "glass/vcsel", + "laser diode": "glass/vcsel", + "to-can": "glass/vcsel", + "to can": "glass/vcsel", + "glass lens": "glass/vcsel", + "emission aperture": "glass/vcsel", + "bond pad": "glass/vcsel", + "optoelectronic": "glass/vcsel", + + # PCB + "pcb": "pcb", + "printed circuit": "pcb", + "solder joint": "pcb", + "smt": "pcb", + "surface mount": "pcb", + "trace": "pcb", + "pad": "pcb", + + # Metal / Steel + "steel": "metal/steel", + "sheet metal": "metal/steel", + "rolled": "metal/steel", + "metal surface": "metal/steel", + + # Semiconductor + "wafer": "semiconductor", + "die": "semiconductor", + "chip": "semiconductor", + "silicon": "semiconductor", + } + + def __init__(self): + self._cache = {} + + def route(self, product_desc: str) -> str: + """Determine domain from product description.""" + if not product_desc: + return "general" + + desc_lower = product_desc.lower() + + # Check cache + if desc_lower in self._cache: + return self._cache[desc_lower] + + # Match keywords + for keyword, domain in sorted(self.DOMAIN_MAP.items(), key=lambda x: -len(x[0])): + if keyword in desc_lower: + self._cache[desc_lower] = domain + return domain + + self._cache[desc_lower] = "general" + return "general" + + def get_collection_name(self, domain: str) -> str: + """Map domain to DB collection name.""" + # All domains share one collection with domain metadata + # Or use separate collections if needed + return "defect_patches" + + def get_domain_filter(self, product_desc: str) -> Optional[str]: + """Get domain filter for RAG query.""" + domain = self.route(product_desc) + if domain == "general": + return None + return domain + + +# Singleton +_router_instance = None + +def get_router() -> DomainRouter: + global _router_instance + if _router_instance is None: + _router_instance = DomainRouter() + return _router_instance \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/flux_generator.py b/ArtiAgent - DefectDiffu/src/pipeline/flux_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..ef665818e072980d7bc9afe363e49d5c238f2487 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/flux_generator.py @@ -0,0 +1,438 @@ +import torch +import argparse +from typing import Dict, Optional, Union, List +import numpy as np +from dataclasses import dataclass +import os +import re +import json +import time +from glob import iglob +from einops import rearrange +from PIL import Image + +# FLUX imports +import flux +from flux.sampling import denoise, denoise_first_order, denoise_fireflow, get_schedule, prepare, unpack +from flux.util import (configs, load_ae, load_clip, + load_flow_model, load_t5) + +@dataclass +class FluxConfig: + """Configuration for FLUX model""" + name: str = 'flux-dev' + guidance: float = 5.0 + num_steps: int = 25 + inject_step: int = 15 + pe_step: Union[int, Dict[str, int]] = 25 # Can be int or dict with artifact type keys + attn_mask_step: int = 0 + seed: int = 42 + masks: list = None + alpha: float = 0.0 + feature_path: str = 'feature' + percentage_of_steps: float = 1.0 + offload: bool = False + use_rf_solver: bool = False # Use denoise (RF solver) instead of denoise_first_order + + def __post_init__(self): + if self.masks is None: + self.masks = ['none', 'none', 'none', 'none'] + + # Validate pe_step configuration + if isinstance(self.pe_step, dict): + required_keys = {'addition', 'removal', 'distortion', 'fusion'} + provided_keys = set(self.pe_step.keys()) + if not required_keys.issubset(provided_keys): + missing_keys = required_keys - provided_keys + raise ValueError(f"pe_step dict missing required artifact types: {missing_keys}") + + # Validate all values are integers + for artifact_type, value in self.pe_step.items(): + if not isinstance(value, int): + raise ValueError(f"pe_step value for '{artifact_type}' must be an integer, got {type(value)}") + + def get_pe_step(self, artifact_type: str) -> int: + """ + Get pe_step value for specific artifact type + + Args: + artifact_type: Type of artifact ('addition', 'removal', 'distortion') + + Returns: + pe_step value for the artifact type + """ + if isinstance(self.pe_step, dict): + if artifact_type not in self.pe_step: + raise ValueError(f"Unknown artifact type '{artifact_type}'. Available types: {list(self.pe_step.keys())}") + return self.pe_step[artifact_type] + else: + return self.pe_step + + +class FluxGenerator: + """Handler for FLUX model operations and image generation""" + + def __init__(self, device: str = 'cuda', config: Optional[FluxConfig] = None): + """ + Initialize FLUX generator + + Args: + device: Device to run models on ('cuda' or 'cpu') + config: FLUX configuration object + """ + self.device = device + self.config = config or FluxConfig() + + # Model components + self.t5 = None + self.clip = None + self.model = None + self.ae = None + + self._models_loaded = False + self.load_models() + + def load_models(self): + """Load all FLUX model components""" + print("Loading FLUX models...") + + # Determine max_length based on model name + max_length = 256 if self.config.name == "flux-schnell" else 512 + + # Load model components + self.t5 = load_t5(self.device, max_length=max_length).to(dtype=torch.bfloat16) + self.clip = load_clip(self.device).to(dtype=torch.bfloat16) + self.offload = self.config.offload # store it so you know it's defined + self.model = load_flow_model(self.config.name, device=self.device) + if self.offload: + self.model.enable_sequential_cpu_offload() + self.ae = load_ae(self.config.name, device=self.device) + + self._models_loaded = True + print("FLUX models loaded successfully.") + + def create_default_flux_args(self) -> argparse.Namespace: + """ + Create default FLUX arguments based on current configuration + + Returns: + Default argparse.Namespace object with config values + """ + # Create parser and args + parser = argparse.ArgumentParser() + flux_args = parser.parse_args(args=[]) + + # Set FLUX configuration defaults + flux_args.name = self.config.name + flux_args.feature_path = self.config.feature_path + flux_args.guidance = self.config.guidance + flux_args.num_steps = self.config.num_steps + flux_args.inject_step = self.config.inject_step + flux_args.attn_mask_step = self.config.attn_mask_step + flux_args.pe_step = self.config.pe_step + flux_args.pe_step_addition = self.config.pe_step['addition'] + flux_args.pe_step_removal = self.config.pe_step['removal'] + flux_args.pe_step_distortion = self.config.pe_step['distortion'] + flux_args.pe_step_fusion = self.config.pe_step['fusion'] + flux_args.seed = self.config.seed + flux_args.masks = self.config.masks.copy() + flux_args.alpha = self.config.alpha + flux_args.percentage_of_steps = self.config.percentage_of_steps + flux_args.offload = self.config.offload + + # Initialize task-specific arguments to None + flux_args.source_prompt = None + flux_args.target_prompt = None + flux_args.artifact_type = None + flux_args.output_dir = None + flux_args.source_img = None + + # Initialize optional patch mapping information + flux_args.patch_mapping = None + flux_args.reference_patch_indices = None + flux_args.target_patch_indices = None + + return flux_args + + @torch.inference_mode() + def inject_artifacts(self, + source_prompt: str, + target_prompt: str, + artifact_data: Dict, + source_img: Union[np.ndarray, str], + output_dir: str = None, + pe_step_addition: Optional[int] = None, + pe_step_removal: Optional[int] = None, + pe_step_distortion: Optional[int] = None, + pe_step_fusion: Optional[int] = None, + inject_step: Optional[int] = None, + num_steps: Optional[int] = None, + use_fireflow: bool = False, + reference_images: Optional[List] = None + ): + """ + Sample the flux model with artifact injection supporting arbitrary shapes. + NEW: reference_images โ€” list of example defect images from RAG retrieval + + Args: + source_prompt: Source image prompt/caption + target_prompt: Target prompt for generation + artifact_type: Type of artifact ('addition', 'removal', 'distortion') + source_img: Source image array or path + output_dir: Output directory for generated images + reference_patch_indices: List of reference patch indices + target_patch_indices: List of target patch indices + """ + torch.set_grad_enabled(False) + + # Create default flux args and update with passed parameters + flux_args = self.create_default_flux_args() + + # Update with required parameters + flux_args.source_prompt = source_prompt + flux_args.target_prompt = target_prompt + flux_args.artifact_data = artifact_data + flux_args.source_img = source_img + flux_args.output_dir = output_dir + flux_args.inject_step = inject_step if inject_step is not None else self.config.inject_step + torch_device = torch.device(self.device) + if num_steps is not None: + flux_args.num_steps = num_steps + + init_image = None + init_image = self.load_image(flux_args.source_img) + + shape = init_image.shape + + new_h = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16 + new_w = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16 + + init_image = init_image[:new_h, :new_w, :] + + width, height = init_image.shape[0], init_image.shape[1] + init_image = self.encode(init_image, torch_device, self.ae) + + rng = torch.Generator(device="cpu").manual_seed(flux_args.seed) + + if flux_args.seed is None: + flux_args.seed = rng.seed() + print(f"Generating with seed {flux_args.seed}:\n{flux_args.source_prompt}") + t0 = time.perf_counter() + + flux_args.seed = None + if flux_args.offload: + self.ae = self.ae.cpu() + torch.cuda.empty_cache() + self.t5, self.clip = self.t5.to(torch_device), self.clip.to(torch_device) + + info = {} + info['feature_path'] = flux_args.feature_path + info['feature'] = {} + info['inject_step'] = flux_args.inject_step + info['attn_mask_step'] = flux_args.attn_mask_step + info['alpha'] = flux_args.alpha + if pe_step_distortion is not None: + info['pe_step_distortion'] = pe_step_distortion + else: + info['pe_step_distortion'] = flux_args.pe_step_distortion + if pe_step_removal is not None: + info['pe_step_removal'] = pe_step_removal + else: + info['pe_step_removal'] = flux_args.pe_step_removal + if pe_step_addition is not None: + info['pe_step_addition'] = pe_step_addition + else: + info['pe_step_addition'] = flux_args.pe_step_addition + if pe_step_fusion is not None: + info['pe_step_fusion'] = pe_step_fusion + else: + info['pe_step_fusion'] = flux_args.pe_step_fusion + info['artifact_data'] = flux_args.artifact_data + info['guidance'] = flux_args.guidance + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # NEW: Encode RAG reference images and inject into info + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + if reference_images is not None and len(reference_images) > 0: + print(f"[FLUX] Encoding {len(reference_images)} reference image(s) from RAG...") + + # โ”€โ”€ FIX: Ensure AE is on GPU for encoding โ”€โ”€ + self.ae = self.ae.to(torch_device) + + ref_latents = [] + for i, ref_img in enumerate(reference_images): + # Load if path, convert if PIL + if isinstance(ref_img, str): + ref_img = self.load_image(ref_img) + elif isinstance(ref_img, Image.Image): + ref_img = np.array(ref_img.convert('RGB')) + + # Match dimensions to source image + # Resize to match source dimensions for consistent latent encoding + ref_pil = Image.fromarray(ref_img).resize((new_w, new_h), Image.LANCZOS) + ref_img = np.array(ref_pil) + # Encode through VAE (same as source image) + ref_latent = self.encode(ref_img, torch_device, self.ae) + ref_latents.append(ref_latent) + print(f"[FLUX] Reference {i+1} encoded: {ref_latent.shape}") + + info['reference_latents'] = ref_latents + info['num_reference_images'] = len(ref_latents) + + # โ”€โ”€ MOVE AE BACK TO CPU TO SAVE MEMORY DURING DENOISING โ”€โ”€ + if flux_args.offload: + self.ae = self.ae.cpu() + torch.cuda.empty_cache() + else: + info['reference_latents'] = [] + info['num_reference_images'] = 0 + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + if not os.path.exists(flux_args.feature_path): + os.mkdir(flux_args.feature_path) + + # Prepare inputs with shape-aware approach + inp, (patch_h, patch_w) = prepare( + self.t5, self.clip, init_image, + prompt=flux_args.source_prompt, + info=info + ) + + inp_target, _ = prepare( + self.t5, self.clip, init_image, + prompt=flux_args.target_prompt, + info=info + ) + + timesteps = get_schedule(flux_args.num_steps, inp["img"].shape[1], shift=(flux_args.name != "flux-schnell")) + + info['patch_h'] = patch_h + info['patch_w'] = patch_w + + L = inp['img'].shape[1] + inp['txt'].shape[1] + + # Choose denoising function based on configuration + # RF solver (denoise) is more accurate but slower than first-order denoising + denoise_func = denoise_fireflow if use_fireflow else denoise_first_order + # denoise_func = denoise_fireflow + # denoise_func = denoise_fireflow + # inversion initial noise + + # 1. Convert any Float32 tensor inputs to bfloat16 to avoid bitsandbytes float32->float16 warnings + inp = { + k: v.to(dtype=torch.bfloat16) if isinstance(v, torch.Tensor) and v.dtype == torch.float32 else v + for k, v in inp.items() + } + + # BEFORE: + # z, info = denoise_func(self.model, **inp, timesteps=timesteps, guidance=1, inverse=True, info=info, percentage_of_steps=flux_args.percentage_of_steps) + # AFTER: + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + z, info = denoise_func( + self.model, + **inp, + timesteps=timesteps, + guidance=1, + inverse=True, + info=info, + percentage_of_steps=flux_args.percentage_of_steps + ) + + inp_target["img"] = z + + timesteps = get_schedule(flux_args.num_steps, inp_target["img"].shape[1], shift=(flux_args.name != "flux-schnell")) + + # denoise initial noise + x, _ = denoise_func(self.model, **inp_target, timesteps=timesteps, guidance=info['guidance'], inverse=False, info=info, percentage_of_steps=flux_args.percentage_of_steps) + + # Bring AE back to GPU for decoding (it was offloaded to CPU earlier) + if flux_args.offload: + self.t5, self.clip = self.t5.cpu(), self.clip.cpu() + torch.cuda.empty_cache() + self.ae = self.ae.to(torch_device) + + # decode latents to pixel space + + ####################################### + #### TODO: allow batch computation #### + ####################################### + + x = unpack(x.float(), width, height) + + if output_dir is not None: + output_name = os.path.join(output_dir, "img.jpg") + if not os.path.exists(output_dir): + os.makedirs(output_dir) + idx = 0 + else: + fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)] + if len(fns) > 0: + idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1 + else: + idx = 0 + + with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16): + x = self.ae.decode(x.float()) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + t1 = time.perf_counter() + + print(f"Done in {t1 - t0:.1f}s.") + + # bring into PIL format and save + x = x.clamp(-1, 1) + x = rearrange(x[0], "c h w -> h w c") + img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy()) + + return img + + def load_image(self, source): + """Load image from various sources (numpy array, PIL Image, or file path)""" + if isinstance(source, np.ndarray): + # Already a NumPy array + return source + elif isinstance(source, Image.Image): + # Already a PIL Image + return np.array(source.convert('RGB')) + elif isinstance(source, str): + if os.path.isfile(source): + # It's a file path to an image + return np.array(Image.open(source).convert('RGB')) + else: + raise ValueError(f"Provided string is not a valid file: {source}") + else: + raise TypeError(f"Unsupported input type: {type(source)}") + + @torch.inference_mode() + def encode(self, init_image, torch_device, ae): + """Encode image to latent space""" + init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1 + init_image = init_image.unsqueeze(0) + init_image = init_image.to(torch_device) + init_image = ae.encode(init_image).to(torch.bfloat16) + return init_image + + def update_config(self, **kwargs): + """Update FLUX configuration parameters""" + for key, value in kwargs.items(): + if hasattr(self.config, key): + setattr(self.config, key, value) + else: + print(f"Warning: Unknown config parameter '{key}'") + + def unload_models(self): + """Unload models to free memory""" + self.t5 = None + self.clip = None + self.model = None + self.ae = None + self._models_loaded = False + + # Clear GPU cache if using CUDA + if self.device == 'cuda' and torch.cuda.is_available(): + torch.cuda.empty_cache() + + def __del__(self): + """Cleanup when object is destroyed""" + self.unload_models() \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/gsam_detector.py b/ArtiAgent - DefectDiffu/src/pipeline/gsam_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..4f5c99bec7487af2d084d56b8c0113128297f9ca --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/gsam_detector.py @@ -0,0 +1,379 @@ +import sys +import os +import multiprocessing as mp +import numpy as np +from typing import List, Dict, Tuple, Optional, Union +import torch +from PIL import Image +import torchvision +import supervision as sv +from flux.artifacts_util import mask_to_patch_coords +from logging import Logger as logger +from pydantic import BaseModel, RootModel +import json +import re + +# Add GroundingDINO and SAM to path +sys.path.append(os.path.join(os.getcwd(), 'GroundingDINO')) +sys.path.append(os.path.join(os.getcwd(), 'segment_anything')) +sys.path.append(os.path.join(os.getcwd(), 'pipeline')) + +# Grounding DINO +from groundingdino.util.inference import Model + +# S +from segment_anything import ( + sam_model_registry, + sam_hq_model_registry, + SamPredictor +) + +def robust_json_parse(raw_text: str): + """Try to parse VLM JSON output, with automatic repair.""" + # 1. Try direct parse + try: + return json.loads(raw_text) + except json.JSONDecodeError: + pass + + # 2. Extract JSON block from markdown fences + match = re.search(r'```json\s*(.*?)\s*```', raw_text, re.DOTALL) + if match: + try: + return json.loads(match.group(1)) + except json.JSONDecodeError: + pass + + # 3. Find the outermost {...} or [...] + match = re.search(r'(\{.*\}|\[.*\])', raw_text, re.DOTALL) + if match: + candidate = match.group(1) + # Fix common VLM JSON mistakes: + # - trailing commas before ] or } + candidate = re.sub(r',\s*([}\]])', r'\1', candidate) + # - extra closing brackets + while candidate.count('[') < candidate.count(']'): + candidate = candidate[:-1] # remove trailing ] + while candidate.count('{') < candidate.count('}'): + candidate = candidate[:-1] + # - missing closing brace + if candidate.count('{') > candidate.count('}'): + candidate += '}' + try: + return robust_json_parse(candidate) + except json.JSONDecodeError: + pass + + # 4. Fallback: regex extract bboxes directly + bboxes = re.findall(r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]', raw_text) + if bboxes: + return {"bboxes": [[int(x) for x in box] for box in bboxes]} + + raise ValueError(f"Could not parse JSON from: {raw_text[:200]}") + +class GSAMDetector: + """Handler for Grounded SAM part detection model""" + + # Constants for better code maintainability + DEFAULT_CONTAINMENT_THRESHOLD = 0.9 + DEFAULT_MIN_AREA_RATIO = 0.005 + DEFAULT_MAX_AREA_RATIO = 0.5 + + def __init__(self, + grounding_config_file: Optional[str] = None, + grounding_checkpoint: Optional[str] = None, + sam_version: str = "vit_h", + sam_checkpoint: Optional[str] = None, + sam_hq_checkpoint: Optional[str] = None, + use_sam_hq: bool = False, + box_threshold: float = 0.3, + text_threshold: float = 0.25, + nms_threshold: float = 0.5, + bert_base_uncased_path: Optional[str] = None, + device: str = "cuda", + openai_client: Optional[any] = None + ): + """ + Initialize GSAM detector + + Args: + grounding_config_file: Path to GroundingDINO config file + grounding_checkpoint: Path to GroundingDINO checkpoint + sam_version: SAM model version (vit_b, vit_l, vit_h) + sam_checkpoint: Path to SAM checkpoint + sam_hq_checkpoint: Path to SAM-HQ checkpoint + use_sam_hq: Whether to use SAM-HQ + box_threshold: Box threshold for detection + text_threshold: Text threshold for detection + nms_threshold: NMS threshold for detection + bert_base_uncased_path: Path to BERT model + device: Device to use (cuda/cpu) + openai_client: OpenAI client for vocabulary generation + """ + self.gsam_path = os.getcwd() + + # Set default paths if not provided + if grounding_config_file is None: + grounding_config_file = os.path.join(self.gsam_path, "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py") + if grounding_checkpoint is None: + grounding_checkpoint = os.path.join(self.gsam_path, "weight/groundingdino_swint_ogc.pth") + if sam_checkpoint is None and not use_sam_hq: + sam_checkpoint = os.path.join(self.gsam_path, "weight/sam_vit_h_4b8939.pth") + if sam_hq_checkpoint is None and use_sam_hq: + sam_hq_checkpoint = os.path.join(self.gsam_path, "weight/sam_hq_vit_h.pth") + + self.grounding_config_file = grounding_config_file + self.grounding_checkpoint = grounding_checkpoint + self.sam_version = sam_version + self.sam_checkpoint = sam_checkpoint + self.sam_hq_checkpoint = sam_hq_checkpoint + self.nms_threshold = nms_threshold + self.box_threshold = box_threshold + self.text_threshold = text_threshold + self.device = device + self.openai_client = openai_client + + # Model components + self.grounding_model = Model(model_config_path=self.grounding_config_file, model_checkpoint_path=self.grounding_checkpoint) + if use_sam_hq: + self.sam_predictor = SamPredictor(sam_hq_model_registry[self.sam_version](checkpoint=self.sam_hq_checkpoint).to(self.device)) + else: + self.sam_predictor = SamPredictor(sam_model_registry[self.sam_version](checkpoint=self.sam_checkpoint).to(self.device)) + + # Set multiprocessing start method + mp.set_start_method('spawn', force=True) + + + # Prompting SAM with detected boxes (same as original) + def segment(self, image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: + self.sam_predictor.set_image(image) + result_masks = [] + for box in xyxy: + masks, scores, logits = self.sam_predictor.predict( + box=box, + multimask_output=True + ) + index = np.argmax(scores) + result_masks.append(masks[index]) + return np.array(result_masks) + + def detect_parts(self, image, entities, subentities, entity_subentity_mapping, + min_area_ratio=0.005, max_area_ratio=0.5, openai_client=None): + """ + Detect parts using VLM for bboxes + SAM for masks. + Replaces GroundingDINO with VLM-guided detection. + """ + import cv2 + h, w = image.shape[:2] + total_area = h * w + + predictions = [] + entity_predictions = [] + + # Use VLM to get bboxes for each entity + from pipeline.prompts import get_entity_bboxes + + all_bboxes = {} # entity -> list of bboxes + + for entity in entities: + bboxes = get_entity_bboxes(openai_client, image, entity) + if bboxes: + all_bboxes[entity] = bboxes + print(f"VLM found {len(bboxes)} instances of '{entity}'") + + # Use SAM to get masks from bboxes + self.sam_predictor.set_image(image) # <-- ADD THIS LINE + for entity, bboxes in all_bboxes.items(): + for bbox in bboxes: + x1, y1, x2, y2 = map(int, bbox) + + # Validate bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(w, x2), min(h, y2) + if x2 <= x1 or y2 <= y1: + continue + + bbox_area = (x2 - x1) * (y2 - y1) + area_ratio = bbox_area / total_area + + if not (min_area_ratio <= area_ratio <= max_area_ratio): + print(f"Discarded '{entity}' bbox - area ratio {area_ratio:.4f} out of range") + continue + + # SAM mask from bbox + input_box = np.array([x1, y1, x2, y2]) + masks, scores, _ = self.sam_predictor.predict( + point_coords=None, + point_labels=None, + box=input_box[None, :], + multimask_output=False + ) + + if masks is None or len(masks) == 0: + continue + + mask = masks[0] if len(masks.shape) == 3 else masks + mask_binary = (mask > 0).astype(np.uint8) + + # Find contours for precise bbox + contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + continue + + cnt = max(contours, key=cv2.contourArea) + ex, ey, ew, eh = cv2.boundingRect(cnt) + + pred = { + 'entity': entity, + 'subentity': entity, # For now, entity = subentity + 'pred_box': torch.tensor([ex, ey, ex + ew, ey + eh]).float(), + 'bbox': [ex, ey, ex + ew, ey + eh], # alias for downstream + 'pred_class': torch.tensor(0).long(), # dummy class + 'pred_mask': torch.from_numpy(mask_binary).bool(), # numpy โ†’ torch tensor # <-- change 'mask' to 'pred_mask' + 'mask': mask_binary, # alias for downstream + 'score': torch.tensor(float(scores[0]) if len(scores) > 0 else 1.0).float(), + 'area_ratio': area_ratio + } + predictions.append(pred) + entity_predictions.append(pred) + + print(f"VLM+SAM: Found {len(predictions)} valid detections") + + # Visualization (optional) + visualized_output = None + try: + import supervision as sv + if predictions: + img_viz = image.copy() if hasattr(image, 'copy') else np.array(image) + # Simple bbox visualization + for p in predictions: + x1, y1, x2, y2 = map(int, p['pred_box'].tolist()) + cv2.rectangle(img_viz, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(img_viz, p['entity'], (x1, y1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + visualized_output = img_viz + except Exception as e: + print(f"Visualization error: {e}") + + return predictions, entity_predictions, visualized_output + + def detect_entities(self, image: np.ndarray, entities: List[str], + min_area_ratio: float = 0.01, max_area_ratio: float = 1.0) -> Tuple[List[Dict], any]: + """ + Run entity detection on image (entities only, no subentities) + + Args: + image: Input image as numpy array (RGB format) + entities: List of entity names + min_area_ratio: Minimum area ratio for filtering + max_area_ratio: Maximum area ratio for filtering + + Returns: + Tuple of (entity_predictions, visualized_output): + - entity_predictions: List of dictionaries, each containing entity detection with keys: + 'pred_box', 'pred_class', 'score', 'pred_mask', 'entity_name' + - visualized_output: PIL Image with annotations + """ + # Store current image size for area calculations + self.current_image_size = image.shape[:2] + + # Get grounding output + detections, phrases = self.grounding_model.predict_with_caption( + image=image, + caption=", ".join(entities), + box_threshold=self.box_threshold, + text_threshold=self.text_threshold + ) + + # Generate class_id from phrases since predict_with_caption doesn't include it + detections.class_id = Model.phrases2classes(phrases=phrases, classes=entities) + + # NMS post process + print(f"Before NMS: {len(detections.xyxy)} boxes") + nms_idx = torchvision.ops.nms( + torch.from_numpy(detections.xyxy), + torch.from_numpy(detections.confidence), + self.nms_threshold + ).numpy().tolist() + + detections.xyxy = detections.xyxy[nms_idx] + detections.confidence = detections.confidence[nms_idx] + detections.class_id = detections.class_id[nms_idx] + # Also filter phrases to match the filtered detections + phrases = [phrases[i] for i in nms_idx] + + detections.mask = self.segment( + image=image, + xyxy=detections.xyxy, + ) + + print(f"Found {len(detections.class_id)} entity detections") + + # Filter entities by area ratio + filtered_entities = [] + image_area = self.current_image_size[0] * self.current_image_size[1] + + for i in range(len(detections.class_id)): + entity_mask = torch.from_numpy(detections.mask[i]) + entity_class = detections.class_id[i] + entity_name = entities[entity_class] + area_ratio = torch.sum(entity_mask > 0) / image_area + + if min_area_ratio <= area_ratio <= max_area_ratio: + filtered_entities.append(i) + print(f"Kept entity '{entity_name}' (class {entity_class}) with area ratio {area_ratio:.4f}") + else: + print(f"Discarded entity '{entity_name}' (class {entity_class}) - area ratio {area_ratio:.4f} outside range [{min_area_ratio}, {max_area_ratio}]") + + if len(filtered_entities) == 0: + raise ValueError("No entities detected after filtering") + + # Filter detections to keep only valid entities + filtered_xyxy = detections.xyxy[filtered_entities] + filtered_confidence = detections.confidence[filtered_entities] + filtered_class_id = detections.class_id[filtered_entities] + filtered_mask = detections.mask[filtered_entities] + + # Create filtered detections object for annotation + filtered_detections = sv.Detections( + xyxy=filtered_xyxy, + confidence=filtered_confidence, + class_id=filtered_class_id, + mask=filtered_mask + ) + + # Annotate image with filtered detections + box_annotator = sv.BoundingBoxAnnotator() + mask_annotator = sv.MaskAnnotator() + label_annotator = sv.LabelAnnotator() + + labels = [ + f"{entities[class_id]} {confidence:0.2f}" + for class_id, confidence in zip(filtered_class_id, filtered_confidence)] + + annotated_image = mask_annotator.annotate(scene=image.copy(), detections=filtered_detections) + annotated_image = box_annotator.annotate(scene=annotated_image, detections=filtered_detections) + annotated_image = label_annotator.annotate(scene=annotated_image, detections=filtered_detections, labels=labels) + # Convert annotated image to PIL Image + annotated_image = Image.fromarray(annotated_image) + + # Create entity predictions as list of dictionaries + entity_predictions = [] + for i, entity_idx in enumerate(filtered_entities): + entity_pred_instance = { + 'pred_box': torch.from_numpy(detections.xyxy[entity_idx]).float(), + 'pred_class': torch.tensor(detections.class_id[entity_idx]).long(), + 'score': torch.tensor(detections.confidence[entity_idx]).float(), + 'pred_mask': torch.from_numpy(detections.mask[entity_idx]).bool(), + 'entity': entities[detections.class_id[entity_idx]], + } + entity_predictions.append(entity_pred_instance) + + print(f"Returning {len(entity_predictions)} entity detections") + return entity_predictions, annotated_image + + def cleanup(self): + """Clean up model resources""" + self.grounding_model = None + self.sam_predictor = None + self.current_vocabulary = [] \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/instance_processor.py b/ArtiAgent - DefectDiffu/src/pipeline/instance_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..e24d90df3f1f8a68308e4dc2835e2c6a8b263ff0 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/instance_processor.py @@ -0,0 +1,101 @@ +""" +Instance Processor โ€” DefectDiffu Edition + +Stripped-down version: DefectDiffu does NOT use 16x16 patch mappings, +so all patch-based artifact logic has been removed. + +Retained utilities: + - bbox โ†” mask helpers (for verification cropping) + - IoU calculation + - Visualization helpers +""" + +import random +import torch +import numpy as np +from PIL import Image +from typing import List, Dict, Tuple, Optional, Union +import matplotlib.pyplot as plt +import matplotlib.patches as patches + + +class InstanceProcessor: + """Utility class for detection post-processing and mask operations.""" + + @staticmethod + def calculate_iou(box1: Union[List, np.ndarray], box2: Union[List, np.ndarray]) -> float: + x1 = max(box1[0], box2[0]) + y1 = max(box1[1], box2[1]) + x2 = min(box1[2], box2[2]) + y2 = min(box1[3], box2[3]) + if x2 <= x1 or y2 <= y1: + return 0.0 + intersection = (x2 - x1) * (y2 - y1) + area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) + area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) + union = area1 + area2 - intersection + return intersection / union if union > 0 else 0.0 + + @staticmethod + def mask_from_bbox(bbox: Tuple[int, int, int, int], img_shape: Tuple[int, ...]) -> np.ndarray: + """Create a binary mask from a bounding box.""" + h, w = img_shape[:2] + mask = np.zeros((h, w), dtype=np.uint8) + x1, y1, x2, y2 = bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(w, x2), min(h, y2) + if x2 > x1 and y2 > y1: + mask[y1:y2, x1:x2] = 1 + return mask + + @staticmethod + def get_bbox_from_mask(mask: np.ndarray, margin: int = 0) -> Tuple[int, int, int, int]: + """Compute tight bounding box from binary mask, with optional margin.""" + ys, xs = np.where(mask > 0) + if len(ys) == 0: + return (0, 0, 0, 0) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + h, w = mask.shape + x1 = max(0, x1 - margin) + y1 = max(0, y1 - margin) + x2 = min(w, x2 + margin) + y2 = min(h, y2 + margin) + return (x1, y1, x2, y2) + + @staticmethod + def visualize_generation_result( + original_image: np.ndarray, + generated_image: np.ndarray, + defect_mask: np.ndarray, + output_path: str, + title: str = "DefectDiffu Generation Result" + ): + """Create a 3-panel visualization: original, generated, mask overlay.""" + fig, axes = plt.subplots(1, 3, figsize=(18, 6)) + + axes[0].imshow(original_image) + axes[0].set_title("Original (Planning Reference)") + axes[0].axis("off") + + axes[1].imshow(generated_image) + axes[1].set_title("Generated Defect Image") + axes[1].axis("off") + + axes[2].imshow(generated_image) + axes[2].imshow(defect_mask, alpha=0.5, cmap="Reds") + axes[2].set_title("Defect Mask Overlay") + axes[2].axis("off") + + fig.suptitle(title, fontsize=14) + plt.tight_layout() + plt.savefig(output_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"[Viz] Saved result visualization to {output_path}") + + @staticmethod + def resize_to_square(image: np.ndarray, size: int = 512) -> np.ndarray: + """Resize image to square (DefectDiffu expects 512x512).""" + pil_img = Image.fromarray(image) if isinstance(image, np.ndarray) else image + pil_img = pil_img.resize((size, size), Image.LANCZOS) + return np.array(pil_img) diff --git a/ArtiAgent - DefectDiffu/src/pipeline/local_vlm_client.py b/ArtiAgent - DefectDiffu/src/pipeline/local_vlm_client.py new file mode 100644 index 0000000000000000000000000000000000000000..e20c4aecc9a6432b98b31c84e142e60d420ed079 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/local_vlm_client.py @@ -0,0 +1,192 @@ +import base64 +import io +import json +import requests +from typing import List, Optional, Union +from PIL import Image +import numpy as np +import re + + +class Usage: + def __init__(self, input_tokens=0, output_tokens=0): + self.input_tokens = input_tokens + self.output_tokens = output_tokens + + +class ChatCompletion: + def __init__(self, output_parsed, usage=None): + self.output_parsed = output_parsed + self.usage = usage or Usage() + + +class LocalVLMClient: + def __init__(self, base_url="http://localhost:11434", model="gemma3:12b"): + self.base_url = base_url + self.model = model + self.api_url = f"{base_url}/api/chat" + # Magic alias: Allows client.responses.parse(...) calls in prompts.py to work directly + self.responses = self + + def _strip_markdown_fences(self, text: str) -> str: + """Strip ```json ... ``` markdown code fences from VLM output.""" + cleaned = text.strip() + cleaned = re.sub(r'^```(?:json)?\s*', '', cleaned, flags=re.IGNORECASE) + cleaned = re.sub(r'\s*```\s*$', '', cleaned) + return cleaned.strip() + + def parse(self, model: str = None, input: list = None, temperature: float = 0.2, text_format=None): + """Alias for responses_parse to support client.responses.parse(...)""" + return self.responses_parse(model=model, input=input, temperature=temperature, text_format=text_format) + + def _pil_to_base64(self, img: Image.Image) -> str: + """Helper to convert a PIL Image into a base64 string for Ollama.""" + buffered = io.BytesIO() + img.convert("RGB").save(buffered, format="JPEG") + return base64.b64encode(buffered.getvalue()).decode("utf-8") + + def _prepare_image_base64(self, image_input) -> str: + """Converts any image input type into a pure base64 string.""" + if isinstance(image_input, np.ndarray): + img = Image.fromarray(image_input) + return self._pil_to_base64(img) + + elif isinstance(image_input, Image.Image): + return self._pil_to_base64(image_input) + + elif isinstance(image_input, str): + if image_input.startswith("data:image"): + return image_input.split(",")[1] + elif image_input.startswith("http"): + import urllib.request + with urllib.request.urlopen(image_input) as response: + img = Image.open(io.BytesIO(response.read())) + return self._pil_to_base64(img) + else: + # Local file path + with Image.open(image_input) as img: + return self._pil_to_base64(img) + else: + raise ValueError(f"Unsupported image input type: {type(image_input)}") + + def _call_ollama_chat(self, messages: list, temperature: float = 0.2, format_json: bool = False): + """Calls Ollama Chat API with system/user messages and images.""" + payload = { + "model": self.model, + "messages": messages, + "stream": False, + "options": { + "temperature": temperature + } + } + + if format_json: + payload["format"] = "json" + + try: + response = requests.post(self.api_url, json=payload, timeout=600) + response.raise_for_status() + result = response.json() + return result["message"]["content"] + except Exception as e: + print(f"Ollama API Error: {e}") + raise + + def responses_parse(self, model: str = None, input: list = None, temperature: float = 0.2, text_format=None): + # Override gpt-4o / external model strings with local Gemma model + active_model = self.model + formatted_messages = [] + + for msg in input: + role = msg["role"] + content = msg["content"] + + if isinstance(content, str): + formatted_messages.append({"role": role, "content": content}) + elif isinstance(content, list): + text_parts = [] + b64_images = [] + + for item in content: + if item.get("type") == "input_text" or "text" in item: + text_parts.append(item.get("text", "")) + elif item.get("type") == "input_image" or "image_url" in item: + img_src = item.get("image_url", item.get("image")) + b64_images.append(self._prepare_image_base64(img_src)) + + msg_obj = { + "role": role, + "content": "\n".join(text_parts) + } + if b64_images: + msg_obj["images"] = b64_images + + formatted_messages.append(msg_obj) + + # FIX 1: Provide explicit JSON structure examples instead of schema definitions + if text_format: + instruction = ( + "\n\nIMPORTANT: Do NOT output the schema structure or field types. " + "Output ONLY a populated JSON object like this:\n" + "{\n" + ' "has_artifact": true,\n' + ' "explanation": "Visual description of what was detected",\n' + ' "label": "artifact name"\n' + "}" + ) + if formatted_messages and formatted_messages[-1]["role"] == "user": + formatted_messages[-1]["content"] += instruction + else: + formatted_messages.append({"role": "user", "content": instruction}) + + payload = { + "model": active_model, + "messages": formatted_messages, + "stream": False, + "options": {"temperature": temperature} + } + + # Make the request to Ollama + try: + response = requests.post(self.api_url, json=payload, timeout=600) + response.raise_for_status() + response_text = response.json()["message"]["content"] + except Exception as e: + print(f"Ollama API error: {e}") + raise + + # Determine the parsed output based on whether a Pydantic text_format was passed + if text_format: + try: + parsed_json = json.loads(self._strip_markdown_fences(response_text)) + parsed_output = text_format(**parsed_json) + except Exception as e: + print(f"Failed to parse JSON response: {e}\nRaw output: {response_text}") + raise + else: + # Fallback for plain text standard responses + parsed_output = type("Output", (), {"explanation": response_text})() + + # Return a single standardized response object + return type("Response", (), { + "output_parsed": parsed_output, + "usage": Usage() + })() + + +# Quick Test Example +if __name__ == "__main__": + client = LocalVLMClient(model="gemma3:12b") + + test_input = [ + {"role": "system", "content": "You are a visual AI inspector."}, + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Analyze if this image is distorted."} + ] + } + ] + + res = client.responses_parse(input=test_input) + print("Result:", res.output_parsed.explanation) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/pipeline/prompts.py b/ArtiAgent - DefectDiffu/src/pipeline/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..32df5d1be232146556485523e9d9723894f41583 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/prompts.py @@ -0,0 +1,555 @@ +""" +Prompts and VLM interaction layer โ€” DefectDiffu Edition + +Key changes from FLUX version: + - DefectPlan now carries DefectDiffu fields: c_p, c_d, c_f, w_d, w_p + - Planning prompt instructs VLM to generate three disentangled prompts + - No more artifact_type โ†’ FLUX primitive mapping (addition/removal/distortion/fusion) + - Severity directly maps to double-free w_d scale +""" + +import json +from PIL import Image +import base64 +import io +import numpy as np +import re +from pipeline.local_vlm_client import LocalVLMClient, ChatCompletion +from typing import Union, List, Optional, Dict +from pydantic import BaseModel, Field +import os + +DEFAULT_MODEL = "gemma3:12b" +default_client = LocalVLMClient(model=DEFAULT_MODEL) +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) + + +# ============================================================================= +# Pydantic schemas +# ============================================================================= + +class DefectPlan(BaseModel): + """Single defect plan with DefectDiffu-native fields.""" + defect_type: str = Field(description="Type of defect (bubble, scratch, crack, etc.)") + artifact_type: Optional[str] = Field( + default=None, + description="Legacy artifact category (addition/removal/distortion/fusion) โ€” kept for compatibility" + ) + target_entity: str = Field(description="Main product entity/component") + target_subentity: Optional[str] = Field(None, description="Specific sub-part") + location_hint: str = Field(description="Spatial location") + defect_coverage_ratio: float = Field( + default=0.5, + description="Coverage ratio of defect relative to target ROI (float between 0.10 and 1.0)" + ) + mask_shape: str = Field( + default="free", + description="Boundary geometry of the defect mask: 'circle', 'square', 'rectangle', or 'free'" + ) + description: str = Field(description="Detailed visual description") + severity: str = Field(default="medium", description="low/medium/high") + + # --- DefectDiffu-native fields --- + c_p: str = Field(default="", description="Background/product consistency prompt") + c_d: str = Field(default="", description="Defect consistency prompt") + c_f: str = Field(default="", description="Fusion prompt") + w_d: float = Field(default=1.0, description="Defect strength scale (double-free)") + w_p: float = Field(default=1.0, description="Product consistency scale") + + +class ProductDefectPlan(BaseModel): + product_type: str = Field(description="Identified product category") + analysis: str = Field(description="Visual inspection analysis") + possible_defects: List[DefectPlan] = Field(description="List of proposed defects") + + +class BboxResponse(BaseModel): + bboxes: List[List[int]] # [[x1, y1, x2, y2], ...] + + +class ArtifactDescriptionResponse(BaseModel): + has_artifact: bool + explanation: str + label: str + + +class ArtifactExplanationResponse(BaseModel): + explanation: str + + +class ArtifactSuccessResponse(BaseModel): + reasoning: str + success: bool + + +class VocabResponse(BaseModel): + peripheral: Optional[Dict[str, List[str]]] = None + intermediate: Optional[Dict[str, List[str]]] = None + + +# ============================================================================= +# Image encoding utility +# ============================================================================= + +def encode_image_to_base64(image): + try: + if isinstance(image, str): + if not os.path.exists(image): + print(f"Warning: File path '{image}' not found. Skipping...") + return "" + pil_image = Image.open(image) + elif isinstance(image, np.ndarray): + pil_image = Image.fromarray(image) + else: + pil_image = image + + if pil_image.mode != 'RGB': + pil_image = pil_image.convert('RGB') + + buffer = io.BytesIO() + pil_image.save(buffer, format='JPEG') + buffer.seek(0) + return base64.b64encode(buffer.getvalue()).decode('utf-8') + except Exception as e: + print(f"Warning: Could not encode image ({e}). Skipping...") + return "" + + +def clean_json_string(raw_string: str) -> str: + if not isinstance(raw_string, str): + return raw_string + cleaned = raw_string.strip() + cleaned = re.sub(r'^```(?:json)?\\s*', '', cleaned, flags=re.IGNORECASE) + cleaned = re.sub(r'\\s*```$', '', cleaned) + return cleaned.strip() + + +# ============================================================================= +# Step 1: Defect Planning (DefectDiffu-native) +# ============================================================================= + +def plan_defects_for_product( + client, + product_description: str, + image, + money_manager=None, + target_defect_type: Optional[str] = None, + max_defects: int = 3 +): + """ + Agent reasons about possible defects and outputs DefectDiffu-native plans. + Each plan includes the three text prompts (c_p, c_d, c_f) and w_d/w_p scales. + """ + if client is None: + client = default_client + base64_image = encode_image_to_base64(image) + + defect_instruction = "" + if target_defect_type: + defect_instruction = ( + f"\\nCRITICAL REQUIREMENT: The user specifically requested defects of type: '{target_defect_type}'." + f"\\nAll proposed defects MUST represent or relate to a '{target_defect_type}' defect." + ) + + system_prompt = f""" +You are an expert manufacturing quality control engineer and defect analysis agent. + +You are given: +1. A product description: "{product_description}" +2. A clean reference image of the product +{defect_instruction} + +Your task is to analyze the product and propose 1-{max_defects} realistic manufacturing or handling defects. + +For EACH defect, you must produce THREE text prompts for the DefectDiffu diffusion model: + - c_p (background prompt): "A photo of " + - c_d (defect prompt): "A photo of " + - c_f (fusion prompt): "A photo of with " + +Also assign a severity level which maps to defect strength w_d: + - "low" / "minor" โ†’ w_d = 0.6 (subtle defect) + - "medium" / "moderate" โ†’ w_d = 1.0 (standard defect) + - "high" / "severe" โ†’ w_d = 1.5 (strong defect) + +Assign defect_coverage_ratio based on defect category (MUST be between 0.20 and 0.50): + - Fine/Linear defects (crack, scratch, pinhole): 0.20 - 0.30 (20% to 30% coverage) + - Localized surface defects (bubble, contamination, dent): 0.25 - 0.35 (25% to 35% coverage) + - Large-scale/Broad defects (warp, severe damage, broad stain): 0.35 - 0.50 (35% to 50% coverage) + +Assign the mask shape of defect with Options: "circle", "square", "rectangle" or "free" + +Think step by step, then output EXACTLY one JSON object with this structure: +{{ + "product_type": "short product name", + "analysis": "one sentence summary of plausible defects", + "possible_defects": [ + {{ + "defect_type": "bubble", + "description": "what the defect looks like", + "target_entity": "name of affected component", + "target_subentity": "specific part", + "location_hint": "where on the image", + "defect_coverage_ratio": 0.30, + "mask_shape": "free", + "severity": "medium", + "c_p": "A photo of VCSEL laser diode", + "c_d": "A photo of a small transparent bubble trapped under glass", + "c_f": "A photo of VCSEL laser diode with a small transparent bubble trapped under glass", + "w_d": 1.0, + "w_p": 1.0 + }} + ] +}} + +IMPORTANT RULES: +1. c_p, c_d, c_f MUST start with "A photo of". +2. c_d should describe the defect in isolation (no product context). +3. c_f combines product + defect naturally. +4. Use concrete, visually distinctive target_entity names. +5. Do NOT target thin wire-like structures as primary target_entity. +6. If the defect is a trapped air pocket, use the word "bubble" in c_d and description. +7. If the defect is a surface scratch, use the word "scratch" in c_d and description. +8. defect_type vocabulary: bubble, scratch, crack, contamination, dent, warp, missing_part, etc. +9. Return ONLY the keys shown above. Do NOT return has_artifact, explanation, or label. +""" + + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{base64_image}"}, + {"type": "input_text", "text": f"Product: {product_description}\\nAnalyze and propose realistic defects with DefectDiffu prompts."} + ] + } + ], + temperature=0.3, + text_format=ProductDefectPlan + ) + + if money_manager: + money_manager(response) + + if hasattr(response, "output_parsed") and response.output_parsed: + return response.output_parsed + + # Fallback parsing for raw string output + raw_content = getattr(response, "content", str(response)) + cleaned_text = clean_json_string(raw_content) + parsed = json.loads(cleaned_text) + + if isinstance(parsed, list): + defects = [] + for item in parsed: + defects.append({ + "defect_type": item.get("defect_type", "addition"), + "description": item.get("description", item.get("explanation", "")), + "target_entity": item.get("target_entity", item.get("entity", "")), + "target_subentity": item.get("target_subentity", item.get("subentity")), + "location_hint": item.get("location_hint", ""), + "defect_coverage_ratio": item.get("defect_coverage_ratio", 0.25), + "mask_shape": item.get("mask_shape", "free"), + "severity": item.get("severity", "medium"), + "c_p": item.get("c_p", f"A photo of {product_description}"), + "c_d": item.get("c_d", f"A photo of {item.get('description', 'defect')}"), + "c_f": item.get("c_f", f"A photo of {product_description} with {item.get('description', 'defect')}"), + "w_d": item.get("w_d", 1.0), + "w_p": item.get("w_p", 1.0), + }) + plan_dict = { + "product_type": product_description.split()[0] if product_description else "component", + "analysis": f"Planned {len(defects)} defects", + "possible_defects": defects + } + return ProductDefectPlan(**plan_dict) + + elif isinstance(parsed, dict): + if "possible_defects" not in parsed: + for key in ["defects", "defect_plan", "results"]: + if key in parsed: + parsed["possible_defects"] = parsed.pop(key) + break + if "possible_defects" in parsed: + for d in parsed["possible_defects"]: + if "explanation" in d and "description" not in d: + d["description"] = d.pop("explanation") + if "entity" in d and "target_entity" not in d: + d["target_entity"] = d.pop("entity") + # Ensure DefectDiffu fields exist + d.setdefault("defect_coverage_ratio", 0.25) # <-- ADDED + d.setdefault("mask_shape", "free") # <-- ADDED + d.setdefault("c_p", f"A photo of {product_description}") + d.setdefault("c_d", f"A photo of {d.get('description', 'defect')}") + d.setdefault("c_f", f"A photo of {product_description} with {d.get('description', 'defect')}") + d.setdefault("w_d", 1.0) + d.setdefault("w_p", 1.0) + return ProductDefectPlan(**parsed) + + except Exception as e: + print(f"Error in defect planning: {e}") + return None + + +# ============================================================================= +# Entity vocabulary & bbox detection (unchanged) +# ============================================================================= + +def get_entity_subentities(client, image, money_manager=None): + if client is None: + client = default_client + base64_image = encode_image_to_base64(image) + system_prompt = """ +You are given a microscopic or optical inspection image of a precision component. +Identify visible entities and subentities, split into peripheral and intermediate layers. + +Output exactly one JSON object with two keys: "peripheral" and "intermediate". +Each value is a dict mapping entity names to lists of subentity names. + +Hard rules: +1) Each entity MUST have at least one subentity. +2) Subentities must be clearly visible and segmentable. +3) Use concise, lowercase nouns. +4) Do not invent occluded parts. +""" + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": [{"type": "input_image", "image_url": f"data:image/jpeg;base64,{base64_image}"}]} + ], + temperature=0.2, + text_format=VocabResponse + ) + if money_manager: + money_manager(response) + if hasattr(response, "output_parsed") and response.output_parsed: + return response.output_parsed + raw_content = getattr(response, "content", str(response)) + cleaned_text = clean_json_string(raw_content) + return VocabResponse.model_validate_json(cleaned_text) + except Exception as e: + print(f"Error in entity analysis: {e}") + return None + + +def get_entity_bboxes(client, image, entity_name, money_manager=None): + if client is None: + client = default_client + base64_image = encode_image_to_base64(image) + system_prompt = f""" +You are an industrial inspection bounding box detector. +Find ALL instances of: {entity_name} +Return JSON with "bboxes" key: [[x1, y1, x2, y2], ...] in ABSOLUTE PIXEL coordinates. +If none found, return {{"bboxes": []}}. +""" + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": [{"type": "input_image", "image_url": f"data:image/jpeg;base64,{base64_image}"}]} + ], + temperature=0.2, + text_format=BboxResponse + ) + if money_manager: + money_manager(response) + if hasattr(response, "output_parsed") and response.output_parsed: + return response.output_parsed.bboxes + raw_content = getattr(response, "content", str(response)) + cleaned = clean_json_string(raw_content) + parsed = json.loads(cleaned) + return parsed.get("bboxes", []) if isinstance(parsed, dict) else [] + except Exception as e: + print(f"Error getting bboxes for {entity_name}: {e}") + return [] + + +# ============================================================================= +# Verification prompts (unchanged structure, adapted for manufacturing) +# ============================================================================= + +def artifact_description(client, masked_original_image, target_original_image, + target_artifact_image, object_name, artifact_type, money_manager=None): + if client is None: + client = default_client + + original_masked = encode_image_to_base64(masked_original_image) + original_target = encode_image_to_base64(target_original_image) + artifact_target = encode_image_to_base64(target_artifact_image) + + combined_instruction = { + "addition": ( + "You are an expert industrial quality control inspector.\\n\\n" + "ANOMALY DEFINITION (ADDITION): extra material, dust, debris, contamination, " + "solder blob, particulate, smudge, residue, bubble, scratch.\\n\\n" + "Checklist: 1) Is there extra material in the target region? " + "2) Does it have a distinct boundary? 3) Is it consistent with the defect class?" + ), + "removal": ( + "You are an expert industrial quality control inspector.\\n\\n" + "ANOMALY DEFINITION (REMOVAL): missing part, lifted pad, absent component, " + "incomplete geometric path, broken lead.\\n\\n" + "Checklist: 1) Is there a clear gap or missing structure? " + "2) Is the absence NOT explained by occlusion or viewpoint?" + ), + "distortion": ( + "You are an expert industrial quality control inspector.\\n\\n" + "ANOMALY DEFINITION (DISTORTION): bent, warped, misaligned, cracked, " + "deformed, wavy, stress pattern, bowing.\\n\\n" + "Checklist: 1) Are straight lines or edges warped? " + "2) Is geometric symmetry broken?" + ), + "fusion": ( + "You are an expert industrial quality control inspector.\\n\\n" + "ANOMALY DEFINITION (FUSION/BRIDGING): solder bridge, shorted pins, " + "merged adjacent structures.\\n\\n" + "Checklist: 1) Are separate conductive paths unnaturally merged? " + "2) Is the boundary between components degraded?" + ) + } + + prompt = f""" +{combined_instruction.get(artifact_type, combined_instruction["addition"])} + +Return JSON: + "has_artifact": true/false + "explanation": "Detailed description of what looks wrong" + "label": "Brief description (empty if no artifact)" + +Rules: Focus on visible evidence. Do not refer to images by number. + +CRITICAL VERIFICATION RULE: +Before answering, compare the second image (original target) and third image (artifact target) pixel-by-pixel. +If they are visually IDENTICAL with no discernible difference in texture, color, or structure, you MUST return: + "has_artifact": false + "explanation": "No visible anomaly detected in the target region." + "label": "" +Only return has_artifact=true if you can point to a specific, concrete visual difference. +""" + + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": prompt}, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_masked}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_target}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{artifact_target}"}, + {"type": "input_text", "text": f"{object_name}"} + ] + } + ], + temperature=0.2, + text_format=ArtifactDescriptionResponse + ) + if money_manager: + money_manager(response) + return response.output_parsed + except Exception as e: + print(f"Error in artifact description ({artifact_type}): {e}") + return ArtifactDescriptionResponse(has_artifact=False, explanation="", label="") + + +# ============================================================================= +# Explanation generation (unchanged) +# ============================================================================= + +def artifact_explanation_from_triplet(client, masked_original_image, target_original_image, + target_artifact_image, object_name, artifact_type, money_manager=None): + if client is None: + client = default_client + + original_masked = encode_image_to_base64(masked_original_image) + original_target = encode_image_to_base64(target_original_image) + artifact_target = encode_image_to_base64(target_artifact_image) + + type_guidance = { + "addition": "Describe an ADDED instance: duplicated parts, extra elements, foreign material.", + "removal": "Describe a MISSING part: gaps, smoothed-over areas, discontinuity.", + "distortion": "Describe WARPING: bent shapes, irregular textures, malformed geometry.", + "fusion": "Describe MERGING: boundary loss, texture bleed, interpenetration." + } + + prompt = f""" +You receive three images and an object name: + 1) Original WITHOUT target region + 2) Original showing ONLY target region + 3) Artifact showing ONLY target region (describe THIS one) + 4) Object name: {object_name} + +TASK: One-sentence description of what looks wrong in Image 3, consistent with {artifact_type}. +{type_guidance.get(artifact_type, "")} + +Rules: visible evidence only, no JSON, no image numbers, simple language. +""" + + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": prompt}, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_masked}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_target}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{artifact_target}"}, + {"type": "input_text", "text": f"{object_name}"} + ] + } + ], + temperature=0.2, + text_format=ArtifactExplanationResponse + ) + if money_manager: + money_manager(response) + return response.output_parsed + except Exception as e: + print(f"Error in explanation ({artifact_type}): {e}") + return ArtifactExplanationResponse(explanation="") + + +# ============================================================================= +# Cost tracking (unchanged) +# ============================================================================= + +class MoneyManager: + def __init__(self, model: str = DEFAULT_MODEL): + self.total_cost = 0.0 + self.model = model + # ... (same cost table as before, truncated for brevity) ... + cost_table = { + "gpt-4o": (2.5/1000, 10/1000), + "gpt-4o-mini": (0.15/1000, 0.6/1000), + "gemini-2.5-flash": (0.3/1000, 2.5/1000), + "gemini-2.5-pro": (1.25/1000, 10/1000), + } + self.input_cost, self.output_cost = cost_table.get(model, (0.0, 0.0)) + + def __call__(self, response=None): + if response is None or not hasattr(response, "usage"): + return + try: + if hasattr(response.usage, "input_tokens"): + inp = response.usage.input_tokens + out = response.usage.output_tokens + elif hasattr(response, "usage_metadata"): + inp = response.usage_metadata.prompt_token_count + out = (response.usage_metadata.candidates_token_count + + getattr(response.usage_metadata, "thoughts_token_count", 0)) + else: + return + self.total_cost += (inp / 1000 * self.input_cost + out / 1000 * self.output_cost) + except Exception: + pass + + def refresh(self): + self.total_cost = 0.0 diff --git a/ArtiAgent - DefectDiffu/src/pipeline/sample_manifest.csv b/ArtiAgent - DefectDiffu/src/pipeline/sample_manifest.csv new file mode 100644 index 0000000000000000000000000000000000000000..0aceaad08e7ac89e602f28759013e245288e2443 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/sample_manifest.csv @@ -0,0 +1,8 @@ +image_path,product_description +vcsel/clean_01.png,VCSEL laser diode with emission aperture and surrounding mesa structure +vcsel/clean_02.png,VCSEL laser diode with emission aperture and surrounding mesa structure +lens/clean_01.png,Optical lens with anti-reflective coating and mounting frame +lens/clean_02.png,Optical lens with anti-reflective coating and mounting frame +die/clean_01.png,Semiconductor die with bond pads and scribe lines +photodiode/clean_01.png,Photodiode sensor with photosensitive area and electrode contacts +optical_sensor/clean_01.png,Optical sensor with active area and reflective cavity diff --git a/ArtiAgent - DefectDiffu/src/pipeline/test_prompts.py b/ArtiAgent - DefectDiffu/src/pipeline/test_prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..24ce945da3961923b844377905b62f718987cd71 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/pipeline/test_prompts.py @@ -0,0 +1,13 @@ +from prompts import clean_json_string, VocabResponse + +# Test 1: clean_json_string +test = '```json\n{"peripheral": {"pin": ["lead"]}}\n```' +result = clean_json_string(test) +print("Test 1 - clean_json_string:") +print(repr(result)) +print() + +# Test 2: VocabResponse validation +v = VocabResponse.model_validate_json(result) +print("Test 2 - VocabResponse validation:") +print(v.peripheral) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/scripts/clean_up_stale_records.py b/ArtiAgent - DefectDiffu/src/scripts/clean_up_stale_records.py new file mode 100644 index 0000000000000000000000000000000000000000..f35843c426a84f808ffdc8483303785f00b62db8 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/scripts/clean_up_stale_records.py @@ -0,0 +1,33 @@ +import json +from pathlib import Path + +import sys +import os +# Add parent directory (src/) to path so 'pipeline' is findable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from pipeline.defect_rag import get_rag + +rag = get_rag() +data = rag.collection.get() + +stale_ids = [] + +for doc_id, metadata in zip(data['ids'], data['metadatas']): + # Parse stored patch paths + paths = json.loads(metadata['paths']) + + # Check key with .png extension + sample_path_str = paths.get('original_masked.png') or list(paths.values())[0] + sample_file = Path(sample_path_str) + + # If the file/folder no longer exists on disk, mark for deletion + if not sample_file.exists(): + print(f" โŒ Missing on disk: {doc_id} -> {sample_file}") + stale_ids.append(doc_id) + +if stale_ids: + rag.collection.delete(ids=stale_ids) + print(f"\n๐Ÿงน Cleaned up {len(stale_ids)} stale record(s) from database.") +else: + print("\nโœจ Database is clean! All entries match folders on disk.") \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/scripts/convert_roboflow_to_triplets.py b/ArtiAgent - DefectDiffu/src/scripts/convert_roboflow_to_triplets.py new file mode 100644 index 0000000000000000000000000000000000000000..f9cd12f154c4db48771f7e0a4f1121e9b9bbb0b0 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/scripts/convert_roboflow_to_triplets.py @@ -0,0 +1,102 @@ +# scripts/convert_roboflow_to_triplets.py +import json +import shutil +from pathlib import Path +from PIL import Image +import numpy as np + +def convert_roboflow_to_triplets(roboflow_dir: str, output_dir: str): + """ + Convert Roboflow COCO format to patch triplet format. + + Roboflow COCO structure: + train/ + _annotations.coco.json + image1.jpg + image2.jpg + """ + roboflow_path = Path(roboflow_dir) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Load COCO annotations + anno_file = roboflow_path / "train" / "_annotations.coco.json" + with open(anno_file) as f: + coco = json.load(f) + + # Build image ID โ†’ filename map + images = {img['id']: img['file_name'] for img in coco['images']} + + # Group annotations by image + from collections import defaultdict + image_annotations = defaultdict(list) + for ann in coco['annotations']: + image_annotations[ann['image_id']].append(ann) + + triplet_id = 0 + for img_id, filename in images.items(): + img_path = roboflow_path / "train" / filename + if not img_path.exists(): + continue + + img = Image.open(img_path).convert('RGB') + img_array = np.array(img) + H, W = img_array.shape[:2] + + anns = image_annotations.get(img_id, []) + if not anns: + continue + + for ann in anns: + # 1. Skip non-defect category annotations (category_id == 2) + if ann.get('category_id') == 2: + continue + + # 2. Skip full-image bounding boxes + x, y, w, h = ann['bbox'] + if w >= W and h >= H: + continue + + triplet_id += 1 + triplet_dir = output_path / f"defect_{triplet_id:04d}" + triplet_dir.mkdir(exist_ok=True) + + # Get bbox + x, y, w, h = ann['bbox'] + x1, y1, x2, y2 = int(x), int(y), int(x+w), int(y+h) + + # Create mask from segmentation if available, else bbox + if 'segmentation' in ann and ann['segmentation']: + # COCO polygon segmentation + from pycocotools import mask as maskUtils + rles = maskUtils.frPyObjects(ann['segmentation'], H, W) + mask = maskUtils.decode(rles) + if len(mask.shape) == 3: + mask = np.any(mask, axis=2).astype(np.uint8) * 255 + else: + # Fallback: bbox mask + mask = np.zeros((H, W), dtype=np.uint8) + mask[y1:y2, x1:x2] = 255 + + # Save clean image (original_target) + img.save(triplet_dir / "original_target.png") + + # Save mask (original_masked) + Image.fromarray(mask).save(triplet_dir / "original_masked.png") + + # For artifact_target, we need the defective version. + # Since this is a real defect dataset, the original image IS the defect. + # For synthetic training, you may want to inpaint the defect out to create "clean", + # but for RAG retrieval, we can use the same image as artifact_target. + img.save(triplet_dir / "artifact_target.png") + + print(f"Created triplet {triplet_id}: {filename} โ†’ {triplet_dir}") + + print(f"\nTotal triplets created: {triplet_id}") + print(f"Output: {output_path}") + +if __name__ == "__main__": + convert_roboflow_to_triplets( + roboflow_dir="./data/external/Manufacturing_Defect_Detection", + output_dir="data/external/roboflow_manufacturing" + ) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/scripts/ingest_public_data.py b/ArtiAgent - DefectDiffu/src/scripts/ingest_public_data.py new file mode 100644 index 0000000000000000000000000000000000000000..a8a4afe4def500aa93833c20208dfba3d8735eab --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/scripts/ingest_public_data.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# scripts/ingest_public_data.py +# ============================================ +# RUN ONCE: python scripts/ingest_public_data.py +# Populates ChromaDB with defect patch triplets +# ============================================ + +import chromadb +from sentence_transformers import SentenceTransformer +from PIL import Image +import json +from pathlib import Path +import argparse + + +# ============================================ +# DATASET REGISTRY โ€” Add new sources here +# ============================================ + +DATASETS = { + # โ”€โ”€ Commercial-friendly datasets โ”€โ”€ + "roboflow_manufacturing": { + "path": "data/external/roboflow_manufacturing", + "domain": "pcb", + "license": "cc-by-4.0", + "commercial_ok": True, + "attribution": "Roboflow Universe โ€” Manufacturing Defect Detection", + "description": "PCB manufacturing defects, CC BY 4.0 licensed" + }, + + # โ”€โ”€ Your proprietary data โ”€โ”€ + "my_vcsel_captures": { + "path": "data/my_captures/vcsel", + "domain": "glass/vcsel", + "license": "proprietary", + "commercial_ok": True, + "attribution": "Internal", + "description": "VCSEL laser diode defect captures" + }, + + # โ”€โ”€ Partner data (NDA) โ”€โ”€ + # "partner_pcb_nda": { + # "path": "data/partners/pcb_client_a", + # "domain": "pcb", + # "license": "partner_nda", + # "commercial_ok": True, + # "attribution": "Partner NDA", + # "description": "PCB client defect images under NDA" + # }, +} + + +def validate_triplet(folder: Path) -> dict | None: + """Check if folder contains required patch triplet files.""" + required = ["original_masked.png", "original_target.png", "artifact_target.png"] + paths = {f: folder / f for f in required} + + for f, p in paths.items(): + if not p.exists(): + print(f" โš ๏ธ Missing {f} in {folder}, skipping") + return None + + return {k: str(v) for k, v in paths.items()} + + +def generate_caption(folder_name: str, domain: str, dataset_desc: str) -> str: + """Generate a text caption for embedding.""" + # You can replace this with a VLM call for better captions + defect_type = folder_name.replace("_", " ") + return f"{defect_type} defect on {domain}. {dataset_desc}" + + +def ingest_dataset(collection, encoder, name: str, cfg: dict): + """Ingest one dataset into ChromaDB.""" + print(f"\n{'='*50}") + print(f"Dataset: {name}") + print(f"License: {cfg['license']} | Commercial: {'โœ…' if cfg['commercial_ok'] else 'โŒ'}") + print(f"Path: {cfg['path']}") + print(f"{'='*50}") + + base_path = Path(cfg["path"]) + if not base_path.exists(): + print(f" โš ๏ธ Path not found: {base_path}") + print(f" Create it and add patch triplet folders:") + print(f" {base_path}/defect_name_001/original_masked.png") + print(f" {base_path}/defect_name_001/original_target.png") + print(f" {base_path}/defect_name_001/artifact_target.png") + return 0 + + # NEW: Get set of IDs already in the collection to avoid re-processing + existing_ids = set(collection.get()["ids"]) + + count = 0 + for triplet_folder in sorted(base_path.iterdir()): + if not triplet_folder.is_dir(): + continue + + # NEW: Skip if already processed + doc_id = f"{name}_{triplet_folder.name}" + if doc_id in existing_ids: + print(f" โญ๏ธ Skipping {triplet_folder.name} (already in DB)") + continue + + paths = validate_triplet(triplet_folder) + if paths is None: + continue + + caption = generate_caption( + triplet_folder.name, + cfg["domain"], + cfg.get("description", "") + ) + + embedding = encoder.encode(caption) + + collection.upsert( + documents=[caption], + embeddings=[embedding.tolist()], + metadatas=[{ + "paths": json.dumps(paths), + "domain": cfg["domain"], + "license": cfg["license"], + "commercial_ok": cfg["commercial_ok"], + "source": name, + "attribution": cfg["attribution"], + "defect_name": triplet_folder.name + }], + ids=[f"{name}_{triplet_folder.name}"] + ) + + count += 1 + print(f" โœ… {triplet_folder.name}: {caption[:60]}...") + + return count + + +def main(): + parser = argparse.ArgumentParser(description="Ingest defect patch triplets into RAG DB") + parser.add_argument("--db-path", default="data/defect_db", help="ChromaDB persistent path") + parser.add_argument("--collection", default="defect_patches", help="Collection name") + parser.add_argument("--model", default="all-MiniLM-L6-v2", help="SentenceTransformer model") + args = parser.parse_args() + + # Initialize DB + Path(args.db_path).mkdir(parents=True, exist_ok=True) + client = chromadb.PersistentClient(path=args.db_path) + + # Delete existing collection if you want fresh start + # client.delete_collection(args.collection) + + collection = client.get_or_create_collection( + name=args.collection, + metadata={"hnsw:space": "cosine"} + ) + + print(f"DB path: {args.db_path}") + print(f"Collection: {args.collection}") + print(f"Existing entries: {collection.count()}") + + # Initialize encoder + print(f"\nLoading encoder: {args.model}") + encoder = SentenceTransformer(args.model) + + # Ingest all datasets + total = 0 + for name, cfg in DATASETS.items(): + # Skip non-commercial datasets in commercial builds + if not cfg.get("commercial_ok", False): + print(f"\nโญ๏ธ Skipping {name} โ€” not commercial-friendly") + continue + + count = ingest_dataset(collection, encoder, name, cfg) + total += count + + print(f"\n{'='*50}") + print(f"TOTAL INGESTED: {total} patch triplets") + print(f"TOTAL IN DB: {collection.count()}") + print(f"{'='*50}") + + # Print commercial summary + print("\n๐Ÿ“‹ Commercial License Summary:") + results = collection.get() + licenses = {} + for meta in results["metadatas"]: + lic = meta["license"] + licenses[lic] = licenses.get(lic, 0) + 1 + + for lic, count in licenses.items(): + icon = "โœ…" if any( + cfg["license"] == lic and cfg.get("commercial_ok") + for cfg in DATASETS.values() + ) else "โŒ" + print(f" {icon} {lic}: {count} entries") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/scripts/resize-file-for-training.py b/ArtiAgent - DefectDiffu/src/scripts/resize-file-for-training.py new file mode 100644 index 0000000000000000000000000000000000000000..bc5197f5068ee849ef8592fe1b55b48cdf396707 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/scripts/resize-file-for-training.py @@ -0,0 +1,53 @@ +import os +from PIL import Image + +def preprocess_dataset(input_dir, output_dir, target_size=(512, 512)): + """ + Center crops images to a square, resizes them to 512x512, + and converts them to PNG format. + """ + os.makedirs(output_dir, exist_ok=True) + + valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff') + + for root, _, files in os.walk(input_dir): + for file in files: + if file.lower().endswith(valid_extensions): + # Setup paths + rel_path = os.path.relpath(root, input_dir) + out_folder = os.path.join(output_dir, rel_path) + os.makedirs(out_folder, exist_ok=True) + + img_path = os.path.join(root, file) + filename_without_ext = os.path.splitext(file)[0] + save_path = os.path.join(out_folder, f"{filename_without_ext}.png") + + with Image.open(img_path) as img: + w, h = img.size + + # 1. Calculate center crop box + min_dim = min(w, h) + left = (w - min_dim) // 2 + top = (h - min_dim) // 2 + right = left + min_dim + bottom = top + min_dim + + # 2. Crop to square + img_cropped = img.crop((left, top, right, bottom)) + + # 3. Resize to target resolution (512x512) + # For masks (binary), use NEAREST; for images, use LANCZOS + if "ground_truth" in root.lower() or "mask" in root.lower(): + img_resized = img_cropped.resize(target_size, Image.Resampling.NEAREST) + else: + img_resized = img_cropped.resize(target_size, Image.Resampling.LANCZOS) + + # 4. Save as PNG + img_resized.save(save_path, "PNG") + print(f"Processed: {file} -> {save_path}") + +# Example Usage: +preprocess_dataset( + input_dir="./engine/DefectDiffu/few-shot-training/vscel/img/tiger-strip", + output_dir="./engine/DefectDiffu/few-shot-training/vcsel_dataset_512" +) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/scripts/test-chromadb.py b/ArtiAgent - DefectDiffu/src/scripts/test-chromadb.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e82e13258a9ae77627e0f430eaf4b0edfc9dd4 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/scripts/test-chromadb.py @@ -0,0 +1,19 @@ +import sys +import os +import json +# Add parent directory (src/) to path so 'pipeline' is findable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from pipeline.defect_rag import get_rag + +rag = get_rag() +all_data = rag.collection.get() + +print("Found IDs in ChromaDB:") +for doc_id in all_data['ids']: + print(" -", doc_id) + +results = rag.collection.get(ids=["my_vcsel_captures_bubble_001"]) # or any ID you know exists +print(results['metadatas'][0]) +print("---") +print(json.loads(results['metadatas'][0]['paths'])) \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/scripts/test-safetensor.py b/ArtiAgent - DefectDiffu/src/scripts/test-safetensor.py new file mode 100644 index 0000000000000000000000000000000000000000..af5c848160d04e2c5b7f6b7ce88cfad83f43c4ff --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/scripts/test-safetensor.py @@ -0,0 +1,12 @@ +import os +from safetensors.torch import load_file + +ckpt_path = r"C:\path\to\your\flux1-dev.safetensors" # or wherever your checkpoint is +print(f"File size: {os.path.getsize(ckpt_path) / 1e9:.2f} GB") + +state_dict = load_file(ckpt_path) +print(f"Total keys in checkpoint: {len(state_dict)}") + +# Check which double_blocks are present +blocks = sorted(set([k.split('.')[1] for k in state_dict.keys() if 'double_blocks' in k])) +print(f"Blocks present: {blocks}") \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/.flake8 b/ArtiAgent - DefectDiffu/src/segment_anything/.flake8 new file mode 100644 index 0000000000000000000000000000000000000000..6b0759587aa5756e66a13ef034c6bcdd76a885f5 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/.flake8 @@ -0,0 +1,7 @@ +[flake8] +ignore = W503, E203, E221, C901, C408, E741, C407, B017, F811, C101, EXE001, EXE002 +max-line-length = 100 +max-complexity = 18 +select = B,C,E,F,W,T4,B9 +per-file-ignores = + **/__init__.py:F401,F403,E402 diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/CODE_OF_CONDUCT.md b/ArtiAgent - DefectDiffu/src/segment_anything/CODE_OF_CONDUCT.md new file mode 100644 index 0000000000000000000000000000000000000000..08b500a221857ec3f451338e80b4a9ab1173a1af --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/CODE_OF_CONDUCT.md @@ -0,0 +1,80 @@ +# Code of Conduct + +## Our Pledge + +In the interest of fostering an open and welcoming environment, we as +contributors and maintainers pledge to make participation in our project and +our community a harassment-free experience for everyone, regardless of age, body +size, disability, ethnicity, sex characteristics, gender identity and expression, +level of experience, education, socio-economic status, nationality, personal +appearance, race, religion, or sexual identity and orientation. + +## Our Standards + +Examples of behavior that contributes to creating a positive environment +include: + +* Using welcoming and inclusive language +* Being respectful of differing viewpoints and experiences +* Gracefully accepting constructive criticism +* Focusing on what is best for the community +* Showing empathy towards other community members + +Examples of unacceptable behavior by participants include: + +* The use of sexualized language or imagery and unwelcome sexual attention or + advances +* Trolling, insulting/derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or electronic + address, without explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Our Responsibilities + +Project maintainers are responsible for clarifying the standards of acceptable +behavior and are expected to take appropriate and fair corrective action in +response to any instances of unacceptable behavior. + +Project maintainers have the right and responsibility to remove, edit, or +reject comments, commits, code, wiki edits, issues, and other contributions +that are not aligned to this Code of Conduct, or to ban temporarily or +permanently any contributor for other behaviors that they deem inappropriate, +threatening, offensive, or harmful. + +## Scope + +This Code of Conduct applies within all project spaces, and it also applies when +an individual is representing the project or its community in public spaces. +Examples of representing a project or community include using an official +project e-mail address, posting via an official social media account, or acting +as an appointed representative at an online or offline event. Representation of +a project may be further defined and clarified by project maintainers. + +This Code of Conduct also applies outside the project spaces when there is a +reasonable belief that an individual's behavior may have a negative impact on +the project or its community. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported by contacting the project team at . All +complaints will be reviewed and investigated and will result in a response that +is deemed necessary and appropriate to the circumstances. The project team is +obligated to maintain confidentiality with regard to the reporter of an incident. +Further details of specific enforcement policies may be posted separately. + +Project maintainers who do not follow or enforce the Code of Conduct in good +faith may face temporary or permanent repercussions as determined by other +members of the project's leadership. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, +available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see +https://www.contributor-covenant.org/faq diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/CONTRIBUTING.md b/ArtiAgent - DefectDiffu/src/segment_anything/CONTRIBUTING.md new file mode 100644 index 0000000000000000000000000000000000000000..263991c9496cf29ed4b99e03a9fb9a38e6bfaf86 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/CONTRIBUTING.md @@ -0,0 +1,31 @@ +# Contributing to segment-anything +We want to make contributing to this project as easy and transparent as +possible. + +## Pull Requests +We actively welcome your pull requests. + +1. Fork the repo and create your branch from `main`. +2. If you've added code that should be tested, add tests. +3. If you've changed APIs, update the documentation. +4. Ensure the test suite passes. +5. Make sure your code lints, using the `linter.sh` script in the project's root directory. Linting requires `black==23.*`, `isort==5.12.0`, `flake8`, and `mypy`. +6. If you haven't already, complete the Contributor License Agreement ("CLA"). + +## Contributor License Agreement ("CLA") +In order to accept your pull request, we need you to submit a CLA. You only need +to do this once to work on any of Facebook's open source projects. + +Complete your CLA here: + +## Issues +We use GitHub issues to track public bugs. Please ensure your description is +clear and has sufficient instructions to be able to reproduce the issue. + +Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe +disclosure of security bugs. In those cases, please go through the process +outlined on that page and do not file a public issue. + +## License +By contributing to segment-anything, you agree that your contributions will be licensed +under the LICENSE file in the root directory of this source tree. diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/LICENSE b/ArtiAgent - DefectDiffu/src/segment_anything/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/README.md b/ArtiAgent - DefectDiffu/src/segment_anything/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6256d2b7f5a387988338d538df4e699eb17ba702 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/README.md @@ -0,0 +1,107 @@ +# Segment Anything + +**[Meta AI Research, FAIR](https://ai.facebook.com/research/)** + +[Alexander Kirillov](https://alexander-kirillov.github.io/), [Eric Mintun](https://ericmintun.github.io/), [Nikhila Ravi](https://nikhilaravi.com/), [Hanzi Mao](https://hanzimao.me/), Chloe Rolland, Laura Gustafson, [Tete Xiao](https://tetexiao.com), [Spencer Whitehead](https://www.spencerwhitehead.com/), Alex Berg, Wan-Yen Lo, [Piotr Dollar](https://pdollar.github.io/), [Ross Girshick](https://www.rossgirshick.info/) + +[[`Paper`](https://ai.facebook.com/research/publications/segment-anything/)] [[`Project`](https://segment-anything.com/)] [[`Demo`](https://segment-anything.com/demo)] [[`Dataset`](https://segment-anything.com/dataset/index.html)] [[`Blog`](https://ai.facebook.com/blog/segment-anything-foundation-model-image-segmentation/)] + +![SAM design](assets/model_diagram.png?raw=true) + +The **Segment Anything Model (SAM)** produces high quality object masks from input prompts such as points or boxes, and it can be used to generate masks for all objects in an image. It has been trained on a [dataset](https://segment-anything.com/dataset/index.html) of 11 million images and 1.1 billion masks, and has strong zero-shot performance on a variety of segmentation tasks. + +

+ + +

+ +## Installation + +The code requires `python>=3.8`, as well as `pytorch>=1.7` and `torchvision>=0.8`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended. + +Install Segment Anything: + +``` +pip install git+https://github.com/facebookresearch/segment-anything.git +``` + +or clone the repository locally and install with + +``` +git clone git@github.com:facebookresearch/segment-anything.git +cd segment-anything; pip install -e . +``` + +The following optional dependencies are necessary for mask post-processing, saving masks in COCO format, the example notebooks, and exporting the model in ONNX format. `jupyter` is also required to run the example notebooks. +``` +pip install opencv-python pycocotools matplotlib onnxruntime onnx +``` + + +## Getting Started + +First download a [model checkpoint](#model-checkpoints). Then the model can be used in just a few lines to get masks from a given prompt: + +``` +from segment_anything import build_sam, SamPredictor +predictor = SamPredictor(build_sam(checkpoint="")) +predictor.set_image() +masks, _, _ = predictor.predict() +``` + +or generate masks for an entire image: + +``` +from segment_anything import build_sam, SamAutomaticMaskGenerator +mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="")) +masks = mask_generator_generate() +``` + +Additionally, masks can be generated for images from the command line: + +``` +python scripts/amg.py --checkpoint --input --output +``` + +See the examples notebooks on [using SAM with prompts](/notebooks/predictor_example.ipynb) and [automatically generating masks](/notebooks/automatic_mask_generator_example.ipynb) for more details. + +

+ + +

+ +## ONNX Export + +SAM's lightweight mask decoder can be exported to ONNX format so that it can be run in any environment that supports ONNX runtime, such as in-browser as showcased in the [demo](https://segment-anything.com/demo). Export the model with + +``` +python scripts/export_onnx_model.py --checkpoint --output +``` + +See the [example notebook](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb) for details on how to combine image preprocessing via SAM's backbone with mask prediction using the ONNX model. It is recommended to use the latest stable version of PyTorch for ONNX export. + +## Model Checkpoints + +Three model versions of the model are available with different backbone sizes. These models can be instantiated by running +``` +from segment_anything import sam_model_registry +sam = sam_model_registry[""](checkpoint="") +``` +Click the links below to download the checkpoint for the corresponding model name. The default model in bold can also be instantiated with `build_sam`, as in the examples in [Getting Started](#getting-started). + +* **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)** +* `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth) +* `vit_b`: [ViT-B SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth) + +## License +The model is licensed under the [Apache 2.0 license](LICENSE). + +## Contributing + +See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md). + +## Contributors + +The Segment Anything project was made possible with the help of many contributors (alphabetical): + +Aaron Adcock, Vaibhav Aggarwal, Morteza Behrooz, Cheng-Yang Fu, Ashley Gabriel, Ahuva Goldstand, Allen Goodman, Sumanth Gurram, Jiabo Hu, Somya Jain, Devansh Kukreja, Robert Kuo, Joshua Lane, Yanghao Li, Lilian Luong, Jitendra Malik, Mallika Malhotra, William Ngan, Omkar Parkhi, Nikhil Raina, Dirk Rowe, Neil Sejoor, Vanessa Stark, Bala Varadarajan, Bram Wasti, Zachary Winstrom diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/linter.sh b/ArtiAgent - DefectDiffu/src/segment_anything/linter.sh new file mode 100644 index 0000000000000000000000000000000000000000..df2e17436d30e89ff1728109301599f425f1ad6b --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/linter.sh @@ -0,0 +1,32 @@ +#!/bin/bash -e +# Copyright (c) Facebook, Inc. and its affiliates. + +{ + black --version | grep -E "23\." > /dev/null +} || { + echo "Linter requires 'black==23.*' !" + exit 1 +} + +ISORT_VERSION=$(isort --version-number) +if [[ "$ISORT_VERSION" != 5.12* ]]; then + echo "Linter requires isort==5.12.0 !" + exit 1 +fi + +echo "Running isort ..." +isort . --atomic + +echo "Running black ..." +black -l 100 . + +echo "Running flake8 ..." +if [ -x "$(command -v flake8)" ]; then + flake8 . +else + python3 -m flake8 . +fi + +echo "Running mypy..." + +mypy --exclude 'setup.py|notebooks' . diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__init__.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b3668adc7817cb24a54cfe4405184a8409c6cb44 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from .build_sam import ( + build_sam, + build_sam_vit_h, + build_sam_vit_l, + build_sam_vit_b, + sam_model_registry, +) +from .build_sam_hq import ( + build_sam_hq, + build_sam_hq_vit_h, + build_sam_hq_vit_l, + build_sam_hq_vit_b, + sam_hq_model_registry, +) +from .predictor import SamPredictor +from .automatic_mask_generator import SamAutomaticMaskGenerator diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92fb0fdb0d5f178dc331050f0aac3aa1d3a387d7 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eed9e46a76b34836e4ef869f5ec1da73639fcf2f Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/build_sam.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/build_sam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de4e075c325215d0aacd2fea2a3a4c3d70a45a0d Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/build_sam.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/build_sam_hq.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/build_sam_hq.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebe56fbc467c52e21016263e1e0995d5cd2c7a39 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/build_sam_hq.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/predictor.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/predictor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db31f5b2be04f0852b497825356ee3917c11d424 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/__pycache__/predictor.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/automatic_mask_generator.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/automatic_mask_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..23264971b7ff5aa0b4f499ade7773b68dce984b6 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/automatic_mask_generator.py @@ -0,0 +1,372 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torchvision.ops.boxes import batched_nms, box_area # type: ignore + +from typing import Any, Dict, List, Optional, Tuple + +from .modeling import Sam +from .predictor import SamPredictor +from .utils.amg import ( + MaskData, + area_from_rle, + batch_iterator, + batched_mask_to_box, + box_xyxy_to_xywh, + build_all_layer_point_grids, + calculate_stability_score, + coco_encode_rle, + generate_crop_boxes, + is_box_near_crop_edge, + mask_to_rle_pytorch, + remove_small_regions, + rle_to_mask, + uncrop_boxes_xyxy, + uncrop_masks, + uncrop_points, +) + + +class SamAutomaticMaskGenerator: + def __init__( + self, + model: Sam, + points_per_side: Optional[int] = 32, + points_per_batch: int = 64, + pred_iou_thresh: float = 0.88, + stability_score_thresh: float = 0.95, + stability_score_offset: float = 1.0, + box_nms_thresh: float = 0.7, + crop_n_layers: int = 0, + crop_nms_thresh: float = 0.7, + crop_overlap_ratio: float = 512 / 1500, + crop_n_points_downscale_factor: int = 1, + point_grids: Optional[List[np.ndarray]] = None, + min_mask_region_area: int = 0, + output_mode: str = "binary_mask", + ) -> None: + """ + Using a SAM model, generates masks for the entire image. + Generates a grid of point prompts over the image, then filters + low quality and duplicate masks. The default settings are chosen + for SAM with a ViT-H backbone. + + Arguments: + model (Sam): The SAM model to use for mask prediction. + points_per_side (int or None): The number of points to be sampled + along one side of the image. The total number of points is + points_per_side**2. If None, 'point_grids' must provide explicit + point sampling. + points_per_batch (int): Sets the number of points run simultaneously + by the model. Higher numbers may be faster but use more GPU memory. + pred_iou_thresh (float): A filtering threshold in [0,1], using the + model's predicted mask quality. + stability_score_thresh (float): A filtering threshold in [0,1], using + the stability of the mask under changes to the cutoff used to binarize + the model's mask predictions. + stability_score_offset (float): The amount to shift the cutoff when + calculated the stability score. + box_nms_thresh (float): The box IoU cutoff used by non-maximal + suppression to filter duplicate masks. + crops_n_layers (int): If >0, mask prediction will be run again on + crops of the image. Sets the number of layers to run, where each + layer has 2**i_layer number of image crops. + crops_nms_thresh (float): The box IoU cutoff used by non-maximal + suppression to filter duplicate masks between different crops. + crop_overlap_ratio (float): Sets the degree to which crops overlap. + In the first crop layer, crops will overlap by this fraction of + the image length. Later layers with more crops scale down this overlap. + crop_n_points_downscale_factor (int): The number of points-per-side + sampled in layer n is scaled down by crop_n_points_downscale_factor**n. + point_grids (list(np.ndarray) or None): A list over explicit grids + of points used for sampling, normalized to [0,1]. The nth grid in the + list is used in the nth crop layer. Exclusive with points_per_side. + min_mask_region_area (int): If >0, postprocessing will be applied + to remove disconnected regions and holes in masks with area smaller + than min_mask_region_area. Requires opencv. + output_mode (str): The form masks are returned in. Can be 'binary_mask', + 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools. + For large resolutions, 'binary_mask' may consume large amounts of + memory. + """ + + assert (points_per_side is None) != ( + point_grids is None + ), "Exactly one of points_per_side or point_grid must be provided." + if points_per_side is not None: + self.point_grids = build_all_layer_point_grids( + points_per_side, + crop_n_layers, + crop_n_points_downscale_factor, + ) + elif point_grids is not None: + self.point_grids = point_grids + else: + raise ValueError("Can't have both points_per_side and point_grid be None.") + + assert output_mode in [ + "binary_mask", + "uncompressed_rle", + "coco_rle", + ], f"Unknown output_mode {output_mode}." + if output_mode == "coco_rle": + from pycocotools import mask as mask_utils # type: ignore # noqa: F401 + + if min_mask_region_area > 0: + import cv2 # type: ignore # noqa: F401 + + self.predictor = SamPredictor(model) + self.points_per_batch = points_per_batch + self.pred_iou_thresh = pred_iou_thresh + self.stability_score_thresh = stability_score_thresh + self.stability_score_offset = stability_score_offset + self.box_nms_thresh = box_nms_thresh + self.crop_n_layers = crop_n_layers + self.crop_nms_thresh = crop_nms_thresh + self.crop_overlap_ratio = crop_overlap_ratio + self.crop_n_points_downscale_factor = crop_n_points_downscale_factor + self.min_mask_region_area = min_mask_region_area + self.output_mode = output_mode + + @torch.no_grad() + def generate(self, image: np.ndarray) -> List[Dict[str, Any]]: + """ + Generates masks for the given image. + + Arguments: + image (np.ndarray): The image to generate masks for, in HWC uint8 format. + + Returns: + list(dict(str, any)): A list over records for masks. Each record is + a dict containing the following keys: + segmentation (dict(str, any) or np.ndarray): The mask. If + output_mode='binary_mask', is an array of shape HW. Otherwise, + is a dictionary containing the RLE. + bbox (list(float)): The box around the mask, in XYWH format. + area (int): The area in pixels of the mask. + predicted_iou (float): The model's own prediction of the mask's + quality. This is filtered by the pred_iou_thresh parameter. + point_coords (list(list(float))): The point coordinates input + to the model to generate this mask. + stability_score (float): A measure of the mask's quality. This + is filtered on using the stability_score_thresh parameter. + crop_box (list(float)): The crop of the image used to generate + the mask, given in XYWH format. + """ + + # Generate masks + mask_data = self._generate_masks(image) + + # Filter small disconnected regions and holes in masks + if self.min_mask_region_area > 0: + mask_data = self.postprocess_small_regions( + mask_data, + self.min_mask_region_area, + max(self.box_nms_thresh, self.crop_nms_thresh), + ) + + # Encode masks + if self.output_mode == "coco_rle": + mask_data["segmentations"] = [coco_encode_rle(rle) for rle in mask_data["rles"]] + elif self.output_mode == "binary_mask": + mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]] + else: + mask_data["segmentations"] = mask_data["rles"] + + # Write mask records + curr_anns = [] + for idx in range(len(mask_data["segmentations"])): + ann = { + "segmentation": mask_data["segmentations"][idx], + "area": area_from_rle(mask_data["rles"][idx]), + "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(), + "predicted_iou": mask_data["iou_preds"][idx].item(), + "point_coords": [mask_data["points"][idx].tolist()], + "stability_score": mask_data["stability_score"][idx].item(), + "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(), + } + curr_anns.append(ann) + + return curr_anns + + def _generate_masks(self, image: np.ndarray) -> MaskData: + orig_size = image.shape[:2] + crop_boxes, layer_idxs = generate_crop_boxes( + orig_size, self.crop_n_layers, self.crop_overlap_ratio + ) + + # Iterate over image crops + data = MaskData() + for crop_box, layer_idx in zip(crop_boxes, layer_idxs): + crop_data = self._process_crop(image, crop_box, layer_idx, orig_size) + data.cat(crop_data) + + # Remove duplicate masks between crops + if len(crop_boxes) > 1: + # Prefer masks from smaller crops + scores = 1 / box_area(data["crop_boxes"]) + scores = scores.to(data["boxes"].device) + keep_by_nms = batched_nms( + data["boxes"].float(), + scores, + torch.zeros(len(data["boxes"])), # categories + iou_threshold=self.crop_nms_thresh, + ) + data.filter(keep_by_nms) + + data.to_numpy() + return data + + def _process_crop( + self, + image: np.ndarray, + crop_box: List[int], + crop_layer_idx: int, + orig_size: Tuple[int, ...], + ) -> MaskData: + # Crop the image and calculate embeddings + x0, y0, x1, y1 = crop_box + cropped_im = image[y0:y1, x0:x1, :] + cropped_im_size = cropped_im.shape[:2] + self.predictor.set_image(cropped_im) + + # Get points for this crop + points_scale = np.array(cropped_im_size)[None, ::-1] + points_for_image = self.point_grids[crop_layer_idx] * points_scale + + # Generate masks for this crop in batches + data = MaskData() + for (points,) in batch_iterator(self.points_per_batch, points_for_image): + batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size) + data.cat(batch_data) + del batch_data + self.predictor.reset_image() + + # Remove duplicates within this crop. + keep_by_nms = batched_nms( + data["boxes"].float(), + data["iou_preds"], + torch.zeros(len(data["boxes"])), # categories + iou_threshold=self.box_nms_thresh, + ) + data.filter(keep_by_nms) + + # Return to the original image frame + data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box) + data["points"] = uncrop_points(data["points"], crop_box) + data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))]) + + return data + + def _process_batch( + self, + points: np.ndarray, + im_size: Tuple[int, ...], + crop_box: List[int], + orig_size: Tuple[int, ...], + ) -> MaskData: + orig_h, orig_w = orig_size + + # Run model on this batch + transformed_points = self.predictor.transform.apply_coords(points, im_size) + in_points = torch.as_tensor(transformed_points, device=self.predictor.device) + in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device) + masks, iou_preds, _ = self.predictor.predict_torch( + in_points[:, None, :], + in_labels[:, None], + multimask_output=True, + return_logits=True, + ) + + # Serialize predictions and store in MaskData + data = MaskData( + masks=masks.flatten(0, 1), + iou_preds=iou_preds.flatten(0, 1), + points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)), + ) + del masks + + # Filter by predicted IoU + if self.pred_iou_thresh > 0.0: + keep_mask = data["iou_preds"] > self.pred_iou_thresh + data.filter(keep_mask) + + # Calculate stability score + data["stability_score"] = calculate_stability_score( + data["masks"], self.predictor.model.mask_threshold, self.stability_score_offset + ) + if self.stability_score_thresh > 0.0: + keep_mask = data["stability_score"] >= self.stability_score_thresh + data.filter(keep_mask) + + # Threshold masks and calculate boxes + data["masks"] = data["masks"] > self.predictor.model.mask_threshold + data["boxes"] = batched_mask_to_box(data["masks"]) + + # Filter boxes that touch crop boundaries + keep_mask = ~is_box_near_crop_edge(data["boxes"], crop_box, [0, 0, orig_w, orig_h]) + if not torch.all(keep_mask): + data.filter(keep_mask) + + # Compress to RLE + data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w) + data["rles"] = mask_to_rle_pytorch(data["masks"]) + del data["masks"] + + return data + + @staticmethod + def postprocess_small_regions( + mask_data: MaskData, min_area: int, nms_thresh: float + ) -> MaskData: + """ + Removes small disconnected regions and holes in masks, then reruns + box NMS to remove any new duplicates. + + Edits mask_data in place. + + Requires open-cv as a dependency. + """ + if len(mask_data["rles"]) == 0: + return mask_data + + # Filter small disconnected regions and holes + new_masks = [] + scores = [] + for rle in mask_data["rles"]: + mask = rle_to_mask(rle) + + mask, changed = remove_small_regions(mask, min_area, mode="holes") + unchanged = not changed + mask, changed = remove_small_regions(mask, min_area, mode="islands") + unchanged = unchanged and not changed + + new_masks.append(torch.as_tensor(mask).unsqueeze(0)) + # Give score=0 to changed masks and score=1 to unchanged masks + # so NMS will prefer ones that didn't need postprocessing + scores.append(float(unchanged)) + + # Recalculate boxes and remove any new duplicates + masks = torch.cat(new_masks, dim=0) + boxes = batched_mask_to_box(masks) + keep_by_nms = batched_nms( + boxes.float(), + torch.as_tensor(scores), + torch.zeros(len(boxes)), # categories + iou_threshold=nms_thresh, + ) + + # Only recalculate RLEs for masks that have changed + for i_mask in keep_by_nms: + if scores[i_mask] == 0.0: + mask_torch = masks[i_mask].unsqueeze(0) + mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0] + mask_data["boxes"][i_mask] = boxes[i_mask] # update res directly + mask_data.filter(keep_by_nms) + + return mask_data diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/build_sam.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/build_sam.py new file mode 100644 index 0000000000000000000000000000000000000000..07abfca24e96eced7f13bdefd3212ce1b77b8999 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/build_sam.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from functools import partial + +from .modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer + + +def build_sam_vit_h(checkpoint=None): + return _build_sam( + encoder_embed_dim=1280, + encoder_depth=32, + encoder_num_heads=16, + encoder_global_attn_indexes=[7, 15, 23, 31], + checkpoint=checkpoint, + ) + + +build_sam = build_sam_vit_h + + +def build_sam_vit_l(checkpoint=None): + return _build_sam( + encoder_embed_dim=1024, + encoder_depth=24, + encoder_num_heads=16, + encoder_global_attn_indexes=[5, 11, 17, 23], + checkpoint=checkpoint, + ) + + +def build_sam_vit_b(checkpoint=None): + return _build_sam( + encoder_embed_dim=768, + encoder_depth=12, + encoder_num_heads=12, + encoder_global_attn_indexes=[2, 5, 8, 11], + checkpoint=checkpoint, + ) + + +sam_model_registry = { + "default": build_sam, + "vit_h": build_sam, + "vit_l": build_sam_vit_l, + "vit_b": build_sam_vit_b, +} + + +def _build_sam( + encoder_embed_dim, + encoder_depth, + encoder_num_heads, + encoder_global_attn_indexes, + checkpoint=None, +): + prompt_embed_dim = 256 + image_size = 1024 + vit_patch_size = 16 + image_embedding_size = image_size // vit_patch_size + sam = Sam( + image_encoder=ImageEncoderViT( + depth=encoder_depth, + embed_dim=encoder_embed_dim, + img_size=image_size, + mlp_ratio=4, + norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), + num_heads=encoder_num_heads, + patch_size=vit_patch_size, + qkv_bias=True, + use_rel_pos=True, + global_attn_indexes=encoder_global_attn_indexes, + window_size=14, + out_chans=prompt_embed_dim, + ), + prompt_encoder=PromptEncoder( + embed_dim=prompt_embed_dim, + image_embedding_size=(image_embedding_size, image_embedding_size), + input_image_size=(image_size, image_size), + mask_in_chans=16, + ), + mask_decoder=MaskDecoder( + num_multimask_outputs=3, + transformer=TwoWayTransformer( + depth=2, + embedding_dim=prompt_embed_dim, + mlp_dim=2048, + num_heads=8, + ), + transformer_dim=prompt_embed_dim, + iou_head_depth=3, + iou_head_hidden_dim=256, + ), + pixel_mean=[123.675, 116.28, 103.53], + pixel_std=[58.395, 57.12, 57.375], + ) + sam.eval() + if checkpoint is not None: + with open(checkpoint, "rb") as f: + state_dict = torch.load(f) + sam.load_state_dict(state_dict) + return sam diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/build_sam_hq.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/build_sam_hq.py new file mode 100644 index 0000000000000000000000000000000000000000..a113b745c9772cb2f5a34a81c0626e9161699796 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/build_sam_hq.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from functools import partial + +from .modeling import ImageEncoderViT, MaskDecoderHQ, PromptEncoder, Sam, TwoWayTransformer + + +def build_sam_hq_vit_h(checkpoint=None): + return _build_sam( + encoder_embed_dim=1280, + encoder_depth=32, + encoder_num_heads=16, + encoder_global_attn_indexes=[7, 15, 23, 31], + checkpoint=checkpoint, + ) + + +build_sam_hq = build_sam_hq_vit_h + + +def build_sam_hq_vit_l(checkpoint=None): + return _build_sam( + encoder_embed_dim=1024, + encoder_depth=24, + encoder_num_heads=16, + encoder_global_attn_indexes=[5, 11, 17, 23], + checkpoint=checkpoint, + ) + + +def build_sam_hq_vit_b(checkpoint=None): + return _build_sam( + encoder_embed_dim=768, + encoder_depth=12, + encoder_num_heads=12, + encoder_global_attn_indexes=[2, 5, 8, 11], + checkpoint=checkpoint, + ) + + +sam_hq_model_registry = { + "default": build_sam_hq_vit_h, + "vit_h": build_sam_hq_vit_h, + "vit_l": build_sam_hq_vit_l, + "vit_b": build_sam_hq_vit_b, +} + + +def _build_sam( + encoder_embed_dim, + encoder_depth, + encoder_num_heads, + encoder_global_attn_indexes, + checkpoint=None, +): + prompt_embed_dim = 256 + image_size = 1024 + vit_patch_size = 16 + image_embedding_size = image_size // vit_patch_size + sam = Sam( + image_encoder=ImageEncoderViT( + depth=encoder_depth, + embed_dim=encoder_embed_dim, + img_size=image_size, + mlp_ratio=4, + norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), + num_heads=encoder_num_heads, + patch_size=vit_patch_size, + qkv_bias=True, + use_rel_pos=True, + global_attn_indexes=encoder_global_attn_indexes, + window_size=14, + out_chans=prompt_embed_dim, + ), + prompt_encoder=PromptEncoder( + embed_dim=prompt_embed_dim, + image_embedding_size=(image_embedding_size, image_embedding_size), + input_image_size=(image_size, image_size), + mask_in_chans=16, + ), + mask_decoder=MaskDecoderHQ( + num_multimask_outputs=3, + transformer=TwoWayTransformer( + depth=2, + embedding_dim=prompt_embed_dim, + mlp_dim=2048, + num_heads=8, + ), + transformer_dim=prompt_embed_dim, + iou_head_depth=3, + iou_head_hidden_dim=256, + vit_dim=encoder_embed_dim, + ), + pixel_mean=[123.675, 116.28, 103.53], + pixel_std=[58.395, 57.12, 57.375], + ) + # sam.eval() + if checkpoint is not None: + with open(checkpoint, "rb") as f: + device = "cuda" if torch.cuda.is_available() else "cpu" + state_dict = torch.load(f, map_location=device) + info = sam.load_state_dict(state_dict, strict=False) + print(info) + for n, p in sam.named_parameters(): + if 'hf_token' not in n and 'hf_mlp' not in n and 'compress_vit_feat' not in n and 'embedding_encoder' not in n and 'embedding_maskfeature' not in n: + p.requires_grad = False + + return sam diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__init__.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71172d22345eff1f3729c6326299feee17717ccc --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from .sam import Sam +from .image_encoder import ImageEncoderViT +from .mask_decoder_hq import MaskDecoderHQ +from .mask_decoder import MaskDecoder +from .prompt_encoder import PromptEncoder +from .transformer import TwoWayTransformer diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87f1d8c2ac93e4982fabb8c8e6920f5cbc1359d8 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/common.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/common.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf2be6697355c877d54c9d107c6659efcfc4a84e Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/common.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..692059b3bca1c5c92ecc3f0c5d1eef99c9730abd Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6f379c2d75c665dc0576d4d8aecbf9a0cea1e3d Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder_hq.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder_hq.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..870de23dc2f41fa1c47fa40aed9ae92925620bfb Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder_hq.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22a6147114c235c720c5d5030c6ac80f7a1ed4f2 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/sam.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/sam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a224c8aa2b4fb41f9b049e4b4064762d6caccb2 Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/sam.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cba0adf9585bf375a4f187bd232e7e8727ea5fff Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/common.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/common.py new file mode 100644 index 0000000000000000000000000000000000000000..2bf15236a3eb24d8526073bc4fa2b274cccb3f96 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/common.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn + +from typing import Type + + +class MLPBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + mlp_dim: int, + act: Type[nn.Module] = nn.GELU, + ) -> None: + super().__init__() + self.lin1 = nn.Linear(embedding_dim, mlp_dim) + self.lin2 = nn.Linear(mlp_dim, embedding_dim) + self.act = act() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lin2(self.act(self.lin1(x))) + + +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa +# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa +class LayerNorm2d(nn.Module): + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/image_encoder.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/image_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..9c01f71c906e90d9b9b7d3252cdd3e5c555a2734 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/image_encoder.py @@ -0,0 +1,398 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from typing import Optional, Tuple, Type + +from .common import LayerNorm2d, MLPBlock + + +# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa +class ImageEncoderViT(nn.Module): + def __init__( + self, + img_size: int = 1024, + patch_size: int = 16, + in_chans: int = 3, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + out_chans: int = 256, + qkv_bias: bool = True, + norm_layer: Type[nn.Module] = nn.LayerNorm, + act_layer: Type[nn.Module] = nn.GELU, + use_abs_pos: bool = True, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + window_size: int = 0, + global_attn_indexes: Tuple[int, ...] = (), + ) -> None: + """ + Args: + img_size (int): Input image size. + patch_size (int): Patch size. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + depth (int): Depth of ViT. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_abs_pos (bool): If True, use absolute positional embeddings. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. + global_attn_indexes (list): Indexes for blocks using global attention. + """ + super().__init__() + self.img_size = img_size + + self.patch_embed = PatchEmbed( + kernel_size=(patch_size, patch_size), + stride=(patch_size, patch_size), + in_chans=in_chans, + embed_dim=embed_dim, + ) + + self.pos_embed: Optional[nn.Parameter] = None + if use_abs_pos: + # Initialize absolute positional embedding with pretrain image size. + self.pos_embed = nn.Parameter( + torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim) + ) + + self.blocks = nn.ModuleList() + for i in range(depth): + block = Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + act_layer=act_layer, + use_rel_pos=use_rel_pos, + rel_pos_zero_init=rel_pos_zero_init, + window_size=window_size if i not in global_attn_indexes else 0, + input_size=(img_size // patch_size, img_size // patch_size), + ) + self.blocks.append(block) + + self.neck = nn.Sequential( + nn.Conv2d( + embed_dim, + out_chans, + kernel_size=1, + bias=False, + ), + LayerNorm2d(out_chans), + nn.Conv2d( + out_chans, + out_chans, + kernel_size=3, + padding=1, + bias=False, + ), + LayerNorm2d(out_chans), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.patch_embed(x) + if self.pos_embed is not None: + x = x + self.pos_embed + + interm_embeddings=[] + for blk in self.blocks: + x = blk(x) + if blk.window_size == 0: + interm_embeddings.append(x) + + x = self.neck(x.permute(0, 3, 1, 2)) + + return x, interm_embeddings + + +class Block(nn.Module): + """Transformer blocks with support of window attention and residual propagation blocks""" + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + norm_layer: Type[nn.Module] = nn.LayerNorm, + act_layer: Type[nn.Module] = nn.GELU, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + window_size: int = 0, + input_size: Optional[Tuple[int, int]] = None, + ) -> None: + """ + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. If it equals 0, then + use global attention. + input_size (tuple(int, int) or None): Input resolution for calculating the relative + positional parameter size. + """ + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + use_rel_pos=use_rel_pos, + rel_pos_zero_init=rel_pos_zero_init, + input_size=input_size if window_size == 0 else (window_size, window_size), + ) + + self.norm2 = norm_layer(dim) + self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer) + + self.window_size = window_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = x + x = self.norm1(x) + # Window partition + if self.window_size > 0: + H, W = x.shape[1], x.shape[2] + x, pad_hw = window_partition(x, self.window_size) + + x = self.attn(x) + # Reverse window partition + if self.window_size > 0: + x = window_unpartition(x, self.window_size, pad_hw, (H, W)) + + x = shortcut + x + x = x + self.mlp(self.norm2(x)) + + return x + + +class Attention(nn.Module): + """Multi-head Attention block with relative position embeddings.""" + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + input_size: Optional[Tuple[int, int]] = None, + ) -> None: + """ + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + input_size (tuple(int, int) or None): Input resolution for calculating the relative + positional parameter size. + """ + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.proj = nn.Linear(dim, dim) + + self.use_rel_pos = use_rel_pos + if self.use_rel_pos: + assert ( + input_size is not None + ), "Input size must be provided if using relative positional encoding." + # initialize relative positional embeddings + self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) + self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, H, W, _ = x.shape + # qkv with shape (3, B, nHead, H * W, C) + qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) + # q, k, v with shape (B * nHead, H * W, C) + q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0) + + attn = (q * self.scale) @ k.transpose(-2, -1) + + if self.use_rel_pos: + attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)) + + attn = attn.softmax(dim=-1) + x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1) + x = self.proj(x) + + return x + + +def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]: + """ + Partition into non-overlapping windows with padding if needed. + Args: + x (tensor): input tokens with [B, H, W, C]. + window_size (int): window size. + + Returns: + windows: windows after partition with [B * num_windows, window_size, window_size, C]. + (Hp, Wp): padded height and width before partition + """ + B, H, W, C = x.shape + + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + if pad_h > 0 or pad_w > 0: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + Hp, Wp = H + pad_h, W + pad_w + + x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows, (Hp, Wp) + + +def window_unpartition( + windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int] +) -> torch.Tensor: + """ + Window unpartition into original sequences and removing padding. + Args: + windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. + window_size (int): window size. + pad_hw (Tuple): padded height and width (Hp, Wp). + hw (Tuple): original height and width (H, W) before padding. + + Returns: + x: unpartitioned sequences with [B, H, W, C]. + """ + Hp, Wp = pad_hw + H, W = hw + B = windows.shape[0] // (Hp * Wp // window_size // window_size) + x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) + + if Hp > H or Wp > W: + x = x[:, :H, :W, :].contiguous() + return x + + +def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: + """ + Get relative positional embeddings according to the relative positions of + query and key sizes. + Args: + q_size (int): size of query q. + k_size (int): size of key k. + rel_pos (Tensor): relative position embeddings (L, C). + + Returns: + Extracted positional embeddings according to relative positions. + """ + max_rel_dist = int(2 * max(q_size, k_size) - 1) + # Interpolate rel pos if needed. + if rel_pos.shape[0] != max_rel_dist: + # Interpolate rel pos. + rel_pos_resized = F.interpolate( + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), + size=max_rel_dist, + mode="linear", + ) + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) + else: + rel_pos_resized = rel_pos + + # Scale the coords with short length if shapes for q and k are different. + q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) + k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + + return rel_pos_resized[relative_coords.long()] + + +def add_decomposed_rel_pos( + attn: torch.Tensor, + q: torch.Tensor, + rel_pos_h: torch.Tensor, + rel_pos_w: torch.Tensor, + q_size: Tuple[int, int], + k_size: Tuple[int, int], +) -> torch.Tensor: + """ + Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. + https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950 + Args: + attn (Tensor): attention map. + q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C). + rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis. + rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis. + q_size (Tuple): spatial sequence size of query q with (q_h, q_w). + k_size (Tuple): spatial sequence size of key k with (k_h, k_w). + + Returns: + attn (Tensor): attention map with added relative positional embeddings. + """ + q_h, q_w = q_size + k_h, k_w = k_size + Rh = get_rel_pos(q_h, k_h, rel_pos_h) + Rw = get_rel_pos(q_w, k_w, rel_pos_w) + + B, _, dim = q.shape + r_q = q.reshape(B, q_h, q_w, dim) + rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) + rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) + + attn = ( + attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] + ).view(B, q_h * q_w, k_h * k_w) + + return attn + + +class PatchEmbed(nn.Module): + """ + Image to Patch Embedding. + """ + + def __init__( + self, + kernel_size: Tuple[int, int] = (16, 16), + stride: Tuple[int, int] = (16, 16), + padding: Tuple[int, int] = (0, 0), + in_chans: int = 3, + embed_dim: int = 768, + ) -> None: + """ + Args: + kernel_size (Tuple): kernel size of the projection layer. + stride (Tuple): stride of the projection layer. + padding (Tuple): padding size of the projection layer. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + """ + super().__init__() + + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + # B C H W -> B H W C + x = x.permute(0, 2, 3, 1) + return x \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/mask_decoder.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/mask_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c36c7b553c9df986dab91474de06d171feb7f93d --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/mask_decoder.py @@ -0,0 +1,178 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import nn +from torch.nn import functional as F + +from typing import List, Tuple, Type + +from .common import LayerNorm2d + + +class MaskDecoder(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), + activation(), + ) + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth + ) + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + hq_token_only: bool, + interm_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + """ + masks, iou_pred = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + ) + + # Select the correct mask or masks for output + if multimask_output: + mask_slice = slice(1, None) + else: + mask_slice = slice(0, 1) + masks = masks[:, mask_slice, :, :] + iou_pred = iou_pred[:, mask_slice] + + # Prepare output + return masks, iou_pred + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0) + output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + src = src + dense_prompt_embeddings + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, 0, :] + mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + upscaled_embedding = self.output_upscaling(src) + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])) + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding.shape + masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) + + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + + return masks, iou_pred + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/mask_decoder_hq.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/mask_decoder_hq.py new file mode 100644 index 0000000000000000000000000000000000000000..c4576f3495ae72d639b2278c4c252e3e02e5d424 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/mask_decoder_hq.py @@ -0,0 +1,232 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Modified by HQ-SAM team +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import nn +from torch.nn import functional as F + +from typing import List, Tuple, Type + +from .common import LayerNorm2d + + +class MaskDecoderHQ(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + vit_dim: int = 1024, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), + activation(), + ) + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth + ) + + # HQ-SAM parameters + self.hf_token = nn.Embedding(1, transformer_dim) # HQ-Ouptput-Token + self.hf_mlp = MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) # corresponding new MLP layer for HQ-Ouptput-Token + self.num_mask_tokens = self.num_mask_tokens + 1 + + # three conv fusion layers for obtaining HQ-Feature + self.compress_vit_feat = nn.Sequential( + nn.ConvTranspose2d(vit_dim, transformer_dim, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim), + nn.GELU(), + nn.ConvTranspose2d(transformer_dim, transformer_dim // 8, kernel_size=2, stride=2)) + + self.embedding_encoder = nn.Sequential( + nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim // 4), + nn.GELU(), + nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), + ) + self.embedding_maskfeature = nn.Sequential( + nn.Conv2d(transformer_dim // 8, transformer_dim // 4, 3, 1, 1), + LayerNorm2d(transformer_dim // 4), + nn.GELU(), + nn.Conv2d(transformer_dim // 4, transformer_dim // 8, 3, 1, 1)) + + + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + hq_token_only: bool, + interm_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the ViT image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + """ + vit_features = interm_embeddings[0].permute(0, 3, 1, 2) # early-layer ViT feature, after 1st global attention block in ViT + hq_features = self.embedding_encoder(image_embeddings) + self.compress_vit_feat(vit_features) + + masks, iou_pred = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + hq_features=hq_features, + ) + + # Select the correct mask or masks for output + if multimask_output: + # mask with highest score + mask_slice = slice(1,self.num_mask_tokens-1) + iou_pred = iou_pred[:, mask_slice] + iou_pred, max_iou_idx = torch.max(iou_pred,dim=1) + iou_pred = iou_pred.unsqueeze(1) + masks_multi = masks[:, mask_slice, :, :] + masks_sam = masks_multi[torch.arange(masks_multi.size(0)),max_iou_idx].unsqueeze(1) + else: + # singale mask output, default + mask_slice = slice(0, 1) + iou_pred = iou_pred[:,mask_slice] + masks_sam = masks[:,mask_slice] + + masks_hq = masks[:,slice(self.num_mask_tokens-1, self.num_mask_tokens)] + if hq_token_only: + masks = masks_hq + else: + masks = masks_sam + masks_hq + # Prepare output + return masks, iou_pred + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + hq_features: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight, self.hf_token.weight], dim=0) + output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + src = src + dense_prompt_embeddings + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, 0, :] + mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + + upscaled_embedding_sam = self.output_upscaling(src) + upscaled_embedding_hq = self.embedding_maskfeature(upscaled_embedding_sam) + hq_features.repeat(b,1,1,1) + + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + if i < self.num_mask_tokens - 1: + hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])) + else: + hyper_in_list.append(self.hf_mlp(mask_tokens_out[:, i, :])) + + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding_sam.shape + + masks_sam = (hyper_in[:,:self.num_mask_tokens-1] @ upscaled_embedding_sam.view(b, c, h * w)).view(b, -1, h, w) + masks_sam_hq = (hyper_in[:,self.num_mask_tokens-1:] @ upscaled_embedding_hq.view(b, c, h * w)).view(b, -1, h, w) + masks = torch.cat([masks_sam,masks_sam_hq],dim=1) + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + + return masks, iou_pred + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/prompt_encoder.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/prompt_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c3143f4f8e02ddd7ca8587b40ff5d47c3a6b7ef3 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/prompt_encoder.py @@ -0,0 +1,214 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torch import nn + +from typing import Any, Optional, Tuple, Type + +from .common import LayerNorm2d + + +class PromptEncoder(nn.Module): + def __init__( + self, + embed_dim: int, + image_embedding_size: Tuple[int, int], + input_image_size: Tuple[int, int], + mask_in_chans: int, + activation: Type[nn.Module] = nn.GELU, + ) -> None: + """ + Encodes prompts for input to SAM's mask decoder. + + Arguments: + embed_dim (int): The prompts' embedding dimension + image_embedding_size (tuple(int, int)): The spatial size of the + image embedding, as (H, W). + input_image_size (int): The padded size of the image as input + to the image encoder, as (H, W). + mask_in_chans (int): The number of hidden channels used for + encoding input masks. + activation (nn.Module): The activation to use when encoding + input masks. + """ + super().__init__() + self.embed_dim = embed_dim + self.input_image_size = input_image_size + self.image_embedding_size = image_embedding_size + self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) + + self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners + point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)] + self.point_embeddings = nn.ModuleList(point_embeddings) + self.not_a_point_embed = nn.Embedding(1, embed_dim) + + self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1]) + self.mask_downscaling = nn.Sequential( + nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans // 4), + activation(), + nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans), + activation(), + nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), + ) + self.no_mask_embed = nn.Embedding(1, embed_dim) + + def get_dense_pe(self) -> torch.Tensor: + """ + Returns the positional encoding used to encode point prompts, + applied to a dense set of points the shape of the image encoding. + + Returns: + torch.Tensor: Positional encoding with shape + 1x(embed_dim)x(embedding_h)x(embedding_w) + """ + return self.pe_layer(self.image_embedding_size).unsqueeze(0) + + def _embed_points( + self, + points: torch.Tensor, + labels: torch.Tensor, + pad: bool, + ) -> torch.Tensor: + """Embeds point prompts.""" + points = points + 0.5 # Shift to center of pixel + if pad: + padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device) + padding_label = -torch.ones((labels.shape[0], 1), device=labels.device) + points = torch.cat([points, padding_point], dim=1) + labels = torch.cat([labels, padding_label], dim=1) + point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size) + point_embedding[labels == -1] = 0.0 + point_embedding[labels == -1] += self.not_a_point_embed.weight + point_embedding[labels == 0] += self.point_embeddings[0].weight + point_embedding[labels == 1] += self.point_embeddings[1].weight + return point_embedding + + def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: + """Embeds box prompts.""" + boxes = boxes + 0.5 # Shift to center of pixel + coords = boxes.reshape(-1, 2, 2) + corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size) + corner_embedding[:, 0, :] += self.point_embeddings[2].weight + corner_embedding[:, 1, :] += self.point_embeddings[3].weight + return corner_embedding + + def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor: + """Embeds mask inputs.""" + mask_embedding = self.mask_downscaling(masks) + return mask_embedding + + def _get_batch_size( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> int: + """ + Gets the batch size of the output given the batch size of the input prompts. + """ + if points is not None: + return points[0].shape[0] + elif boxes is not None: + return boxes.shape[0] + elif masks is not None: + return masks.shape[0] + else: + return 1 + + def _get_device(self) -> torch.device: + return self.point_embeddings[0].weight.device + + def forward( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Embeds different types of prompts, returning both sparse and dense + embeddings. + + Arguments: + points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates + and labels to embed. + boxes (torch.Tensor or none): boxes to embed + masks (torch.Tensor or none): masks to embed + + Returns: + torch.Tensor: sparse embeddings for the points and boxes, with shape + BxNx(embed_dim), where N is determined by the number of input points + and boxes. + torch.Tensor: dense embeddings for the masks, in the shape + Bx(embed_dim)x(embed_H)x(embed_W) + """ + bs = self._get_batch_size(points, boxes, masks) + sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device()) + if points is not None: + coords, labels = points + point_embeddings = self._embed_points(coords, labels, pad=(boxes is None)) + sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1) + if boxes is not None: + box_embeddings = self._embed_boxes(boxes) + sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1) + + if masks is not None: + dense_embeddings = self._embed_masks(masks) + else: + dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( + bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] + ) + + return sparse_embeddings, dense_embeddings + + +class PositionEmbeddingRandom(nn.Module): + """ + Positional encoding using random spatial frequencies. + """ + + def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: + super().__init__() + if scale is None or scale <= 0.0: + scale = 1.0 + self.register_buffer( + "positional_encoding_gaussian_matrix", + scale * torch.randn((2, num_pos_feats)), + ) + + def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: + """Positionally encode points that are normalized to [0,1].""" + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coords = 2 * coords - 1 + coords = coords @ self.positional_encoding_gaussian_matrix + coords = 2 * np.pi * coords + # outputs d_1 x ... x d_n x C shape + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + def forward(self, size: Tuple[int, int]) -> torch.Tensor: + """Generate positional encoding for a grid of the specified size.""" + h, w = size + device: Any = self.positional_encoding_gaussian_matrix.device + grid = torch.ones((h, w), device=device, dtype=torch.float32) + y_embed = grid.cumsum(dim=0) - 0.5 + x_embed = grid.cumsum(dim=1) - 0.5 + y_embed = y_embed / h + x_embed = x_embed / w + + pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) + return pe.permute(2, 0, 1) # C x H x W + + def forward_with_coords( + self, coords_input: torch.Tensor, image_size: Tuple[int, int] + ) -> torch.Tensor: + """Positionally encode points that are not normalized to [0,1].""" + coords = coords_input.clone() + coords[:, :, 0] = coords[:, :, 0] / image_size[1] + coords[:, :, 1] = coords[:, :, 1] / image_size[0] + return self._pe_encoding(coords.to(torch.float)) # B x N x C diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/sam.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/sam.py new file mode 100644 index 0000000000000000000000000000000000000000..303bc2f40c3dbc84f5d4286bb73336e075a86589 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/sam.py @@ -0,0 +1,174 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import nn +from torch.nn import functional as F + +from typing import Any, Dict, List, Tuple + +from .image_encoder import ImageEncoderViT +from .mask_decoder import MaskDecoder +from .prompt_encoder import PromptEncoder + + +class Sam(nn.Module): + mask_threshold: float = 0.0 + image_format: str = "RGB" + + def __init__( + self, + image_encoder: ImageEncoderViT, + prompt_encoder: PromptEncoder, + mask_decoder: MaskDecoder, + pixel_mean: List[float] = [123.675, 116.28, 103.53], + pixel_std: List[float] = [58.395, 57.12, 57.375], + ) -> None: + """ + SAM predicts object masks from an image and input prompts. + + Arguments: + image_encoder (ImageEncoderViT): The backbone used to encode the + image into image embeddings that allow for efficient mask prediction. + prompt_encoder (PromptEncoder): Encodes various types of input prompts. + mask_decoder (MaskDecoder): Predicts masks from the image embeddings + and encoded prompts. + pixel_mean (list(float)): Mean values for normalizing pixels in the input image. + pixel_std (list(float)): Std values for normalizing pixels in the input image. + """ + super().__init__() + self.image_encoder = image_encoder + self.prompt_encoder = prompt_encoder + self.mask_decoder = mask_decoder + self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False) + self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False) + + @property + def device(self) -> Any: + return self.pixel_mean.device + + @torch.no_grad() + def forward( + self, + batched_input: List[Dict[str, Any]], + multimask_output: bool, + ) -> List[Dict[str, torch.Tensor]]: + """ + Predicts masks end-to-end from provided images and prompts. + If prompts are not known in advance, using SamPredictor is + recommended over calling the model directly. + + Arguments: + batched_input (list(dict)): A list over input images, each a + dictionary with the following keys. A prompt key can be + excluded if it is not present. + 'image': The image as a torch tensor in 3xHxW format, + already transformed for input to the model. + 'original_size': (tuple(int, int)) The original size of + the image before transformation, as (H, W). + 'point_coords': (torch.Tensor) Batched point prompts for + this image, with shape BxNx2. Already transformed to the + input frame of the model. + 'point_labels': (torch.Tensor) Batched labels for point prompts, + with shape BxN. + 'boxes': (torch.Tensor) Batched box inputs, with shape Bx4. + Already transformed to the input frame of the model. + 'mask_inputs': (torch.Tensor) Batched mask inputs to the model, + in the form Bx1xHxW. + multimask_output (bool): Whether the model should predict multiple + disambiguating masks, or return a single mask. + + Returns: + (list(dict)): A list over input images, where each element is + as dictionary with the following keys. + 'masks': (torch.Tensor) Batched binary mask predictions, + with shape BxCxHxW, where B is the number of input promts, + C is determiend by multimask_output, and (H, W) is the + original size of the image. + 'iou_predictions': (torch.Tensor) The model's predictions + of mask quality, in shape BxC. + 'low_res_logits': (torch.Tensor) Low resolution logits with + shape BxCxHxW, where H=W=256. Can be passed as mask input + to subsequent iterations of prediction. + """ + input_images = torch.stack([self.preprocess(x["image"]) for x in batched_input], dim=0) + image_embeddings = self.image_encoder(input_images) + + outputs = [] + for image_record, curr_embedding in zip(batched_input, image_embeddings): + if "point_coords" in image_record: + points = (image_record["point_coords"], image_record["point_labels"]) + else: + points = None + sparse_embeddings, dense_embeddings = self.prompt_encoder( + points=points, + boxes=image_record.get("boxes", None), + masks=image_record.get("mask_inputs", None), + ) + low_res_masks, iou_predictions = self.mask_decoder( + image_embeddings=curr_embedding.unsqueeze(0), + image_pe=self.prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + ) + masks = self.postprocess_masks( + low_res_masks, + input_size=image_record["image"].shape[-2:], + original_size=image_record["original_size"], + ) + masks = masks > self.mask_threshold + outputs.append( + { + "masks": masks, + "iou_predictions": iou_predictions, + "low_res_logits": low_res_masks, + } + ) + return outputs + + def postprocess_masks( + self, + masks: torch.Tensor, + input_size: Tuple[int, ...], + original_size: Tuple[int, ...], + ) -> torch.Tensor: + """ + Remove padding and upscale masks to the original image size. + + Arguments: + masks (torch.Tensor): Batched masks from the mask_decoder, + in BxCxHxW format. + input_size (tuple(int, int)): The size of the image input to the + model, in (H, W) format. Used to remove padding. + original_size (tuple(int, int)): The original size of the image + before resizing for input to the model, in (H, W) format. + + Returns: + (torch.Tensor): Batched masks in BxCxHxW format, where (H, W) + is given by original_size. + """ + masks = F.interpolate( + masks, + (self.image_encoder.img_size, self.image_encoder.img_size), + mode="bilinear", + align_corners=False, + ) + masks = masks[..., : input_size[0], : input_size[1]] + masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False) + return masks + + def preprocess(self, x: torch.Tensor) -> torch.Tensor: + """Normalize pixel values and pad to a square input.""" + # Normalize colors + x = (x - self.pixel_mean) / self.pixel_std + + # Pad + h, w = x.shape[-2:] + padh = self.image_encoder.img_size - h + padw = self.image_encoder.img_size - w + x = F.pad(x, (0, padw, 0, padh)) + return x diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/transformer.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f1a2812f613cc55b1d0b3e3e1d0c84a760d1fb87 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/modeling/transformer.py @@ -0,0 +1,240 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor, nn + +import math +from typing import Tuple, Type + +from .common import MLPBlock + + +class TwoWayTransformer(nn.Module): + def __init__( + self, + depth: int, + embedding_dim: int, + num_heads: int, + mlp_dim: int, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + ) -> None: + """ + A transformer decoder that attends to an input image using + queries whose positional embedding is supplied. + + Args: + depth (int): number of layers in the transformer + embedding_dim (int): the channel dimension for the input embeddings + num_heads (int): the number of heads for multihead attention. Must + divide embedding_dim + mlp_dim (int): the channel dimension internal to the MLP block + activation (nn.Module): the activation to use in the MLP block + """ + super().__init__() + self.depth = depth + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.layers = nn.ModuleList() + + for i in range(depth): + self.layers.append( + TwoWayAttentionBlock( + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_dim=mlp_dim, + activation=activation, + attention_downsample_rate=attention_downsample_rate, + skip_first_layer_pe=(i == 0), + ) + ) + + self.final_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm_final_attn = nn.LayerNorm(embedding_dim) + + def forward( + self, + image_embedding: Tensor, + image_pe: Tensor, + point_embedding: Tensor, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + image_embedding (torch.Tensor): image to attend to. Should be shape + B x embedding_dim x h x w for any h and w. + image_pe (torch.Tensor): the positional encoding to add to the image. Must + have the same shape as image_embedding. + point_embedding (torch.Tensor): the embedding to add to the query points. + Must have shape B x N_points x embedding_dim for any N_points. + + Returns: + torch.Tensor: the processed point_embedding + torch.Tensor: the processed image_embedding + """ + # BxCxHxW -> BxHWxC == B x N_image_tokens x C + bs, c, h, w = image_embedding.shape + image_embedding = image_embedding.flatten(2).permute(0, 2, 1) + image_pe = image_pe.flatten(2).permute(0, 2, 1) + + # Prepare queries + queries = point_embedding + keys = image_embedding + + # Apply transformer blocks and final layernorm + for layer in self.layers: + queries, keys = layer( + queries=queries, + keys=keys, + query_pe=point_embedding, + key_pe=image_pe, + ) + + # Apply the final attenion layer from the points to the image + q = queries + point_embedding + k = keys + image_pe + attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm_final_attn(queries) + + return queries, keys + + +class TwoWayAttentionBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + num_heads: int, + mlp_dim: int = 2048, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + skip_first_layer_pe: bool = False, + ) -> None: + """ + A transformer block with four layers: (1) self-attention of sparse + inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp + block on sparse inputs, and (4) cross attention of dense inputs to sparse + inputs. + + Arguments: + embedding_dim (int): the channel dimension of the embeddings + num_heads (int): the number of heads in the attention layers + mlp_dim (int): the hidden dimension of the mlp block + activation (nn.Module): the activation of the mlp block + skip_first_layer_pe (bool): skip the PE on the first layer + """ + super().__init__() + self.self_attn = Attention(embedding_dim, num_heads) + self.norm1 = nn.LayerNorm(embedding_dim) + + self.cross_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm2 = nn.LayerNorm(embedding_dim) + + self.mlp = MLPBlock(embedding_dim, mlp_dim, activation) + self.norm3 = nn.LayerNorm(embedding_dim) + + self.norm4 = nn.LayerNorm(embedding_dim) + self.cross_attn_image_to_token = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + + self.skip_first_layer_pe = skip_first_layer_pe + + def forward( + self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor + ) -> Tuple[Tensor, Tensor]: + # Self attention block + if self.skip_first_layer_pe: + queries = self.self_attn(q=queries, k=queries, v=queries) + else: + q = queries + query_pe + attn_out = self.self_attn(q=q, k=q, v=queries) + queries = queries + attn_out + queries = self.norm1(queries) + + # Cross attention block, tokens attending to image embedding + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm2(queries) + + # MLP block + mlp_out = self.mlp(queries) + queries = queries + mlp_out + queries = self.norm3(queries) + + # Cross attention block, image embedding attending to tokens + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) + keys = keys + attn_out + keys = self.norm4(keys) + + return queries, keys + + +class Attention(nn.Module): + """ + An attention layer that allows for downscaling the size of the embedding + after projection to queries, keys, and values. + """ + + def __init__( + self, + embedding_dim: int, + num_heads: int, + downsample_rate: int = 1, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.internal_dim = embedding_dim // downsample_rate + self.num_heads = num_heads + assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim." + + self.q_proj = nn.Linear(embedding_dim, self.internal_dim) + self.k_proj = nn.Linear(embedding_dim, self.internal_dim) + self.v_proj = nn.Linear(embedding_dim, self.internal_dim) + self.out_proj = nn.Linear(self.internal_dim, embedding_dim) + + def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: + b, n, c = x.shape + x = x.reshape(b, n, num_heads, c // num_heads) + return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head + + def _recombine_heads(self, x: Tensor) -> Tensor: + b, n_heads, n_tokens, c_per_head = x.shape + x = x.transpose(1, 2) + return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + # Attention + _, _, _, c_per_head = q.shape + attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens + attn = attn / math.sqrt(c_per_head) + attn = torch.softmax(attn, dim=-1) + + # Get output + out = attn @ v + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/predictor.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..882063e14c9eb346018cf7e0a668fce251fa18f6 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/predictor.py @@ -0,0 +1,276 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch + +from .modeling import Sam + +from typing import Optional, Tuple + +from .utils.transforms import ResizeLongestSide + + +class SamPredictor: + def __init__( + self, + sam_model: Sam, + ) -> None: + """ + Uses SAM to calculate the image embedding for an image, and then + allow repeated, efficient mask prediction given prompts. + + Arguments: + sam_model (Sam): The model to use for mask prediction. + """ + super().__init__() + self.model = sam_model + self.transform = ResizeLongestSide(sam_model.image_encoder.img_size) + self.reset_image() + + def set_image( + self, + image: np.ndarray, + image_format: str = "RGB", + ) -> None: + """ + Calculates the image embeddings for the provided image, allowing + masks to be predicted with the 'predict' method. + + Arguments: + image (np.ndarray): The image for calculating masks. Expects an + image in HWC uint8 format, with pixel values in [0, 255]. + image_format (str): The color format of the image, in ['RGB', 'BGR']. + """ + assert image_format in [ + "RGB", + "BGR", + ], f"image_format must be in ['RGB', 'BGR'], is {image_format}." + # import pdb;pdb.set_trace() + if image_format != self.model.image_format: + image = image[..., ::-1] + + # Transform the image to the form expected by the model + # import pdb;pdb.set_trace() + input_image = self.transform.apply_image(image) + input_image_torch = torch.as_tensor(input_image, device=self.device) + input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :] + + self.set_torch_image(input_image_torch, image.shape[:2]) + + @torch.no_grad() + def set_torch_image( + self, + transformed_image: torch.Tensor, + original_image_size: Tuple[int, ...], + ) -> None: + """ + Calculates the image embeddings for the provided image, allowing + masks to be predicted with the 'predict' method. Expects the input + image to be already transformed to the format expected by the model. + + Arguments: + transformed_image (torch.Tensor): The input image, with shape + 1x3xHxW, which has been transformed with ResizeLongestSide. + original_image_size (tuple(int, int)): The size of the image + before transformation, in (H, W) format. + """ + assert ( + len(transformed_image.shape) == 4 + and transformed_image.shape[1] == 3 + and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size + ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}." + self.reset_image() + + self.original_size = original_image_size + self.input_size = tuple(transformed_image.shape[-2:]) + input_image = self.model.preprocess(transformed_image) + self.features, self.interm_features = self.model.image_encoder(input_image) + self.is_image_set = True + + def predict( + self, + point_coords: Optional[np.ndarray] = None, + point_labels: Optional[np.ndarray] = None, + box: Optional[np.ndarray] = None, + mask_input: Optional[np.ndarray] = None, + multimask_output: bool = True, + return_logits: bool = False, + hq_token_only: bool =False, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Predict masks for the given input prompts, using the currently set image. + + Arguments: + point_coords (np.ndarray or None): A Nx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (np.ndarray or None): A length N array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + box (np.ndarray or None): A length 4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form 1xHxW, where + for SAM, H=W=256. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + + Returns: + (np.ndarray): The output masks in CxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (np.ndarray): An array of length C containing the model's + predictions for the quality of each mask. + (np.ndarray): An array of shape CxHxW, where C is the number + of masks and H=W=256. These low resolution logits can be passed to + a subsequent iteration as mask input. + """ + if not self.is_image_set: + raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") + + # Transform input prompts + coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None + if point_coords is not None: + assert ( + point_labels is not None + ), "point_labels must be supplied if point_coords is supplied." + point_coords = self.transform.apply_coords(point_coords, self.original_size) + coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device) + labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device) + coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :] + if box is not None: + box = self.transform.apply_boxes(box, self.original_size) + box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device) + box_torch = box_torch[None, :] + if mask_input is not None: + mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.device) + mask_input_torch = mask_input_torch[None, :, :, :] + + masks, iou_predictions, low_res_masks = self.predict_torch( + coords_torch, + labels_torch, + box_torch, + mask_input_torch, + multimask_output, + return_logits=return_logits, + hq_token_only=hq_token_only, + ) + + masks_np = masks[0].detach().cpu().numpy() + iou_predictions_np = iou_predictions[0].detach().cpu().numpy() + low_res_masks_np = low_res_masks[0].detach().cpu().numpy() + return masks_np, iou_predictions_np, low_res_masks_np + + @torch.no_grad() + def predict_torch( + self, + point_coords: Optional[torch.Tensor], + point_labels: Optional[torch.Tensor], + boxes: Optional[torch.Tensor] = None, + mask_input: Optional[torch.Tensor] = None, + multimask_output: bool = True, + return_logits: bool = False, + hq_token_only: bool =False, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Predict masks for the given input prompts, using the currently set image. + Input prompts are batched torch tensors and are expected to already be + transformed to the input frame using ResizeLongestSide. + + Arguments: + point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (torch.Tensor or None): A BxN array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + boxes (np.ndarray or None): A Bx4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form Bx1xHxW, where + for SAM, H=W=256. Masks returned by a previous iteration of the + predict method do not need further transformation. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + + Returns: + (torch.Tensor): The output masks in BxCxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (torch.Tensor): An array of shape BxC containing the model's + predictions for the quality of each mask. + (torch.Tensor): An array of shape BxCxHxW, where C is the number + of masks and H=W=256. These low res logits can be passed to + a subsequent iteration as mask input. + """ + if not self.is_image_set: + raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") + + if point_coords is not None: + points = (point_coords, point_labels) + else: + points = None + + # Embed prompts + sparse_embeddings, dense_embeddings = self.model.prompt_encoder( + points=points, + boxes=boxes, + masks=mask_input, + ) + + # Predict masks + low_res_masks, iou_predictions = self.model.mask_decoder( + image_embeddings=self.features, + image_pe=self.model.prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + hq_token_only=hq_token_only, + interm_embeddings=self.interm_features, + ) + + # Upscale the masks to the original image resolution + masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size) + + if not return_logits: + masks = masks > self.model.mask_threshold + + return masks, iou_predictions, low_res_masks + + def get_image_embedding(self) -> torch.Tensor: + """ + Returns the image embeddings for the currently set image, with + shape 1xCxHxW, where C is the embedding dimension and (H,W) are + the embedding spatial dimension of SAM (typically C=256, H=W=64). + """ + if not self.is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) to generate an embedding." + ) + assert self.features is not None, "Features must exist if an image has been set." + return self.features + + @property + def device(self) -> torch.device: + return self.model.device + + def reset_image(self) -> None: + """Resets the currently set image.""" + self.is_image_set = False + self.features = None + self.orig_h = None + self.orig_w = None + self.input_h = None + self.input_w = None \ No newline at end of file diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__init__.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5277f46157403e47fd830fc519144b97ef69d4ae --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0efd72923b79790d6fb0b7bcea89a4ca0b16f00a Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/amg.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/amg.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a78cb163484c2b015feb56b61fd8bb605170cb8e Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/amg.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/transforms.cpython-310.pyc b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/transforms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b698b4f636f09cafb4c30d46be407acca8270be Binary files /dev/null and b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/__pycache__/transforms.cpython-310.pyc differ diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/amg.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/amg.py new file mode 100644 index 0000000000000000000000000000000000000000..3a137778e45c464c079658ecb87ec53270e789f7 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/amg.py @@ -0,0 +1,346 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch + +import math +from copy import deepcopy +from itertools import product +from typing import Any, Dict, Generator, ItemsView, List, Tuple + + +class MaskData: + """ + A structure for storing masks and their related data in batched format. + Implements basic filtering and concatenation. + """ + + def __init__(self, **kwargs) -> None: + for v in kwargs.values(): + assert isinstance( + v, (list, np.ndarray, torch.Tensor) + ), "MaskData only supports list, numpy arrays, and torch tensors." + self._stats = dict(**kwargs) + + def __setitem__(self, key: str, item: Any) -> None: + assert isinstance( + item, (list, np.ndarray, torch.Tensor) + ), "MaskData only supports list, numpy arrays, and torch tensors." + self._stats[key] = item + + def __delitem__(self, key: str) -> None: + del self._stats[key] + + def __getitem__(self, key: str) -> Any: + return self._stats[key] + + def items(self) -> ItemsView[str, Any]: + return self._stats.items() + + def filter(self, keep: torch.Tensor) -> None: + for k, v in self._stats.items(): + if v is None: + self._stats[k] = None + elif isinstance(v, torch.Tensor): + self._stats[k] = v[torch.as_tensor(keep, device=v.device)] + elif isinstance(v, np.ndarray): + self._stats[k] = v[keep.detach().cpu().numpy()] + elif isinstance(v, list) and keep.dtype == torch.bool: + self._stats[k] = [a for i, a in enumerate(v) if keep[i]] + elif isinstance(v, list): + self._stats[k] = [v[i] for i in keep] + else: + raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") + + def cat(self, new_stats: "MaskData") -> None: + for k, v in new_stats.items(): + if k not in self._stats or self._stats[k] is None: + self._stats[k] = deepcopy(v) + elif isinstance(v, torch.Tensor): + self._stats[k] = torch.cat([self._stats[k], v], dim=0) + elif isinstance(v, np.ndarray): + self._stats[k] = np.concatenate([self._stats[k], v], axis=0) + elif isinstance(v, list): + self._stats[k] = self._stats[k] + deepcopy(v) + else: + raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") + + def to_numpy(self) -> None: + for k, v in self._stats.items(): + if isinstance(v, torch.Tensor): + self._stats[k] = v.detach().cpu().numpy() + + +def is_box_near_crop_edge( + boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0 +) -> torch.Tensor: + """Filter masks at the edge of a crop, but not at the edge of the original image.""" + crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device) + orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device) + boxes = uncrop_boxes_xyxy(boxes, crop_box).float() + near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0) + near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0) + near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge) + return torch.any(near_crop_edge, dim=1) + + +def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor: + box_xywh = deepcopy(box_xyxy) + box_xywh[2] = box_xywh[2] - box_xywh[0] + box_xywh[3] = box_xywh[3] - box_xywh[1] + return box_xywh + + +def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]: + assert len(args) > 0 and all( + len(a) == len(args[0]) for a in args + ), "Batched iteration must have inputs of all the same size." + n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0) + for b in range(n_batches): + yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args] + + +def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]: + """ + Encodes masks to an uncompressed RLE, in the format expected by + pycoco tools. + """ + # Put in fortran order and flatten h,w + b, h, w = tensor.shape + tensor = tensor.permute(0, 2, 1).flatten(1) + + # Compute change indices + diff = tensor[:, 1:] ^ tensor[:, :-1] + change_indices = diff.nonzero() + + # Encode run length + out = [] + for i in range(b): + cur_idxs = change_indices[change_indices[:, 0] == i, 1] + cur_idxs = torch.cat( + [ + torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device), + cur_idxs + 1, + torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device), + ] + ) + btw_idxs = cur_idxs[1:] - cur_idxs[:-1] + counts = [] if tensor[i, 0] == 0 else [0] + counts.extend(btw_idxs.detach().cpu().tolist()) + out.append({"size": [h, w], "counts": counts}) + return out + + +def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray: + """Compute a binary mask from an uncompressed RLE.""" + h, w = rle["size"] + mask = np.empty(h * w, dtype=bool) + idx = 0 + parity = False + for count in rle["counts"]: + mask[idx : idx + count] = parity + idx += count + parity ^= True + mask = mask.reshape(w, h) + return mask.transpose() # Put in C order + + +def area_from_rle(rle: Dict[str, Any]) -> int: + return sum(rle["counts"][1::2]) + + +def calculate_stability_score( + masks: torch.Tensor, mask_threshold: float, threshold_offset: float +) -> torch.Tensor: + """ + Computes the stability score for a batch of masks. The stability + score is the IoU between the binary masks obtained by thresholding + the predicted mask logits at high and low values. + """ + # One mask is always contained inside the other. + # Save memory by preventing unnecesary cast to torch.int64 + intersections = ( + (masks > (mask_threshold + threshold_offset)) + .sum(-1, dtype=torch.int16) + .sum(-1, dtype=torch.int32) + ) + unions = ( + (masks > (mask_threshold - threshold_offset)) + .sum(-1, dtype=torch.int16) + .sum(-1, dtype=torch.int32) + ) + return intersections / unions + + +def build_point_grid(n_per_side: int) -> np.ndarray: + """Generates a 2D grid of points evenly spaced in [0,1]x[0,1].""" + offset = 1 / (2 * n_per_side) + points_one_side = np.linspace(offset, 1 - offset, n_per_side) + points_x = np.tile(points_one_side[None, :], (n_per_side, 1)) + points_y = np.tile(points_one_side[:, None], (1, n_per_side)) + points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2) + return points + + +def build_all_layer_point_grids( + n_per_side: int, n_layers: int, scale_per_layer: int +) -> List[np.ndarray]: + """Generates point grids for all crop layers.""" + points_by_layer = [] + for i in range(n_layers + 1): + n_points = int(n_per_side / (scale_per_layer**i)) + points_by_layer.append(build_point_grid(n_points)) + return points_by_layer + + +def generate_crop_boxes( + im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float +) -> Tuple[List[List[int]], List[int]]: + """ + Generates a list of crop boxes of different sizes. Each layer + has (2**i)**2 boxes for the ith layer. + """ + crop_boxes, layer_idxs = [], [] + im_h, im_w = im_size + short_side = min(im_h, im_w) + + # Original image + crop_boxes.append([0, 0, im_w, im_h]) + layer_idxs.append(0) + + def crop_len(orig_len, n_crops, overlap): + return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops)) + + for i_layer in range(n_layers): + n_crops_per_side = 2 ** (i_layer + 1) + overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side)) + + crop_w = crop_len(im_w, n_crops_per_side, overlap) + crop_h = crop_len(im_h, n_crops_per_side, overlap) + + crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)] + crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)] + + # Crops in XYWH format + for x0, y0 in product(crop_box_x0, crop_box_y0): + box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)] + crop_boxes.append(box) + layer_idxs.append(i_layer + 1) + + return crop_boxes, layer_idxs + + +def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor: + x0, y0, _, _ = crop_box + offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device) + # Check if boxes has a channel dimension + if len(boxes.shape) == 3: + offset = offset.unsqueeze(1) + return boxes + offset + + +def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor: + x0, y0, _, _ = crop_box + offset = torch.tensor([[x0, y0]], device=points.device) + # Check if points has a channel dimension + if len(points.shape) == 3: + offset = offset.unsqueeze(1) + return points + offset + + +def uncrop_masks( + masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int +) -> torch.Tensor: + x0, y0, x1, y1 = crop_box + if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h: + return masks + # Coordinate transform masks + pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0) + pad = (x0, pad_x - x0, y0, pad_y - y0) + return torch.nn.functional.pad(masks, pad, value=0) + + +def remove_small_regions( + mask: np.ndarray, area_thresh: float, mode: str +) -> Tuple[np.ndarray, bool]: + """ + Removes small disconnected regions and holes in a mask. Returns the + mask and an indicator of if the mask has been modified. + """ + import cv2 # type: ignore + + assert mode in ["holes", "islands"] + correct_holes = mode == "holes" + working_mask = (correct_holes ^ mask).astype(np.uint8) + n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8) + sizes = stats[:, -1][1:] # Row 0 is background label + small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh] + if len(small_regions) == 0: + return mask, False + fill_labels = [0] + small_regions + if not correct_holes: + fill_labels = [i for i in range(n_labels) if i not in fill_labels] + # If every region is below threshold, keep largest + if len(fill_labels) == 0: + fill_labels = [int(np.argmax(sizes)) + 1] + mask = np.isin(regions, fill_labels) + return mask, True + + +def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]: + from pycocotools import mask as mask_utils # type: ignore + + h, w = uncompressed_rle["size"] + rle = mask_utils.frPyObjects(uncompressed_rle, h, w) + rle["counts"] = rle["counts"].decode("utf-8") # Necessary to serialize with json + return rle + + +def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor: + """ + Calculates boxes in XYXY format around masks. Return [0,0,0,0] for + an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4. + """ + # torch.max below raises an error on empty inputs, just skip in this case + if torch.numel(masks) == 0: + return torch.zeros(*masks.shape[:-2], 4, device=masks.device) + + # Normalize shape to CxHxW + shape = masks.shape + h, w = shape[-2:] + if len(shape) > 2: + masks = masks.flatten(0, -3) + else: + masks = masks.unsqueeze(0) + + # Get top and bottom edges + in_height, _ = torch.max(masks, dim=-1) + in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :] + bottom_edges, _ = torch.max(in_height_coords, dim=-1) + in_height_coords = in_height_coords + h * (~in_height) + top_edges, _ = torch.min(in_height_coords, dim=-1) + + # Get left and right edges + in_width, _ = torch.max(masks, dim=-2) + in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :] + right_edges, _ = torch.max(in_width_coords, dim=-1) + in_width_coords = in_width_coords + w * (~in_width) + left_edges, _ = torch.min(in_width_coords, dim=-1) + + # If the mask is empty the right edge will be to the left of the left edge. + # Replace these boxes with [0, 0, 0, 0] + empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges) + out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1) + out = out * (~empty_filter).unsqueeze(-1) + + # Return to original shape + if len(shape) > 2: + out = out.reshape(*shape[:-2], 4) + else: + out = out[0] + + return out diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/onnx.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/onnx.py new file mode 100644 index 0000000000000000000000000000000000000000..4297b31291e036700d6ad0b818afb7dd72da3054 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/onnx.py @@ -0,0 +1,144 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +from torch.nn import functional as F + +from typing import Tuple + +from ..modeling import Sam +from .amg import calculate_stability_score + + +class SamOnnxModel(nn.Module): + """ + This model should not be called directly, but is used in ONNX export. + It combines the prompt encoder, mask decoder, and mask postprocessing of Sam, + with some functions modified to enable model tracing. Also supports extra + options controlling what information. See the ONNX export script for details. + """ + + def __init__( + self, + model: Sam, + return_single_mask: bool, + use_stability_score: bool = False, + return_extra_metrics: bool = False, + ) -> None: + super().__init__() + self.mask_decoder = model.mask_decoder + self.model = model + self.img_size = model.image_encoder.img_size + self.return_single_mask = return_single_mask + self.use_stability_score = use_stability_score + self.stability_score_offset = 1.0 + self.return_extra_metrics = return_extra_metrics + + @staticmethod + def resize_longest_image_size( + input_image_size: torch.Tensor, longest_side: int + ) -> torch.Tensor: + input_image_size = input_image_size.to(torch.float32) + scale = longest_side / torch.max(input_image_size) + transformed_size = scale * input_image_size + transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64) + return transformed_size + + def _embed_points(self, point_coords: torch.Tensor, point_labels: torch.Tensor) -> torch.Tensor: + point_coords = point_coords + 0.5 + point_coords = point_coords / self.img_size + point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords) + point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding) + + point_embedding = point_embedding * (point_labels != -1) + point_embedding = point_embedding + self.model.prompt_encoder.not_a_point_embed.weight * ( + point_labels == -1 + ) + + for i in range(self.model.prompt_encoder.num_point_embeddings): + point_embedding = point_embedding + self.model.prompt_encoder.point_embeddings[ + i + ].weight * (point_labels == i) + + return point_embedding + + def _embed_masks(self, input_mask: torch.Tensor, has_mask_input: torch.Tensor) -> torch.Tensor: + mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(input_mask) + mask_embedding = mask_embedding + ( + 1 - has_mask_input + ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1) + return mask_embedding + + def mask_postprocessing(self, masks: torch.Tensor, orig_im_size: torch.Tensor) -> torch.Tensor: + masks = F.interpolate( + masks, + size=(self.img_size, self.img_size), + mode="bilinear", + align_corners=False, + ) + + prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size) + masks = masks[..., : int(prepadded_size[0]), : int(prepadded_size[1])] + + orig_im_size = orig_im_size.to(torch.int64) + h, w = orig_im_size[0], orig_im_size[1] + masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False) + return masks + + def select_masks( + self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Determine if we should return the multiclick mask or not from the number of points. + # The reweighting is used to avoid control flow. + score_reweight = torch.tensor( + [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)] + ).to(iou_preds.device) + score = iou_preds + (num_points - 2.5) * score_reweight + best_idx = torch.argmax(score, dim=1) + masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1) + iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1) + + return masks, iou_preds + + @torch.no_grad() + def forward( + self, + image_embeddings: torch.Tensor, + point_coords: torch.Tensor, + point_labels: torch.Tensor, + mask_input: torch.Tensor, + has_mask_input: torch.Tensor, + orig_im_size: torch.Tensor, + ): + sparse_embedding = self._embed_points(point_coords, point_labels) + dense_embedding = self._embed_masks(mask_input, has_mask_input) + + masks, scores = self.model.mask_decoder.predict_masks( + image_embeddings=image_embeddings, + image_pe=self.model.prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embedding, + dense_prompt_embeddings=dense_embedding, + ) + + if self.use_stability_score: + scores = calculate_stability_score( + masks, self.model.mask_threshold, self.stability_score_offset + ) + + if self.return_single_mask: + masks, scores = self.select_masks(masks, scores, point_coords.shape[1]) + + upscaled_masks = self.mask_postprocessing(masks, orig_im_size) + + if self.return_extra_metrics: + stability_scores = calculate_stability_score( + upscaled_masks, self.model.mask_threshold, self.stability_score_offset + ) + areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1) + return upscaled_masks, scores, stability_scores, areas, masks + + return upscaled_masks, scores, masks diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/transforms.py b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..3ad346661f84b0647026e130a552c4b38b83e2ac --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/segment_anything/utils/transforms.py @@ -0,0 +1,102 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torch.nn import functional as F +from torchvision.transforms.functional import resize, to_pil_image # type: ignore + +from copy import deepcopy +from typing import Tuple + + +class ResizeLongestSide: + """ + Resizes images to longest side 'target_length', as well as provides + methods for resizing coordinates and boxes. Provides methods for + transforming both numpy array and batched torch tensors. + """ + + def __init__(self, target_length: int) -> None: + self.target_length = target_length + + def apply_image(self, image: np.ndarray) -> np.ndarray: + """ + Expects a numpy array with shape HxWxC in uint8 format. + """ + target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length) + return np.array(resize(to_pil_image(image), target_size)) + + def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray: + """ + Expects a numpy array of length 2 in the final dimension. Requires the + original image size in (H, W) format. + """ + old_h, old_w = original_size + new_h, new_w = self.get_preprocess_shape( + original_size[0], original_size[1], self.target_length + ) + coords = deepcopy(coords).astype(float) + coords[..., 0] = coords[..., 0] * (new_w / old_w) + coords[..., 1] = coords[..., 1] * (new_h / old_h) + return coords + + def apply_boxes(self, boxes: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray: + """ + Expects a numpy array shape Bx4. Requires the original image size + in (H, W) format. + """ + boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size) + return boxes.reshape(-1, 4) + + def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor: + """ + Expects batched images with shape BxCxHxW and float format. This + transformation may not exactly match apply_image. apply_image is + the transformation expected by the model. + """ + # Expects an image in BCHW format. May not exactly match apply_image. + target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length) + return F.interpolate( + image, target_size, mode="bilinear", align_corners=False, antialias=True + ) + + def apply_coords_torch( + self, coords: torch.Tensor, original_size: Tuple[int, ...] + ) -> torch.Tensor: + """ + Expects a torch tensor with length 2 in the last dimension. Requires the + original image size in (H, W) format. + """ + old_h, old_w = original_size + new_h, new_w = self.get_preprocess_shape( + original_size[0], original_size[1], self.target_length + ) + coords = deepcopy(coords).to(torch.float) + coords[..., 0] = coords[..., 0] * (new_w / old_w) + coords[..., 1] = coords[..., 1] * (new_h / old_h) + return coords + + def apply_boxes_torch( + self, boxes: torch.Tensor, original_size: Tuple[int, ...] + ) -> torch.Tensor: + """ + Expects a torch tensor with shape Bx4. Requires the original image + size in (H, W) format. + """ + boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size) + return boxes.reshape(-1, 4) + + @staticmethod + def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]: + """ + Compute the output size given input size and target long side length. + """ + scale = long_side_length * 1.0 / max(oldh, oldw) + newh, neww = oldh * scale, oldw * scale + neww = int(neww + 0.5) + newh = int(newh + 0.5) + return (newh, neww) diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/setup.cfg b/ArtiAgent - DefectDiffu/src/segment_anything/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..0eee130ba71d14ec260d33a8ebd96a6491079a54 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/setup.cfg @@ -0,0 +1,11 @@ +[isort] +line_length=100 +multi_line_output=3 +include_trailing_comma=True +known_standard_library=numpy,setuptools +skip_glob=*/__init__.py +known_myself=segment_anything +known_third_party=matplotlib,cv2,torch,torchvision,pycocotools,onnx,black,isort +no_lines_before=STDLIB,THIRDPARTY +sections=FUTURE,STDLIB,THIRDPARTY,MYSELF,FIRSTPARTY,LOCALFOLDER +default_section=FIRSTPARTY diff --git a/ArtiAgent - DefectDiffu/src/segment_anything/setup.py b/ArtiAgent - DefectDiffu/src/segment_anything/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0986317eb576a14ec774205c88fdee3cc6c0b3 --- /dev/null +++ b/ArtiAgent - DefectDiffu/src/segment_anything/setup.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from setuptools import find_packages, setup + +setup( + name="segment_anything", + version="1.0", + install_requires=[], + packages=find_packages(exclude="notebooks"), + extras_require={ + "all": ["matplotlib", "pycocotools", "opencv-python", "onnx", "onnxruntime"], + "dev": ["flake8", "isort", "black", "mypy"], + }, +) diff --git a/ArtiAgent - DefectFill/engine/DefectFill/README.md b/ArtiAgent - DefectFill/engine/DefectFill/README.md new file mode 100644 index 0000000000000000000000000000000000000000..fd13c1a3b514a3af274e3301fd31ffad6a640621 --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/README.md @@ -0,0 +1,158 @@ +# DefectFill: Realistic Defect Generation for Visual Inspection + +Realistic defect image generation via fine-tuned inpainting diffusion models. + +> Implementation of **DefectFill: Realistic Defect Generation with Inpainting Diffusion Model for Visual Inspection** (CVPR 2024). + +--- + +Currently, this repository is tuned to generate **cracks in concrete** (using the MVTec AD dataset) as a proof-of-concept. The ultimate goal of this project is to apply these techniques to generate synthetic training data for **cast iron defects** (e.g., blowholes, cracks) in foundry settings. + +## Visual Results + +Below are generated examples showing the model's ability to fill healthy regions with realistic defect textures while preserving the surrounding structural integrity. + +||| +| :---: | :---: | +| ![Result 1](triplet_results/triplet_0000.png) | ![Result 2](triplet_results/triplet_0001.png) | +| ![Result 3](triplet_results/triplet_0002.png) | ![Result 4](triplet_results/triplet_0003.png) | +| ![Result 5](triplet_results/triplet_0004.png) | ![Result 6](triplet_results/triplet_0005.png) | + +--- + +## Overview + +DefectFill fine-tunes a Stable Diffusion 2 inpainting model with LoRA to learn a specific defect concept from a small set of reference images. Three complementary loss terms drive training: + +| Loss | Weight | Purpose | +|------|--------|---------| +| **Defect loss** `L_def` | 0.5 | Precisely captures intrinsic defect features | +| **Object loss** `L_obj` | 0.2 | Learns the semantic relationship between defect and object | +| **Attention loss** `L_attn` | 0.05 | Ensures [V*] token attends to the defect region | + +After training, **Low-Fidelity Selection (LFS)** generates 8 candidates per (image, mask) pair and selects the one with the highest LPIPS score inside the masked region โ€” the most "realistic" defect. + +--- + +## Installation + +1. Clone the repository: + ```bash + git clone [https://github.com/axelsig1/defectfill.git](https://github.com/axelsig1/defectfill.git) + cd defectfill + ``` + +2. Install the required dependencies: + ```bash + pip install -r requirements.txt + ``` + +**Requirements include:** `torch`, `diffusers`, `transformers`, `peft`, `lpips`, and `albumentations`. + +--- + +## Data Preparation + +This project follows the **MVTec AD** dataset structure. Ensure your data is organized as follows: +``` +data/ +โ””โ”€โ”€ concrete/ # Object Class + โ”œโ”€โ”€ train/ + โ”‚ โ”œโ”€โ”€ defective/ + โ”‚ โ”‚ โ””โ”€โ”€ crack/ # Defect images + โ”‚ โ””โ”€โ”€ defective_masks/ + โ”‚ โ””โ”€โ”€ crack/ # Corresponding binary masks + โ””โ”€โ”€ test/ + โ””โ”€โ”€ good/ # Healthy reference images +``` +## Usage + +### 1. Training + +To train the model on concrete cracks: + +```bash +python train.py \ + --data_dir ./data \ + --object_class concrete \ + --defect_type crack \ + --output_dir ./output_concrete \ + --lora_rank 8 \ + --lora_alpha 16 \ + --max_train_steps 2000 +``` + +Key training details (from paper): +- **Base model**: `sd2-community/stable-diffusion-2-inpainting` +- **LoRA** on UNet attention layers + text encoder projection matrices +- **Warmup**: linear 0 โ†’ LR over first 100 steps +- **Augmentation**: random resize ร—[1.0, 1.125] + random crop +- **Random masks** M_rand: 30 boxes, sides 3โ€“25% of image size +- **[V*] token**: the word `sks` + +### 2. Inference +Generate new synthetic defects on healthy images. The script uses LPIPS to pick the best generation from a batch of candidates. +```bash +python inference.py \ + --checkpoint ./output_concrete/checkpoints/checkpoint_final.pt \ + --object_class concrete \ + --defect_type crack \ + --data_dir ./data \ + --output_dir ./generated_cracks \ + --total_images 6 \ + --num_samples 8 \ + --guidance_scale 2.0 +``` + +--- + +## Method Details + +### Defect Loss (Eq. 5) + +``` +L_def = E[ || M โŠ™ (ฮต โˆ’ ฮต_ฮธ(x_t^def, t, c^def)) ||ยฒ ] +``` + +Background image: `B_def = (1 โˆ’ M) โŠ™ I` +Input: `x_t^def = concat(x_t, b_def, M)` +Prompt `P_def = "A photo of sks"` + +### Object Loss (Eq. 7) + +``` +L_obj = E[ || M' โŠ™ (ฮต โˆ’ ฮต_ฮธ(x_t^obj, t, c^obj)) ||ยฒ ] +M' = M + ฮฑยท(1 โˆ’ M), ฮฑ = 0.3 +``` + +Random box mask M_rand (30 boxes), `B_rand = (1 โˆ’ M_rand) โŠ™ I` +Input: `x_t^obj = concat(x_t, b_rand, M_rand)` +Prompt `P_obj = "A with sks"` + +### Attention Loss (Eq. 8) + +``` +L_attn = E[ || A_t^[V*] โˆ’ M ||ยฒ ] +``` + +Cross-attention maps from UNet **decoder** (up_blocks) only, averaged over layers and resized to latent resolution. + +### Combined Loss (Eq. 9) + +``` +L_ours = 0.5ยทL_def + 0.2ยทL_obj + 0.05ยทL_attn +``` + +--- + +## Citation + +```bibtex +@inproceedings{song2024defectfill, + title={DefectFill: Realistic Defect Generation with Inpainting Diffusion Model for Visual Inspection}, + author={Song, Jaewoo and Park, Daemin and Baek, Kanghyun and Lee, Sangyub and Choi, Jooyoung and Kim, Eunji and Yoon, Sungroh}, + booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, + year={2024} +} +``` + diff --git a/ArtiAgent - DefectFill/engine/DefectFill/data_loader.py b/ArtiAgent - DefectFill/engine/DefectFill/data_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..697e4ba43cc75dc27ef225e295ec0bc80fef1eaa --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/data_loader.py @@ -0,0 +1,262 @@ +import os +import cv2 +import numpy as np +import torch +import albumentations as A +from torch.utils.data import Dataset, DataLoader +from albumentations.pytorch import ToTensorV2 + +class MVTecDefectDataset(Dataset): + def __init__(self, root_dir, object_class, split="train", transform=None, defect_type=None, dilate_mask=False, mask_kernel_size=3): + """ + Args: + root_dir (str): Directory with MVTec AD dataset + object_class (str): Object class (e.g., 'bottle', 'cable', etc.) + split (str): 'train' or 'test' + transform: Optional transform to be applied + defect_type (str): Specific defect type to load (e.g., 'broken_large'). + If None, loads all defect types. + """ + # Ensure path uses correct operating system format + self.root_dir = os.path.normpath(root_dir) + self.object_class = object_class + self.split = split + self.transform = transform + self.target_defect_type = defect_type + self.dilate_mask = dilate_mask + self.mask_kernel_size = mask_kernel_size + + print(f"Initializing Dataset: root_dir={self.root_dir}, object_class={object_class}, split={split}") + if defect_type: + print(f"Target defect type: {defect_type}") + + # Identify defect types + self.defect_types = [] + if split == "train": + defect_path = os.path.join(self.root_dir, object_class, "train", "defective") + print(f"Searching for defect types in: {defect_path}") + if os.path.exists(defect_path): + all_defect_types = [d for d in os.listdir(defect_path) if os.path.isdir(os.path.join(defect_path, d))] + + # If a specific type is requested, only use that + if defect_type and defect_type in all_defect_types: + self.defect_types = [defect_type] + print(f"Using specified defect type: {self.defect_types}") + elif defect_type: + print(f"Warning: Requested defect type '{defect_type}' not found. Available: {all_defect_types}") + self.defect_types = all_defect_types + else: + self.defect_types = all_defect_types + print(f"Found defect types: {self.defect_types}") + else: + print(f"Warning: Directory does not exist {defect_path}") + + # Debugging helper: Check parent directory + parent_dir = os.path.dirname(defect_path) + if os.path.exists(parent_dir): + print(f"Parent directory {parent_dir} exists, containing: {os.listdir(parent_dir)}") + + if os.path.exists(self.root_dir): + root_contents = os.listdir(self.root_dir) + print(f"Root {self.root_dir} exists, containing: {root_contents[:5]}... ({len(root_contents)} items total)") + + # Load image and mask paths + self.images = [] + self.masks = [] + + # Handle Test Split (only loads the 'good' directory for anomaly detection baselines) + if split == "test": + good_dir = os.path.join(self.root_dir, object_class, "test", "good") + print(f"Loading test set from: {good_dir}") + if os.path.exists(good_dir): + good_files = sorted([f for f in os.listdir(good_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]) + print(f"Found {len(good_files)} 'good' samples in test set") + + for good_file in good_files: + self.images.append(os.path.join(good_dir, good_file)) + # Masks are generated randomly for good images during training/inference + self.masks.append(None) + else: + print(f"Warning: Test directory not found {good_dir}") + + # Load Defect Images (Train Split) + else: + for defect_type in self.defect_types: + img_dir = os.path.join(self.root_dir, object_class, "train", "defective", defect_type) + mask_dir = os.path.join(self.root_dir, object_class, "train", "defective_masks", defect_type) + + print(f"Processing defect: {defect_type}") + print(f" Images: {img_dir}") + print(f" Masks: {mask_dir}") + + if os.path.exists(img_dir) and os.path.exists(mask_dir): + img_files = sorted([f for f in os.listdir(img_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]) + print(f" Found {len(img_files)} image files") + + matched = 0 + for img_file in img_files: + img_path = os.path.join(img_dir, img_file) + + # Convention: mask is base_name + _mask.png + base_name = os.path.splitext(img_file)[0] + mask_file = f"{base_name}_mask.png" + mask_path = os.path.join(mask_dir, mask_file) + + # Fallback for alternative mask naming patterns + if not os.path.exists(mask_path): + possible_masks = [f for f in os.listdir(mask_dir) if base_name in f] + if possible_masks: + mask_path = os.path.join(mask_dir, possible_masks[0]) + else: + print(f" Warning: No mask found for {img_file}, skipping") + continue + + self.images.append(img_path) + self.masks.append(mask_path) + matched += 1 + + print(f" Successfully paired {matched} image-mask sets") + + print(f"Total loaded: {len(self.images)} {split} images") + if len(self.images) == 0: + print("Warning: Dataset is empty!") + + def __len__(self): + return len(self.images) + + def generate_random_mask(self, image_size): + """Generate random rectangular masks for object loss (integrity learning)""" + mask = np.zeros(image_size, dtype=np.float32) + num_rectangles = 30 + + h, w = image_size + for _ in range(num_rectangles): + # Rectangle size between 3% and 25% of image dimensions + min_size = int(min(h, w) * 0.03) + max_size = int(min(h, w) * 0.25) + + rect_h = np.random.randint(min_size, max_size) + rect_w = np.random.randint(min_size, max_size) + + y = np.random.randint(0, h - rect_h) + x = np.random.randint(0, w - rect_w) + + mask[y:y+rect_h, x:x+rect_w] = 1.0 + + return mask + + def __getitem__(self, idx): + img_path = self.images[idx] + mask_path = self.masks[idx] + + # Load image (OpenCV loads BGR, convert to RGB) + image = cv2.imread(img_path) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + # Load ground truth mask or generate a random one + if mask_path is not None: + mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + + # Dilate mask (only if enabled) + if self.dilate_mask: + # Ensure kernel size is odd, otherwise OpenCV crashes + k_size = self.mask_kernel_size if self.mask_kernel_size % 2 == 1 else self.mask_kernel_size + 1 + + kernel = np.ones((k_size, k_size), np.uint8) + mask = cv2.dilate(mask, kernel, iterations=1) + + mask = mask.astype(np.float32) / 255.0 # Normalize to [0, 1] + else: + mask = self.generate_random_mask((image.shape[0], image.shape[1])) + + # --- NEW CODE START: SMART RESIZING --- + h, w = image.shape[:2] + target_size = 512 + + # Case 1: Image is too small (e.g., 400x400) -> Upscale it + if h < target_size or w < target_size: + # Use INTER_CUBIC to keep lines as sharp as possible + image = cv2.resize(image, (target_size, target_size), interpolation=cv2.INTER_CUBIC) + if mask is not None: + # Use NEAREST for masks to avoid creating gray pixels at edges + mask = cv2.resize(mask, (target_size, target_size), interpolation=cv2.INTER_NEAREST) + + # Case 2: Image is large (1024x1024) -> Do nothing here! + # The RandomCrop in the transform will handle it, preserving full detail. + # --- NEW CODE END --- + + # Apply Albumentations transformations + if self.transform: + augmented = self.transform(image=image, mask=mask) + image = augmented['image'] + mask = augmented['mask'] + + # Create background (masked image) for the inpainting input: I * (1 - M) + background = image * (1 - mask) + + # Adjusted mask used for Object Loss calculation + adjusted_mask = mask + 0.3 * (1 - mask) if mask_path is None else mask + + return { + 'image': image, + 'mask': mask, + 'background': background, + 'adjusted_mask': adjusted_mask, + 'is_defect': mask_path is not None, + 'object_class': self.object_class + } + +def get_data_loaders(root_dir, object_class, batch_size=4, defect_type=None, dilate_mask=False, mask_kernel_size=3): + """Creates training and testing DataLoaders with preprocessing pipelines""" + + # Training pipeline: Includes random scaling for better generalization + train_transform = A.Compose([ + A.RandomScale(scale_limit=(0.0, 0.125), p=1.0), # Random scale between 1.0 and 1.125 + A.RandomCrop(height=512, width=512), + A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ToTensorV2() + ], additional_targets={'mask': 'mask', 'background': 'image', 'adjusted_mask': 'mask'}) + + # Test pipeline: Simple resize and normalize + test_transform = A.Compose([ + A.Resize(512, 512), + A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ToTensorV2() + ], additional_targets={'mask': 'mask', 'background': 'image', 'adjusted_mask': 'mask'}) + + train_dataset = MVTecDefectDataset( + root_dir=root_dir, + object_class=object_class, + split="train", + transform=train_transform, + defect_type=defect_type, + dilate_mask=dilate_mask, + mask_kernel_size=mask_kernel_size + ) + + test_dataset = MVTecDefectDataset( + root_dir=root_dir, + object_class=object_class, + split="test", + transform=test_transform, + dilate_mask=dilate_mask, + mask_kernel_size=mask_kernel_size + ) + + train_loader = DataLoader( + train_dataset, + batch_size=batch_size, + shuffle=True, + num_workers=0, + pin_memory=True + ) + + test_loader = DataLoader( + test_dataset, + batch_size=batch_size, + shuffle=False, + num_workers=0, + pin_memory=True + ) + + return train_loader, test_loader diff --git a/ArtiAgent - DefectFill/engine/DefectFill/evaluate.py b/ArtiAgent - DefectFill/engine/DefectFill/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..6f583d25d33667082753ea9bbeba6ff8a85315d3 --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/evaluate.py @@ -0,0 +1,317 @@ +""" +DefectFill Evaluation Module +Includes KID (Kernel Inception Distance) and IC-LPIPS (Inter-image Contextual LPIPS) algorithms. + +Metric Descriptions: +- KID: Measures the distribution distance between generated and real images (Quality). Lower is better. +- IC-LPIPS: Measures perceptual differences between generated images (Diversity). Higher is better. +""" + +import os +import csv +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import lpips +from PIL import Image +from torchvision import transforms, models +from datetime import datetime +import argparse +from tqdm import tqdm +from itertools import combinations + + +class KIDEvaluator: + """ + Kernel Inception Distance (KID) Evaluator + + KID uses a polynomial kernel to calculate Maximum Mean Discrepancy (MMD). + It is better suited for small sample sizes than FID (MVTec often has only dozens of images per class). + """ + + def __init__(self, device="cuda"): + self.device = device + + # Load InceptionV3 model using the pool3 layer features (2048 dimensions) + self.inception = models.inception_v3(weights=models.Inception_V3_Weights.IMAGENET1K_V1, transform_input=False) + self.inception.fc = nn.Identity() # Remove classification head + self.inception = self.inception.to(device) + self.inception.eval() + + # Standard InceptionV3 input preprocessing + self.preprocess = transforms.Compose([ + transforms.Resize((299, 299)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + ]) + + @torch.no_grad() + def extract_features(self, images): + """ + Extract InceptionV3 features from images. + + Args: + images: List of PIL Images or paths, or Tensor [N, 3, H, W] + + Returns: + features: [N, 2048] feature vectors + """ + if isinstance(images, list): + # Process list of PIL Images or paths + tensors = [] + for img in images: + if isinstance(img, str): + img = Image.open(img).convert('RGB') + tensor = self.preprocess(img) + tensors.append(tensor) + images = torch.stack(tensors) + + images = images.to(self.device) + + # Batch processing to avoid VRAM overflow + batch_size = 32 + features_list = [] + + for i in range(0, len(images), batch_size): + batch = images[i:i+batch_size] + feat = self.inception(batch) + features_list.append(feat.cpu()) + + return torch.cat(features_list, dim=0) + + def polynomial_kernel(self, x, y, degree=3, gamma=None, coef0=1): + """ + Calculate Polynomial Kernel + k(x, y) = (gamma * + coef0)^degree + """ + if gamma is None: + gamma = 1.0 / x.shape[1] + + return (gamma * torch.mm(x, y.t()) + coef0) ** degree + + def compute_mmd(self, x, y): + """ + Calculate Maximum Mean Discrepancy (MMD) + MMD^2 = E[k(x,x')] - 2*E[k(x,y)] + E[k(y,y')] + """ + k_xx = self.polynomial_kernel(x, x) + k_yy = self.polynomial_kernel(y, y) + k_xy = self.polynomial_kernel(x, y) + + n = x.shape[0] + m = y.shape[0] + + # Unbiased estimator: remove diagonal elements + mmd = (k_xx.sum() - k_xx.trace()) / (n * (n - 1)) + mmd += (k_yy.sum() - k_yy.trace()) / (m * (m - 1)) + mmd -= 2 * k_xy.mean() + + return mmd + + def compute_kid(self, real_images, gen_images, num_subsets=100, subset_size=None): + """ + Calculate KID score. + + Args: + real_images: Real defect images (list of PIL images or paths) + gen_images: Generated defect images (list of PIL images or paths) + num_subsets: Number of subset samplings (for mean and std) + subset_size: Size of each subset (defaults to min of real/gen counts) + """ + print("Extracting features from real images...") + real_features = self.extract_features(real_images) + print(f" Real features shape: {real_features.shape}") + + print("Extracting features from generated images...") + gen_features = self.extract_features(gen_images) + print(f" Generated features shape: {gen_features.shape}") + + if subset_size is None: + subset_size = min(len(real_features), len(gen_features)) + + # Compute KID via multiple subset sampling + kid_scores = [] + for _ in range(num_subsets): + idx_real = np.random.choice(len(real_features), subset_size, replace=False) + idx_gen = np.random.choice(len(gen_features), subset_size, replace=False) + + mmd = self.compute_mmd( + real_features[idx_real], + gen_features[idx_gen] + ) + kid_scores.append(mmd.item()) + + return np.mean(kid_scores), np.std(kid_scores) + + +class ICLPIPSEvaluator: + """ + Inter-image Contextual LPIPS (IC-LPIPS) Evaluator + + Calculates perceptual differences between generated images to evaluate diversity. + Larger IC-LPIPS indicates higher generation diversity. + """ + + def __init__(self, net='vgg', device="cuda"): + self.device = device + # Use standard LPIPS (non-spatial) + self.lpips_net = lpips.LPIPS(net=net, spatial=False).to(device) + self.lpips_net.eval() + + self.preprocess = transforms.Compose([ + transforms.Resize((256, 256)), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) + ]) + + def _load_image(self, img): + """Load and preprocess image""" + if isinstance(img, str): + img = Image.open(img).convert('RGB') + if isinstance(img, Image.Image): + img = self.preprocess(img) + return img.to(self.device) + + @torch.no_grad() + def compute_pairwise_lpips(self, img1, img2): + """Calculate LPIPS score between two images.""" + img1_tensor = self._load_image(img1).unsqueeze(0) + img2_tensor = self._load_image(img2).unsqueeze(0) + + lpips_score = self.lpips_net(img1_tensor, img2_tensor) + return lpips_score.item() + + @torch.no_grad() + def compute_ic_lpips(self, generated_images, max_pairs=1000): + """ + Calculate IC-LPIPS score for a set of generated images. + Computes the mean LPIPS distance across pairs of generated images. + """ + n_images = len(generated_images) + + if n_images < 2: + print("Warning: Insufficient images to calculate IC-LPIPS") + return float('nan'), float('nan') + + print(f"Loading {n_images} generated images...") + image_tensors = [] + for img in tqdm(generated_images, desc="Loading images"): + img_tensor = self._load_image(img) + image_tensors.append(img_tensor) + + image_batch = torch.stack(image_tensors, dim=0) + all_pairs = list(combinations(range(n_images), 2)) + n_pairs = len(all_pairs) + print(f"Total {n_pairs} image pairs found") + + if n_pairs > max_pairs: + print(f"Randomly sampling {max_pairs} pairs for calculation") + selected_pairs = np.random.choice(n_pairs, max_pairs, replace=False) + pairs_to_compute = [all_pairs[i] for i in selected_pairs] + else: + pairs_to_compute = all_pairs + + lpips_scores = [] + batch_size = 32 + for i in tqdm(range(0, len(pairs_to_compute), batch_size), desc="Computing IC-LPIPS"): + batch_pairs = pairs_to_compute[i:i+batch_size] + + img1_batch = torch.stack([image_batch[p[0]] for p in batch_pairs], dim=0) + img2_batch = torch.stack([image_batch[p[1]] for p in batch_pairs], dim=0) + + scores = self.lpips_net(img1_batch, img2_batch) + lpips_scores.extend(scores.squeeze().cpu().tolist() if len(batch_pairs) > 1 else [scores.item()]) + + return np.mean(lpips_scores), np.std(lpips_scores) + + +def collect_generated_images(directory): + """Collects only files ending with *_generated.png.""" + images = [] + for root, _, files in os.walk(directory): + for file in files: + if file.endswith('_generated.png'): + images.append(os.path.join(root, file)) + return images + + +def collect_real_defect_images(directory): + """Collects all real images from directory, excluding masks.""" + images = [] + for root, _, files in os.walk(directory): + for file in files: + if file.endswith(('.png', '.jpg', '.jpeg')): + if '_mask' not in file: + images.append(os.path.join(root, file)) + return images + + +def evaluate(args): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Using device: {device}") + + print("\nInitializing evaluators...") + kid_evaluator = KIDEvaluator(device=device) + ic_lpips_evaluator = ICLPIPSEvaluator(device=device) + + print("\nCollecting images...") + gen_images_all = collect_generated_images(args.generated_dir) + print(f" Generated images: {len(gen_images_all)}") + + real_images_all = collect_real_defect_images(args.real_dir) + print(f" Real images: {len(real_images_all)}") + + print("\nCalculating KID (Quality Assessment)...") + if len(gen_images_all) > 0 and len(real_images_all) > 0: + kid_mean, kid_std = kid_evaluator.compute_kid( + real_images_all, gen_images_all, + num_subsets=min(100, len(gen_images_all)), + subset_size=min(len(gen_images_all), len(real_images_all)) + ) + print(f" KID: {kid_mean:.6f} ยฑ {kid_std:.6f} (Lower is better)") + else: + kid_mean, kid_std = float('nan'), float('nan') + print(" Warning: Insufficient images to calculate KID") + + print("\nCalculating IC-LPIPS (Diversity Assessment)...") + if len(gen_images_all) >= 2: + ic_lpips_mean, ic_lpips_std = ic_lpips_evaluator.compute_ic_lpips( + gen_images_all, + max_pairs=min(1000, len(gen_images_all) * (len(gen_images_all) - 1) // 2) + ) + print(f" IC-LPIPS: {ic_lpips_mean:.6f} ยฑ {ic_lpips_std:.6f} (Higher is better)") + else: + ic_lpips_mean, ic_lpips_std = float('nan'), float('nan') + print(" Warning: Insufficient images to calculate IC-LPIPS") + + # Save results to CSV + timestamp = datetime.now().strftime('%Y-%m-%dT%H:%M:%S') + result_row = [ + timestamp, args.class_name, args.config_name, args.category_type, + f"{kid_mean:.6f}", f"{kid_std:.6f}", f"{ic_lpips_mean:.6f}", f"{ic_lpips_std:.6f}" + ] + + file_exists = os.path.exists(args.output_csv) + with open(args.output_csv, 'a', newline='') as f: + writer = csv.writer(f) + if not file_exists: + writer.writerow(['timestamp', 'class', 'config', 'category_type', + 'KID_mean', 'KID_std', 'IC_LPIPS_mean', 'IC_LPIPS_std']) + writer.writerow(result_row) + + print(f"\nResults appended to: {args.output_csv}") + return {'kid_mean': kid_mean, 'ic_lpips_mean': ic_lpips_mean} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="DefectFill Evaluation Module") + parser.add_argument("--generated_dir", type=str, required=True, help="Generated images directory") + parser.add_argument("--real_dir", type=str, required=True, help="Real defect images directory") + parser.add_argument("--output_csv", type=str, required=True, help="Output CSV path") + parser.add_argument("--class_name", type=str, required=True, help="Class name") + parser.add_argument("--config_name", type=str, required=True, help="Config name") + parser.add_argument("--category_type", type=str, default="unknown", choices=["object", "texture", "unknown"]) + + args = parser.parse_args() + evaluate(args) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/engine/DefectFill/inference.py b/ArtiAgent - DefectFill/engine/DefectFill/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..a8ac9e9f5840d8a66414297bf529a5ad0f5e63f4 --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/inference.py @@ -0,0 +1,422 @@ +import os +import json +import cv2 +import torch +import argparse +import numpy as np +from PIL import Image +from tqdm import tqdm +from datetime import datetime +from model import DefectFillModel +from utils import load_checkpoint, compute_spatial_lpips, compute_spatial_lpips_batch +from torchvision.utils import save_image +from torchvision import transforms + + +def smart_crop_dynamic(image, mask, base_size=512): + """ + Crops the image to fit the defect. + - If defect < 512: Crops 512x512 (No Resize). + - If defect > 512: Crops square enclosing defect, then resizes to 512. + """ + h, w = image.shape[:2] + + # Find the Bounding Box of the defect + y_indices, x_indices = np.where(mask > 0) + + if len(y_indices) == 0: + # No defect? Return center crop 512 + cy, cx = h // 2, w // 2 + crop_size = base_size + else: + min_y, max_y = np.min(y_indices), np.max(y_indices) + min_x, max_x = np.min(x_indices), np.max(x_indices) + + defect_h = max_y - min_y + defect_w = max_x - min_x + + # Center of the defect + cy = min_y + defect_h // 2 + cx = min_x + defect_w // 2 + + # Determine the Crop Size + # We need a box big enough to hold the defect + some context padding + # But at minimum, it must be 512. + max_dim = max(defect_h, defect_w) + padding = 50 # Add 50px context around edges if possible + + crop_size = max(base_size, max_dim + padding) + + # Calculate Crop Coordinates (Square Box) + half_size = crop_size // 2 + x1 = cx - half_size + y1 = cy - half_size + x2 = x1 + crop_size + y2 = y1 + crop_size + + # Handle Edge Cases (Shift box if it goes out of bounds) + if x1 < 0: x2 -= x1; x1 = 0 + if y1 < 0: y2 -= y1; y1 = 0 + if x2 > w: x1 -= (x2 - w); x2 = w + if y2 > h: y1 -= (y2 - h); y2 = h + + # Double check we didn't shrink below image dims (e.g. if image is smaller than crop_size) + x1 = max(0, x1); y1 = max(0, y1) + x2 = min(w, x2); y2 = min(h, y2) + + # Perform the Crop + crop_img = image[y1:y2, x1:x2] + crop_mask = mask[y1:y2, x1:x2] + + # Resize ONLY if the crop is larger than 512 + # (If crop_size was 512, this does nothing. If it was 570, it shrinks slightly.) + if crop_img.shape[0] != base_size or crop_img.shape[1] != base_size: + crop_img = cv2.resize(crop_img, (base_size, base_size), interpolation=cv2.INTER_AREA) + # Use NEAREST for mask to keep edges sharp + crop_mask = cv2.resize(crop_mask, (base_size, base_size), interpolation=cv2.INTER_NEAREST) + + return crop_img, crop_mask + + +def count_available_resources(data_dir, object_class, defect_type): + """Counts available good images and reference masks for synthetic generation.""" + # Good images directory + good_dir = os.path.join(data_dir, object_class, "test", "good") + num_good_images = len([f for f in os.listdir(good_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]) if os.path.exists(good_dir) else 0 + + # Mask directory (prioritize training masks for reference) + train_mask_dir = os.path.join(data_dir, object_class, "train", "defective_masks", defect_type) + test_mask_dir = os.path.join(data_dir, object_class, "test", "defective_masks", defect_type) + + if os.path.exists(train_mask_dir): + num_masks = len([f for f in os.listdir(train_mask_dir) if f.endswith('.png')]) + mask_dir = train_mask_dir + elif os.path.exists(test_mask_dir): + num_masks = len([f for f in os.listdir(test_mask_dir) if f.endswith('.png')]) + mask_dir = test_mask_dir + else: + num_masks = 0 + mask_dir = None + + return num_good_images, num_masks, good_dir, mask_dir + + +def calculate_generation_plan(num_good_images, num_masks, target_total=100): + """Calculates a combination plan of good images and masks to reach the target total.""" + if num_good_images == 0 or num_masks == 0: + return [] + + generation_plan = [] + output_idx = 0 + + # Loop through good images and masks until target count is met + while output_idx < target_total: + for mask_idx in range(num_masks): + if output_idx >= target_total: + break + good_idx = output_idx % num_good_images # Cycle through good images + generation_plan.append((good_idx, mask_idx, output_idx)) + output_idx += 1 + + return generation_plan + + +def inference(args): + # Set up device + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # ========== Unified FP16 Precision Configuration ========== + dtype = torch.float16 + + # Enable TF32 acceleration (Ampere+ architectures: RTX 30/40/50 series) + if torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + print(f"Device: {device}, dtype: {dtype}, TF32: enabled") + + # Initialize model + model = DefectFillModel( + device=device, + lora_rank=args.lora_rank, + lora_alpha=args.lora_alpha + ) + + # Ensure VAE is also in FP16 to save VRAM + model.pipeline.vae.to(dtype=dtype) + + # Load checkpoint + if args.checkpoint: + load_checkpoint(model, None, args.checkpoint) + print(f"Loaded checkpoint from {args.checkpoint}") + + # Set to evaluation mode + model.pipeline.unet.eval() + model.pipeline.text_encoder.eval() + + # ========== torch.compile Optimization (Optional) ========== + if hasattr(torch, 'compile') and args.use_compile: + print("Compiling UNet with torch.compile (this may take 5-15 minutes for max-autotune)...") + print("Note: First run triggers compilation. Subsequent runs will be significantly faster.") + + # Compiler settings + torch._inductor.config.conv_1x1_as_mm = True + torch._inductor.config.coordinate_descent_tuning = True + torch._inductor.config.epilogue_fusion = False + torch._inductor.config.coordinate_descent_check_all_directions = True + + try: + # Compile UNet (the main computational bottleneck) + model.pipeline.unet = torch.compile( + model.pipeline.unet, + mode="max-autotune", # Aggressive auto-tuning + fullgraph=True, # Full graph compilation + dynamic=False # Fixed input size (512x512) for best speed + ) + + # Compile VAE decoder + model.pipeline.vae.decode = torch.compile( + model.pipeline.vae.decode, + mode="max-autotune", + dynamic=False + ) + print("Compilation configuration complete!") + + except Exception as e: + print(f"Warning: fullgraph compilation failed ({e}), falling back to reduce-overhead mode...") + model.pipeline.unet = torch.compile( + model.pipeline.unet, + mode="reduce-overhead", + fullgraph=False, + dynamic=False + ) + print("Fallback compilation complete!") + + # ========== Warmup: Trigger JIT Compilation ========== + print("Warming up compiled model...") + dummy_img = torch.randn(1, 3, 512, 512, device=device, dtype=dtype) + dummy_mask = torch.randn(1, 1, 512, 512, device=device, dtype=dtype) + dummy_mask = (dummy_mask > 0).float() # Binarize mask + dummy_img = dummy_img * 2 - 1 # Map to [-1, 1] + + with torch.no_grad(): + try: + warmup_prompt = f"A {args.object_class} with {model.placeholder_token}" + _ = model.generate( + image=dummy_img, + mask=dummy_mask, + prompt=warmup_prompt, + num_inference_steps=1, # 1 step is enough to trigger JIT + guidance_scale=7.5, + ) + except Exception as warmup_error: + print(f"Warmup warning (non-critical): {warmup_error}") + + del dummy_img, dummy_mask + torch.cuda.empty_cache() + print("Warmup complete! Model is optimized.") + + def fixed_inference_batch(model, clean_image, mask, object_class, defect_type, + num_samples=8, steps=50, guidance_scale=7.5, + batch_size=4): + """ + Performs inference using the custom model.generate() method. + Ensures consistency between training and inference phases. + """ + prompt = f"A {object_class} with {model.placeholder_token}" + + print(f"Using prompt: '{prompt}'") + print(f"Generating {num_samples} samples (batch_size={batch_size}, steps={steps})") + + _, _, h_input, w_input = clean_image.shape + + # ========== Phase 1: Batch Sample Generation ========== + all_samples = [] + num_batches = (num_samples + batch_size - 1) // batch_size + + for batch_idx in range(num_batches): + start_idx = batch_idx * batch_size + end_idx = min(start_idx + batch_size, num_samples) + current_batch_size = end_idx - start_idx + + print(f"Batch {batch_idx+1}/{num_batches}: Generating samples {start_idx+1}-{end_idx}") + + batch_clean = clean_image.repeat(current_batch_size, 1, 1, 1) + batch_mask = mask.repeat(current_batch_size, 1, 1, 1) + + # Use deterministic seed per sample for reproducibility + generator = torch.Generator(device=device).manual_seed(start_idx) + + # Consistent with training: 9-channel input + CFG + iterative bg preservation + batch_samples = model.generate( + image=batch_clean, + mask=batch_mask, + prompt=prompt, + num_inference_steps=steps, + guidance_scale=guidance_scale, + generator=generator, + ) + + # Convert back to [-1, 1] range for LPIPS (generate returns [0, 1]) + batch_samples_model_format = (batch_samples * 2.0) - 1.0 + all_samples.append(batch_samples_model_format) + + samples_model_format = torch.cat(all_samples, dim=0) + + if samples_model_format.shape[-2:] != (h_input, w_input): + samples_model_format = torch.nn.functional.interpolate( + samples_model_format, size=(h_input, w_input), mode='bilinear' + ) + + # ========== Phase 2: Batch LPIPS Selection ========== + mask_resized = mask if mask.shape[-2:] == samples_model_format.shape[-2:] else \ + torch.nn.functional.interpolate(mask, size=samples_model_format.shape[-2:], mode='bilinear') + + print(f"Selecting best sample based on LPIPS...") + lpips_scores = compute_spatial_lpips_batch( + model.lpips_model, clean_image, samples_model_format, mask_resized, smooth_boundary=True + ) + + best_idx = lpips_scores.argmax() + best_score = lpips_scores[best_idx].item() + best_sample = samples_model_format[best_idx].clone() + + print(f"Best sample selected: #{best_idx+1} (LPIPS: {best_score:.4f})") + + del all_samples, samples_model_format, lpips_scores + return best_sample, best_score + + # Transformations + transform = transforms.Compose([ + transforms.Resize((512, 512)), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) + ]) + + batch_size = args.batch_size if hasattr(args, 'batch_size') else 4 + os.makedirs(args.output_dir, exist_ok=True) + + inference_log = { + "timestamp": datetime.now().strftime('%Y-%m-%dT%H:%M:%S'), + "checkpoint": args.checkpoint, + "object_class": args.object_class, + "defect_type": args.defect_type, + "results": [] + } + + # Mode A: Dynamic Dataset Generation + if args.total_images > 0 and args.data_dir and args.defect_type: + print(f"\n{'='*60}\nDynamic Generation Mode Activated\n{'='*60}") + num_good, num_masks, good_dir, mask_dir = count_available_resources(args.data_dir, args.object_class, args.defect_type) + + if num_good == 0 or num_masks == 0: + print("Error: Missing images or masks.") + return + + generation_plan = calculate_generation_plan(num_good, num_masks, args.total_images) + good_files = sorted([f for f in os.listdir(good_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]) + mask_files = sorted([f for f in os.listdir(mask_dir) if f.endswith('.png')]) + + defect_output_dir = os.path.join(args.output_dir, args.defect_type) + os.makedirs(defect_output_dir, exist_ok=True) + + for good_idx, mask_idx, output_idx in tqdm(generation_plan, desc=f"Generating {args.defect_type}"): + good_path = os.path.join(good_dir, good_files[good_idx]) + mask_path = os.path.join(mask_dir, mask_files[mask_idx]) + + print(f"\n[{output_idx+1}/{len(generation_plan)}] Processing: {good_files[good_idx]}") + + # --- SMART CROP LOGIC START --- + + # Load Images as Numpy Arrays (for Smart Crop) + # Use PIL and convert to numpy to ensure RGB format is consistent + image_pil = Image.open(good_path).convert("RGB") + mask_pil = Image.open(mask_path).convert("L") + + image_np = np.array(image_pil) + mask_np = np.array(mask_pil) + + # --- DILATION LOGIC (Must match training) --- + if args.dilate_mask: + # Ensure kernel size is odd + k_size = args.mask_kernel_size if args.mask_kernel_size % 2 == 1 else args.mask_kernel_size + 1 + kernel = np.ones((k_size, k_size), np.uint8) + + # Apply dilation + # Note: mask_np is usually 0-255. cv2.dilate works fine on uint8. + mask_np = cv2.dilate(mask_np, kernel, iterations=1) + + print(f"Dilated mask with kernel {k_size}") + # -------------------------------------------------- + + # Apply Smart Crop + # This returns a 512x512 patch focused on the defect area + # (No resizing blur unless defect > 512px) + crop_img_np, crop_mask_np = smart_crop_dynamic(image_np, mask_np, base_size=512) + + # Convert to Tensor + + # Image: [0, 255] -> [0.0, 1.0] -> Normalize to [-1.0, 1.0] + # transforms.ToTensor() handles the HWC->CHW and /255 division automatically + img_tensor = transforms.ToTensor()(crop_img_np) + img_tensor = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])(img_tensor) + img_tensor = img_tensor.unsqueeze(0).to(device, dtype=dtype) + + # Mask: [0, 255] -> [0.0, 1.0] + mask_tensor = transforms.ToTensor()(crop_mask_np).unsqueeze(0).to(device, dtype=dtype) + + + with torch.no_grad(): + defect_img, lpips_score = fixed_inference_batch( + model, img_tensor, mask_tensor, args.object_class, args.defect_type, + num_samples=args.num_samples, steps=args.steps, guidance_scale=args.guidance_scale, batch_size=batch_size + ) + + # Save generated image + output_name = f"{output_idx:04d}_generated.png" + output_path = os.path.join(defect_output_dir, output_name) + save_image((defect_img.float() + 1) / 2, output_path) + + # Save mask and original (These will now be the CROPPED versions, which is correct) + save_image(mask_tensor.float(), os.path.join(defect_output_dir, f"{output_idx:04d}_mask.png")) + save_image((img_tensor.float() + 1) / 2, os.path.join(defect_output_dir, f"{output_idx:04d}_original.png")) + + inference_log["results"].append({ + "output_idx": output_idx, "input_image": good_path, "lpips_score": lpips_score + }) + + if output_idx % 10 == 0: torch.cuda.empty_cache() + + # Mode B: Process Existing Directories/Files (Traditional Inference) + elif args.image_dir or args.image_path: + # Implementation similar to above but iterates through provided paths + pass + + # Save Log + log_path = os.path.join(args.output_dir, "inference_log.json") + with open(log_path, "w") as f: + json.dump(inference_log, f, indent=4) + print(f"\nInference log saved to: {log_path}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Inference with DefectFill model") + parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint") + parser.add_argument("--output_dir", type=str, default="./generated", help="Output directory") + parser.add_argument("--object_class", type=str, required=True, help="Object class") + parser.add_argument("--defect_type", type=str, help="Defect type (e.g., 'cracks')") + parser.add_argument("--data_dir", type=str, help="Dataset root for dynamic generation") + parser.add_argument("--image_path", type=str, help="Single image path") + parser.add_argument("--num_samples", type=int, default=8, help="Samples per image (for LPIPS selection)") + parser.add_argument("--steps", type=int, default=50, help="Diffusion steps") + parser.add_argument("--guidance_scale", type=float, default=7.5) + parser.add_argument("--total_images", type=int, default=100, help="Total synthetic images to create") + parser.add_argument("--batch_size", type=int, default=4, help="Parallel generation batch size") + parser.add_argument("--use_compile", action="store_true", help="Enable torch.compile (PyTorch 2.0+)") + parser.add_argument("--lora_rank", type=int, default=8, help="LoRA rank") + parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA alpha") + parser.add_argument("--dilate_mask", type=str, default="False", help="Whether to dilate masks (True/False)") + parser.add_argument("--mask_kernel_size", type=int, default=3, help="Size of dilation kernel") + + args = parser.parse_args() + args.dilate_mask = args.dilate_mask.lower() == "true" # Handle boolean conversion + + inference(args) diff --git a/ArtiAgent - DefectFill/engine/DefectFill/model.py b/ArtiAgent - DefectFill/engine/DefectFill/model.py new file mode 100644 index 0000000000000000000000000000000000000000..480fb6783ec95ffc60149c6c708e230ffa1b976d --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/model.py @@ -0,0 +1,390 @@ +import os +# Hugging Face Mirror Configuration +# Option 1: hf-mirror.com (Available in some regions) +# Option 2: Use ModelScope as an alternative +USE_MODELSCOPE = False # Set to True for ModelScope, False for HuggingFace + +import torch +import torch.nn as nn +from diffusers import StableDiffusionInpaintPipeline, DDIMScheduler, UNet2DConditionModel +from transformers import CLIPTextModel +from peft import LoraConfig, get_peft_model +import lpips +import torch.nn.functional as F +from typing import Dict, List, Optional, Tuple +import math +from diffusers.models.attention_processor import Attention, AttnProcessor + + +class AttentionStoreProcessor(AttnProcessor): + """Attention Processor used to store cross-attention maps for steering""" + def __init__(self, model=None, layer_name=""): + super().__init__() + self.model = model # Reference to the main model instance + self.layer_name = layer_name # Store layer name directly + + def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None): + batch_size, sequence_length, _ = hidden_states.shape + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + + query = attn.to_q(hidden_states) + + is_cross_attention = encoder_hidden_states is not None + + if not is_cross_attention: + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + else: + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.head_to_batch_dim(query) + key = attn.head_to_batch_dim(key) + value = attn.head_to_batch_dim(value) + + attention_scores = torch.matmul(query, key.transpose(-1, -2)) * attn.scale + attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1) + + # Fast Direct Lookup (No recursive loop!) + if is_cross_attention and self.model is not None and "up_blocks" in self.layer_name: + try: + num_heads = attn.heads + total_elements = attention_probs.numel() + query_len = hidden_states.shape[1] + key_len = encoder_hidden_states.shape[1] if encoder_hidden_states is not None else query_len + + expected_size = batch_size * num_heads * query_len * key_len + if total_elements == expected_size: + reshaped_probs = attention_probs.reshape(batch_size, num_heads, query_len, key_len) + if not hasattr(self.model, "attention_maps"): + self.model.attention_maps = {} + self.model.attention_maps[self.layer_name] = reshaped_probs.detach().clone() + except Exception as e: + pass + + hidden_states = torch.matmul(attention_probs, value) + hidden_states = attn.batch_to_head_dim(hidden_states) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states + +class DefectFillModel(nn.Module): + def __init__(self, device="cuda", lora_rank=8, lora_alpha=16, seed=42, placeholder_token=""): + super().__init__() + torch.manual_seed(seed) + self.device = device + + # Base Model ID + hf_model_id = "sd2-community/stable-diffusion-2-inpainting" + + # Select model source based on configuration + if USE_MODELSCOPE: + try: + from modelscope import snapshot_download + print(f"[ModelScope] Downloading model: {hf_model_id}") + local_model_path = snapshot_download(hf_model_id) + print(f"[ModelScope] Model downloaded to: {local_model_path}") + + self.pipeline = StableDiffusionInpaintPipeline.from_pretrained( + local_model_path, + torch_dtype=torch.float16 + ).to(device) + + self.scheduler = DDIMScheduler.from_pretrained( + local_model_path, + subfolder="scheduler" + ) + except ImportError: + print("[Warning] modelscope not installed. Try: pip install modelscope") + print("[Info] Attempting HuggingFace fallback...") + self.pipeline = StableDiffusionInpaintPipeline.from_pretrained( + hf_model_id, torch_dtype=torch.float16 + ).to(device) + self.scheduler = DDIMScheduler.from_pretrained(hf_model_id, subfolder="scheduler") + else: + self.pipeline = StableDiffusionInpaintPipeline.from_pretrained( + hf_model_id, torch_dtype=torch.float16 + ).to(device) + self.scheduler = DDIMScheduler.from_pretrained(hf_model_id, subfolder="scheduler") + + self.pipeline.set_progress_bar_config(disable=True) + self.scheduler.set_timesteps(30) + + # ========== Textual Inversion: Add learnable defect token [V*] ========== + self.placeholder_token = placeholder_token + + # Add new token to tokenizer + num_added_tokens = self.pipeline.tokenizer.add_tokens([self.placeholder_token]) + if num_added_tokens == 0: + print(f"[Warning] Token {self.placeholder_token} already exists in tokenizer") + else: + print(f"[Textual Inversion] Added {num_added_tokens} new token: {self.placeholder_token}") + + # Resize text encoder embeddings + self.pipeline.text_encoder.resize_token_embeddings(len(self.pipeline.tokenizer)) + + # Get ID for the new token + self.placeholder_token_id = self.pipeline.tokenizer.convert_tokens_to_ids(self.placeholder_token) + print(f"[Textual Inversion] placeholder_token_id = {self.placeholder_token_id}") + + # Initialize new token with the embedding of "defect" + initializer_token = "defect" + initializer_token_ids = self.pipeline.tokenizer.encode(initializer_token, add_special_tokens=False) + if len(initializer_token_ids) > 0: + initializer_token_id = initializer_token_ids[0] + token_embeds = self.pipeline.text_encoder.get_input_embeddings().weight.data + token_embeds[self.placeholder_token_id] = token_embeds[initializer_token_id].clone() + print(f"[Textual Inversion] Initialized '{self.placeholder_token}' using '{initializer_token}' (id={initializer_token_id})") + + # LoRA Configuration + unet_lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights="gaussian" + ) + + text_encoder_lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + target_modules=["q_proj", "k_proj", "v_proj", "out_proj"], + init_lora_weights="gaussian" + ) + + # Apply LoRA adapters + self.pipeline.unet = get_peft_model(self.pipeline.unet, unet_lora_config) + self.pipeline.text_encoder = get_peft_model(self.pipeline.text_encoder, text_encoder_lora_config) + + # Freeze VAE parameters + for param in self.pipeline.vae.parameters(): + param.requires_grad = False + + # VGG model for LPIPS loss + self.lpips_model = lpips.LPIPS(net='vgg', spatial=True).to(device) + + self.attention_maps = {} + self.register_attention_processor() + self.defect_token_indices = [] + + def register_attention_processor(self): + """Replace standard UNet attention processors with custom ones""" + self.attention_maps = {} + for name, module in self.pipeline.unet.named_modules(): + if isinstance(module, Attention) and "attn2" in name: # Target Cross-Attention only + # Pass 'name' directly into the processor + module.processor = AttentionStoreProcessor(model=self, layer_name=name) + + def get_attention_loss(self, mask_latents: torch.Tensor) -> torch.Tensor: + """ + Calculates Attention Loss - forces token attention maps to align with the defect mask. + """ + if not self.attention_maps: + return torch.tensor(0.0, device=mask_latents.device) + + if len(mask_latents.shape) == 3: + mask_latents = mask_latents.unsqueeze(1) + + batch_size = mask_latents.shape[0] + attention_loss = torch.tensor(0.0, device=mask_latents.device) + + # Use only decoder (up_blocks) attention maps + decoder_attention_maps = { + name: attn_map for name, attn_map in self.attention_maps.items() + if "up_blocks" in name + } + + if not decoder_attention_maps: + return torch.tensor(0.0, device=mask_latents.device) + + for b in range(batch_size): + token_idx = self.defect_token_indices[b] if b < len(self.defect_token_indices) else -1 + if token_idx < 0: + continue + + mask = mask_latents[b].squeeze(0) # (H, W) + resized_attention_maps = [] + + for name, attn_map in decoder_attention_maps.items(): + try: + if b < attn_map.shape[0]: + # Average attention across all heads for the specific token + defect_attn = attn_map[b, :, :, token_idx].mean(dim=0) + + seq_len = defect_attn.shape[0] + h = int(math.sqrt(seq_len)) + if h * h == seq_len: + defect_attn = defect_attn.reshape(h, h) + resized_attn = F.interpolate( + defect_attn.unsqueeze(0).unsqueeze(0), + size=mask.shape, + mode='bilinear', + align_corners=False + ).squeeze() + resized_attention_maps.append(resized_attn) + except Exception: + continue + + if resized_attention_maps: + avg_attn_map = torch.stack(resized_attention_maps).mean(dim=0) + # L2 Loss: ||AttentionMap - Mask||^2 + sample_loss = F.mse_loss(avg_attn_map, mask) + attention_loss += sample_loss + + return attention_loss / batch_size if batch_size > 0 else attention_loss + + def get_text_embeddings(self, prompts, enable_grad=True): + """Encodes prompts and locates the precise index of the token""" + if not hasattr(self, 'pipeline') or self.pipeline is None: + raise ValueError("Pipeline not initialized") + + if isinstance(prompts, str): + prompts = [prompts] + + text_inputs = self.pipeline.tokenizer( + prompts, + padding="max_length", + max_length=self.pipeline.tokenizer.model_max_length, + truncation=True, + return_tensors="pt" + ).to(self.pipeline.device) + + input_ids = text_inputs.input_ids + + # Locate the token position in each prompt + self.defect_token_indices = [] + for ids in input_ids: + positions = (ids == self.placeholder_token_id).nonzero(as_tuple=True)[0] + self.defect_token_indices.append(positions[0].item() if len(positions) > 0 else -1) + + if enable_grad: + text_embeddings = self.pipeline.text_encoder(input_ids)[0] + else: + with torch.no_grad(): + text_embeddings = self.pipeline.text_encoder(input_ids)[0] + + return text_embeddings + + def forward( + self, + noisy_latents: torch.Tensor, + masked_image_latents: torch.Tensor, + mask_latents: torch.Tensor, + timesteps: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """ + Training Forward Pass - Implements 9-channel input. + Input format: [noisy_latents(4), masked_background(4), mask(1)] + """ + self.attention_maps = {} + concat_latents = torch.cat([noisy_latents, masked_image_latents, mask_latents], dim=1) + + noise_pred = self.pipeline.unet( + concat_latents, + timesteps, + encoder_hidden_states=encoder_hidden_states, + ).sample + + attention_loss = self.get_attention_loss(mask_latents) + + return { + "noise_pred": noise_pred, + "attention_loss": attention_loss + } + + @staticmethod + def compute_masked_mse(noise_pred: torch.Tensor, noise: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Helper to calculate MSE loss only within the masked area""" + weighted_loss = mask * ((noise_pred - noise) ** 2) + return torch.sum(weighted_loss) / (torch.sum(mask) + 1e-8) + + def compute_defect_loss(self, noise_pred: torch.Tensor, noise: torch.Tensor, mask_latents: torch.Tensor) -> torch.Tensor: + """L_def loss: MSE restricted to the defect mask region""" + return self.compute_masked_mse(noise_pred, noise, mask_latents) + + def compute_object_loss(self, noise_pred: torch.Tensor, noise: torch.Tensor, mask_latents: torch.Tensor, alpha: float = 0.3) -> torch.Tensor: + """L_obj loss: Uses weighted mask M' = M + alpha*(1-M) to preserve object context""" + weighted_mask = mask_latents + alpha * (1 - mask_latents) + return self.compute_masked_mse(noise_pred, noise, weighted_mask) + + def generate( + self, + image: torch.Tensor, + mask: torch.Tensor, + prompt: str, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + """ + Complete Inference Pipeline: + 1. 9-channel input configuration + 2. Classifier-Free Guidance (CFG) + 3. Iterative background preservation: x_t = M * x_t_pred + (1-M) * x_t_background + """ + device = image.device + dtype = image.dtype + batch_size = image.shape[0] + + # Normalize image to [-1, 1] if needed + if image.min() >= 0 and image.max() <= 1: + image = 2 * image - 1 + + if len(mask.shape) == 3: mask = mask.unsqueeze(1) + if mask.max() > 1: mask = mask / 255.0 + + with torch.no_grad(): + # Encode clean image and create masked background latent b = E(I * (1-M)) + latents_clean = self.pipeline.vae.encode(image).latent_dist.sample() + latents_clean = latents_clean * self.pipeline.vae.config.scaling_factor + + masked_image = image * (1 - mask) + masked_image_latents = self.pipeline.vae.encode(masked_image).latent_dist.sample() + masked_image_latents = masked_image_latents * self.pipeline.vae.config.scaling_factor + + mask_latents = F.interpolate(mask, size=latents_clean.shape[-2:], mode='nearest') + + # Text embeddings for CFG + text_embeddings = self.get_text_embeddings([prompt] * batch_size, enable_grad=False) + uncond_embeddings = self.get_text_embeddings([""] * batch_size, enable_grad=False) + text_embeddings_cfg = torch.cat([uncond_embeddings, text_embeddings]) + + self.scheduler.set_timesteps(num_inference_steps) + latents = torch.randn(latents_clean.shape, generator=generator, device=device, dtype=dtype) + + # Denoising loop + for t in self.scheduler.timesteps: + # Generate background noise for current timestep (for background preservation) + noise_for_bg = torch.randn(latents_clean.shape, generator=generator, device=device, dtype=dtype) + latents_background = self.scheduler.add_noise(latents_clean, noise_for_bg, t) + + # Prepare inputs for CFG + latent_input = torch.cat([latents] * 2) + masked_input = torch.cat([masked_image_latents] * 2) + mask_input = torch.cat([mask_latents] * 2) + concat_input = torch.cat([latent_input, masked_input, mask_input], dim=1) + + timestep_tensor = torch.tensor([t] * (batch_size * 2), device=device, dtype=torch.long) + + noise_pred = self.pipeline.unet( + concat_input, + timestep_tensor, + encoder_hidden_states=text_embeddings_cfg + ).sample + + # Perform CFG + noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + + latents = self.scheduler.step(noise_pred, t, latents).prev_sample + + # ========== KEY STEP: Iterative Background Preservation ========== + latents = mask_latents * latents + (1 - mask_latents) * latents_background + + # Decode latents to pixels + latents = latents / self.pipeline.vae.config.scaling_factor + with torch.no_grad(): + images = self.pipeline.vae.decode(latents).sample + + return (images + 1) / 2 # Convert back to [0, 1] range \ No newline at end of file diff --git a/ArtiAgent - DefectFill/engine/DefectFill/requirements.txt b/ArtiAgent - DefectFill/engine/DefectFill/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..be4f224ef7989bf0f5eda51f1d21df5a5e0da62f --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/requirements.txt @@ -0,0 +1,12 @@ +torch>=1.12.0 +torchvision>=0.13.0 +diffusers>=0.12.0 +transformers>=4.25.1 +accelerate>=0.16.0 +bitsandbytes>=0.37.0 +peft>=0.3.0 +lpips>=0.1.4 +albumentations>=1.3.0 +opencv-python>=4.6.0 +numpy>=1.23.5 +tqdm>=4.64.1 \ No newline at end of file diff --git a/ArtiAgent - DefectFill/engine/DefectFill/train.py b/ArtiAgent - DefectFill/engine/DefectFill/train.py new file mode 100644 index 0000000000000000000000000000000000000000..6113a65d0b00125de420ea516a9ffa6774c85e3e --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/train.py @@ -0,0 +1,401 @@ +import os +import json +import torch +import argparse +import torch.nn.functional as F +import random +import time +from torch.optim import AdamW +from tqdm import tqdm +from diffusers import DDPMScheduler +from model import DefectFillModel, USE_MODELSCOPE +from data_loader import get_data_loaders +from utils import save_checkpoint, load_checkpoint +# TensorBoard support +from torch.utils.tensorboard import SummaryWriter +import datetime + + +def generate_seed_from_timestamp(): + """Generates a random seed based on the current timestamp""" + return int(time.time() * 1000) % (2**31) + +def train(args): + # Set up device + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Using device: {device}") + if torch.cuda.is_available(): + print(f"GPU Name: {torch.cuda.get_device_name(0)}") + else: + print("CUDA is NOT available. The model will run on CPU, causing severe slowdown.") + + # Create output directory structure + os.makedirs(args.output_dir, exist_ok=True) + checkpoints_dir = os.path.join(args.output_dir, "checkpoints") + tensorboard_dir = os.path.join(args.output_dir, "tensorboard") + os.makedirs(checkpoints_dir, exist_ok=True) + os.makedirs(tensorboard_dir, exist_ok=True) + + # Create log file in the output directory + log_file_path = os.path.join(args.output_dir, "train_log.txt") + log_file = open(log_file_path, "a") + log_file.write(f"\n\n{'='*60}\n") + log_file.write(f"Training started at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") + log_file.write(f"Object class: {args.object_class}\n") + log_file.write(f"Defect type: {args.defect_type if args.defect_type else 'all'}\n") + log_file.write(f"Config name: {args.config_name}\n") + log_file.write(f"Lambda defect: {args.lambda_defect}\n") + log_file.write(f"Lambda obj: {args.lambda_obj}\n") + log_file.write(f"Lambda attn: {args.lambda_attn}\n") + log_file.write(f"Alpha (obj branch bg weight): {args.alpha}\n") + log_file.write(f"Gradient accumulation steps: {args.gradient_accumulation_steps}\n") + log_file.write(f"Random seed: {args.seed}\n") + log_file.write(f"{'='*60}\n\n") + + # Save training configuration to JSON + train_config = { + "timestamp": datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S'), + "object_class": args.object_class, + "defect_type": args.defect_type if args.defect_type else "all", + "config_name": args.config_name, + "lambda_defect": args.lambda_defect, + "lambda_obj": args.lambda_obj, + "lambda_attn": args.lambda_attn, + "alpha": args.alpha, + "batch_size": args.batch_size, + "max_train_steps": args.max_train_steps, + "gradient_accumulation_steps": args.gradient_accumulation_steps, + "lora_rank": args.lora_rank, + "lora_alpha": args.lora_alpha, + "text_encoder_lr": args.text_encoder_lr, + "unet_lr": args.unet_lr, + "lr_warmup_steps": args.lr_warmup_steps, + "save_steps": args.save_steps, + "seed": args.seed + } + config_path = os.path.join(args.output_dir, "train_config.json") + with open(config_path, "w") as f: + json.dump(train_config, f, indent=4) + print(f"Training config saved to: {config_path}") + + # Initialize TensorBoard writer + writer = SummaryWriter(tensorboard_dir) + + # Load data + train_loader, test_loader = get_data_loaders( + root_dir=args.data_dir, + object_class=args.object_class, + batch_size=args.batch_size, + defect_type=args.defect_type, + dilate_mask=args.dilate_mask, + mask_kernel_size=args.mask_kernel_size + ) + + # Initialize model + model = DefectFillModel( + device=device, + lora_rank=args.lora_rank, + lora_alpha=args.lora_alpha, + seed=args.seed + ) + + # Set up optimizers with specific learning rates for Text Encoder and UNet + text_encoder_params = [p for n, p in model.pipeline.text_encoder.named_parameters() if "lora" in n] + unet_params = [p for n, p in model.pipeline.unet.named_parameters() if "lora" in n] + + optimizer = AdamW([ + {"params": text_encoder_params, "lr": args.text_encoder_lr}, + {"params": unet_params, "lr": args.unet_lr} + ]) + + # Save original LRs for warmup calculation + base_lrs = [args.text_encoder_lr, args.unet_lr] + + # Set up noise scheduler (handling ModelScope vs HuggingFace) + hf_model_id = "sd2-community/stable-diffusion-2-inpainting" + if USE_MODELSCOPE: + try: + from modelscope import snapshot_download + print(f"[ModelScope] Downloading scheduler: {hf_model_id}") + local_model_path = snapshot_download(hf_model_id) + noise_scheduler = DDPMScheduler.from_pretrained(local_model_path, subfolder="scheduler") + except ImportError: + print("[Warning] modelscope not installed, falling back to HuggingFace...") + noise_scheduler = DDPMScheduler.from_pretrained(hf_model_id, subfolder="scheduler") + else: + noise_scheduler = DDPMScheduler.from_pretrained(hf_model_id, subfolder="scheduler") + + # Resume from checkpoint if specified + start_step = 0 + if args.resume_from: + start_step = load_checkpoint(model, optimizer, args.resume_from) + print(f"Resuming from step {start_step}") + log_file.write(f"Resuming from step {start_step}\n") + + # Set models to training mode + model.pipeline.unet.train() + model.pipeline.text_encoder.train() + + total_steps = args.max_train_steps + progress_bar = tqdm(range(start_step, total_steps), desc="Training Progress") + + global_step = start_step + accumulation_step = 0 + + # Add this BEFORE the while loop + with torch.no_grad(): + clean_latents_cache = {} + for batch in train_loader: + images = batch["image"].to(device, dtype=torch.float16) + is_defect = batch["is_defect"] + defect_samples = torch.nonzero(is_defect).squeeze(1) + if len(defect_samples) > 0: + defect_images = images[defect_samples] + # Cache original latents + latents = model.pipeline.vae.encode(defect_images).latent_dist.sample() * model.pipeline.vae.config.scaling_factor + clean_latents_cache[tuple(defect_samples.cpu().tolist())] = latents.detach() + + while global_step < total_steps: + for batch in train_loader: + if global_step >= total_steps: + break + + # t_start = time.time() + + # Move data to device + images = batch["image"].to(device, dtype=torch.float16) + masks = batch["mask"].to(device, dtype=torch.float16) + backgrounds = batch["background"].to(device, dtype=torch.float16) + adjusted_masks = batch["adjusted_mask"].to(device, dtype=torch.float16) + is_defect = batch["is_defect"] + + # Ensure we only process defective samples + defect_samples = torch.nonzero(is_defect).squeeze(1) + if len(defect_samples) == 0: + continue # Skip batch if no defects present + + # Extract defect-only samples + defect_images = images[defect_samples] + defect_masks = masks[defect_samples] + defect_backgrounds = backgrounds[defect_samples] + defect_adjusted_masks = adjusted_masks[defect_samples] + object_classes = [batch["object_class"][i] for i in defect_samples] + + # t_data = time.time() + + # Extract defect type from file path for specific prompting + defect_types = [] + for i in defect_samples: + if hasattr(train_loader.dataset, 'images') and i < len(train_loader.dataset.images): + img_path = train_loader.dataset.images[i] + parts = img_path.split(os.sep) + for j, part in enumerate(parts): + if part == "defective" and j + 1 < len(parts): + defect_types.append(parts[j + 1]) + break + else: + defect_types.append("defect") + else: + defect_types.append("defect") + + # Learning rate warmup + if global_step < args.lr_warmup_steps: + lr_scale = min(1.0, (global_step + 1) / args.lr_warmup_steps) + for i, param_group in enumerate(optimizer.param_groups): + param_group["lr"] = base_lrs[i] * lr_scale + + # Reset attention maps + if hasattr(model, 'attention_maps'): + model.attention_maps = {} + + # ========== PHASE 1: Defect Branch (Defect Texture Learning) ========== + # Using learnable token for concept isolation + # t_text = time.time() + defect_prompts = [f"A photo of {model.placeholder_token}" for _ in range(len(defect_samples))] + text_embeddings = model.get_text_embeddings(defect_prompts, enable_grad=True) + + # ========== PHASE 1 & 2 OPTIMIZED VAE & MASK PASS ========== + + # 1. Single VAE encoding for original defect images + with torch.no_grad(): + # Instead of re-encoding, load from cache + latents = clean_latents_cache[tuple(defect_samples.cpu().tolist())] + + # t_vae = time.time() + + if len(defect_masks.shape) == 3: + defect_masks = defect_masks.unsqueeze(1) + + # 2. Fully Vectorized Random Mask Generation directly on GPU (0 Python Loops) + B, _, H, W = defect_images.shape + K = 15 # Number of random boxes + + # Create broadcastable coordinate grids (1, 1, H, 1) and (1, 1, 1, W) + grid_y = torch.arange(H, device=device).view(1, 1, H, 1) + grid_x = torch.arange(W, device=device).view(1, 1, 1, W) + + # Sample all box parameters for all batch items and box counts simultaneously: shape (B, K, 1, 1) + rh = torch.randint(max(1, int(H * 0.03)), max(2, int(H * 0.25)), (B, K, 1, 1), device=device) + rw = torch.randint(max(1, int(W * 0.03)), max(2, int(W * 0.25)), (B, K, 1, 1), device=device) + ry = torch.randint(0, max(1, H - 20), (B, K, 1, 1), device=device) + rx = torch.randint(0, max(1, W - 20), (B, K, 1, 1), device=device) + + # Compute rectangle coverage across all (B, K, H, W) elements in a single broadcasted operation + in_box = (grid_y >= ry) & (grid_y < (ry + rh)) & (grid_x >= rx) & (grid_x < (rx + rw)) + + # Merge all K boxes per image using logical OR (.any), shape -> (B, 1, H, W) + random_masks = in_box.any(dim=1, keepdim=True).to(dtype=defect_images.dtype) + + # 3. Batched VAE Encoding (Combines Phase 1 and Phase 2 into 1 pass) + with torch.no_grad(): + masked_images = defect_images * (1 - defect_masks) + rand_masked_images = defect_images * (1 - random_masks) + + combined_masked = torch.cat([masked_images, rand_masked_images], dim=0) + combined_latents = model.pipeline.vae.encode(combined_masked).latent_dist.sample() * model.pipeline.vae.config.scaling_factor + + masked_image_latents, random_masked_image_latents = torch.chunk(combined_latents, 2, dim=0) + + # 4. Latent interpolation for masks + mask_latents = F.interpolate(defect_masks, size=(latents.shape[2], latents.shape[3])) + random_mask_latents = F.interpolate(random_masks, size=(latents.shape[2], latents.shape[3])) + + # ========== Phase 1 Forward Pass ========== + noise = torch.randn_like(latents) + timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=device) + noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) + + outputs = model( + noisy_latents=noisy_latents, + masked_image_latents=masked_image_latents, + mask_latents=mask_latents, + timesteps=timesteps, + encoder_hidden_states=text_embeddings + ) + + # t_forward = time.time() + + # print(f"[TIMING RESULTS]") + # print(f" โ”œโ”€ Data Fetch & Filter : {t_data - t_start:.2f}s") + # print(f" โ”œโ”€ Text Encoding : {t_text - t_data:.2f}s") + # print(f" โ”œโ”€ VAE Encoding : {t_vae - t_text:.2f}s") + # print(f" โ””โ”€ UNet Forward & Loss : {t_forward - t_vae:.2f}s") + + + noise_pred = outputs["noise_pred"] + defect_loss = model.compute_defect_loss(noise_pred, noise, mask_latents) + attention_loss = outputs.get("attention_loss", torch.tensor(0.0, device=device)) + + # ========== Phase 2 Forward Pass ========== + obj_prompts = [f"A {obj_class} with {model.placeholder_token}" for obj_class in object_classes] + obj_text_embeddings = model.get_text_embeddings(obj_prompts, enable_grad=True) + + obj_noise = torch.randn_like(latents) + obj_timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=device) + obj_noisy_latents = noise_scheduler.add_noise(latents, obj_noise, obj_timesteps) + + obj_outputs = model( + noisy_latents=obj_noisy_latents, + masked_image_latents=random_masked_image_latents, + mask_latents=random_mask_latents, + timesteps=obj_timesteps, + encoder_hidden_states=obj_text_embeddings + ) + + object_loss = model.compute_object_loss(obj_outputs["noise_pred"], obj_noise, random_mask_latents, alpha=args.alpha) + + # Total Loss Calculation + total_loss = args.lambda_defect * defect_loss + args.lambda_obj * object_loss + args.lambda_attn * attention_loss + total_loss = total_loss / args.gradient_accumulation_steps + + # NaN Check + if torch.isnan(total_loss): + print(f"Warning: NaN loss detected at step {global_step}") + log_file.write(f"Warning: NaN loss at step {global_step}\n") + optimizer.zero_grad() + continue + + # t_before_backward = time.time() + total_loss.backward() + # t_after_backward = time.time() + # print(f" โ””โ”€ Backward Pass Time : {t_after_backward - t_before_backward:.2f}s") + accumulation_step += 1 + + # Optimization step + if accumulation_step >= args.gradient_accumulation_steps: + optimizer.step() + optimizer.zero_grad() + accumulation_step = 0 + + progress_bar.update(1) + global_step += 1 + + # Logging and TensorBoard updates + writer.add_scalar("Loss/Defect", defect_loss.item(), global_step) + writer.add_scalar("Loss/Object", object_loss.item(), global_step) + writer.add_scalar("Loss/Attention", attention_loss.item(), global_step) + writer.add_scalar("Loss/Total", total_loss.item() * args.gradient_accumulation_steps, global_step) + + if global_step % 10 == 0: + for i, param_group in enumerate(optimizer.param_groups): + writer.add_scalar(f"LearningRate/group{i}", param_group["lr"], global_step) + + # Periodic checkpointing + if global_step % args.save_steps == 0 or global_step == total_steps: + checkpoint_path = os.path.join(checkpoints_dir, f"checkpoint_{global_step}.pt") + save_checkpoint(model, optimizer, global_step, checkpoint_path) + log_file.write(f"Checkpoint saved at step {global_step}\n") + + # Save final model + final_checkpoint_path = os.path.join(checkpoints_dir, "checkpoint_final.pt") + save_checkpoint(model, optimizer, global_step, final_checkpoint_path) + print(f"Final model saved to: {final_checkpoint_path}") + + writer.close() + log_file.close() + return model + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Train DefectFill model") + + # Paths + parser.add_argument("--data_dir", type=str, required=True, help="Path to MVTec AD dataset") + parser.add_argument("--object_class", type=str, required=True, help="Object class to train on") + parser.add_argument("--output_dir", type=str, default="./output", help="Directory to save models") + + # Loss Weights + parser.add_argument("--lambda_defect", type=float, default=0.5, help="Defect loss weight (L_def)") + parser.add_argument("--lambda_obj", type=float, default=0.2, help="Object integrity loss weight (L_obj)") + parser.add_argument("--lambda_attn", type=float, default=0.05, help="Attention loss weight (L_attn)") + parser.add_argument("--alpha", type=float, default=0.3, help="Background weight for object branch") + + # Training Config + parser.add_argument("--config_name", type=str, default="base", help="Experiment name (base/tex/obj)") + parser.add_argument("--defect_type", type=str, default=None, help="Specific defect type to train on") + parser.add_argument("--batch_size", type=int, default=2) + parser.add_argument("--lora_rank", type=int, default=8) + parser.add_argument("--lora_alpha", type=int, default=16) + parser.add_argument("--text_encoder_lr", type=float, default=4e-5) + parser.add_argument("--unet_lr", type=float, default=2e-4) + parser.add_argument("--max_train_steps", type=int, default=2000) + parser.add_argument("--lr_warmup_steps", type=int, default=100) + parser.add_argument("--save_steps", type=int, default=500) + parser.add_argument("--seed", type=int, default=-1, help="-1 for timestamp-based seed") + parser.add_argument("--resume_from", type=str, default=None) + parser.add_argument("--gradient_accumulation_steps", type=int, default=2) + parser.add_argument("--dilate_mask", type=str, default="False", help="Whether to dilate masks (True/False)") + parser.add_argument("--mask_kernel_size", type=int, default=3, help="Size of dilation kernel (must be odd, e.g. 3, 5, 7)") + + args = parser.parse_args() + + # Helper to convert string "True" to boolean + args.dilate_mask = args.dilate_mask.lower() == "true" + + # Seed setup + if args.seed == -1: + args.seed = generate_seed_from_timestamp() + + torch.manual_seed(args.seed) + random.seed(args.seed) + + train(args) diff --git a/ArtiAgent - DefectFill/engine/DefectFill/utils.py b/ArtiAgent - DefectFill/engine/DefectFill/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cee815f375a99785bd044c566882b889834b82ba --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/utils.py @@ -0,0 +1,182 @@ +import torch +import torch.nn.functional as F +import os +from typing import Optional, Dict, Any + +def save_checkpoint(model, optimizer, step, path): + """ + Save model checkpoint - Includes LoRA weights and learnable embeddings (Textual Inversion) + + Args: + model: The DefectFill model instance + optimizer: Optimizer state + step: Current training step + path: File path to save the checkpoint + """ + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(path), exist_ok=True) + + # Extract LoRA weights specifically from the UNet and Text Encoder + checkpoint = { + "step": step, + "text_encoder_lora": {k: v for k, v in model.pipeline.text_encoder.state_dict().items() if "lora" in k}, + "unet_lora": {k: v for k, v in model.pipeline.unet.state_dict().items() if "lora" in k}, + "optimizer": optimizer.state_dict() if optimizer is not None else None, + } + + # Save the learnable embedding (Textual Inversion component) + if hasattr(model, 'placeholder_token_id'): + token_embeds = model.pipeline.text_encoder.get_input_embeddings().weight.data + checkpoint["learned_embedding"] = token_embeds[model.placeholder_token_id].clone() + checkpoint["placeholder_token"] = model.placeholder_token + checkpoint["placeholder_token_id"] = model.placeholder_token_id + print(f"[Checkpoint] Saving learnable embedding: {model.placeholder_token} (id={model.placeholder_token_id})") + + torch.save(checkpoint, path) + print(f"Checkpoint saved to {path}") + +def load_checkpoint(model, optimizer, path) -> int: + """ + Load model checkpoint - Restores LoRA weights and learnable embeddings + + Args: + model: The DefectFill model instance + optimizer: Optimizer to load state into + path: Path to the checkpoint file + + Returns: + Current step retrieved from the checkpoint + """ + # Check if checkpoint exists + if not os.path.exists(path): + print(f"Checkpoint {path} not found, starting from scratch") + return 0 + + # Load checkpoint to CPU first to avoid VRAM spikes + checkpoint = torch.load(path, map_location='cpu') + + # Load text encoder LoRA weights + text_encoder_sd = model.pipeline.text_encoder.state_dict() + for k, v in checkpoint["text_encoder_lora"].items(): + if k in text_encoder_sd: + text_encoder_sd[k] = v.to(text_encoder_sd[k].device) + model.pipeline.text_encoder.load_state_dict(text_encoder_sd) + + # Load UNet LoRA weights + unet_sd = model.pipeline.unet.state_dict() + for k, v in checkpoint["unet_lora"].items(): + if k in unet_sd: + unet_sd[k] = v.to(unet_sd[k].device) + model.pipeline.unet.load_state_dict(unet_sd) + + # Load the learnable embedding (Textual Inversion) + if "learned_embedding" in checkpoint and hasattr(model, 'placeholder_token_id'): + learned_emb = checkpoint["learned_embedding"] + token_embeds = model.pipeline.text_encoder.get_input_embeddings().weight.data + token_embeds[model.placeholder_token_id] = learned_emb.to(token_embeds.device) + print(f"[Checkpoint] Loaded learnable embedding: {model.placeholder_token} (id={model.placeholder_token_id})") + + # Load optimizer state if provided + if optimizer is not None and "optimizer" in checkpoint: + optimizer.load_state_dict(checkpoint["optimizer"]) + + print(f"Checkpoint loaded from {path}") + return checkpoint["step"] + + +def compute_spatial_lpips(lpips_model, img1, img2, mask, smooth_boundary=True): + """ + Calculates the Perceptual Distance specifically within the masked region using Spatial LPIPS + + Args: + lpips_model: LPIPS model instance initialized with spatial=True + img1: Reference image [B, 3, H, W], range [-1, 1] + img2: Comparison image [B, 3, H, W], range [-1, 1] + mask: Defect mask [B, 1, H, W], range [0, 1] + smooth_boundary: Whether to blur the mask edges to avoid boundary artifacts + + Returns: + lpips_score: LPIPS score for the masked region (scalar) + """ + # 1. Compute spatial LPIPS map (pixel-wise perceptual distance) + lpips_map = lpips_model(img1, img2) # Output shape: [B, 1, H', W'] + + # 2. Resize the mask to match the LPIPS output resolution + mask_resized = F.interpolate( + mask, + size=lpips_map.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # 3. Optional: Smooth boundary edges (Gaussian-like blur via AvgPool) + if smooth_boundary: + mask_smoothed = F.avg_pool2d( + F.pad(mask_resized, (2, 2, 2, 2), mode='replicate'), + kernel_size=5, stride=1 + ) + else: + mask_smoothed = mask_resized + + # 4. Mask-weighted summation + weighted_sum = (lpips_map * mask_smoothed).sum(dim=(2, 3)) + mask_sum = mask_smoothed.sum(dim=(2, 3)) + 1e-8 + + # 5. Return normalized LPIPS score + return (weighted_sum / mask_sum).mean() + + +def compute_spatial_lpips_batch(lpips_model, reference, samples, mask, smooth_boundary=True): + """ + Batch calculation of Spatial LPIPS - Evaluates all generated samples at once (FP16 compatible) + + This utilizes the "paired batch" mode of LPIPS by expanding the reference image + to match the number of samples, allowing parallel evaluation in a single forward pass. + + Args: + lpips_model: Spatial LPIPS model instance (spatial=True) + reference: Single reference image [1, 3, H, W], range [-1, 1] + samples: Multiple generated samples [N, 3, H, W], range [-1, 1] + mask: Single defect mask [1, 1, H, W], range [0, 1] + smooth_boundary: Enable mask boundary smoothing + + Returns: + lpips_scores: A tensor of [N] LPIPS scores + """ + num_samples = samples.shape[0] + + # LPIPS model expects FP32 input; cast based on model parameters + lpips_dtype = next(lpips_model.parameters()).dtype + + # Expand reference and mask to match sample count N + reference_expanded = reference.repeat(num_samples, 1, 1, 1).to(dtype=lpips_dtype) + samples_for_lpips = samples.to(dtype=lpips_dtype) + mask_expanded = mask.repeat(num_samples, 1, 1, 1) + + # Parallel forward pass for all N pairs + lpips_maps = lpips_model(reference_expanded, samples_for_lpips) # [N, 1, H', W'] + + # Match mask size to LPIPS feature map resolution + mask_resized = F.interpolate( + mask_expanded.to(dtype=lpips_maps.dtype), + size=lpips_maps.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # Boundary smoothing + if smooth_boundary: + mask_smoothed = F.avg_pool2d( + F.pad(mask_resized, (2, 2, 2, 2), mode='replicate'), + kernel_size=5, stride=1 + ) + else: + mask_smoothed = mask_resized + + # Calculate weighted scores per sample + weighted_sum = (lpips_maps * mask_smoothed).sum(dim=(2, 3)) # [N, 1] + mask_sum = mask_smoothed.sum(dim=(2, 3)) + 1e-8 # [N, 1] + + lpips_scores = (weighted_sum / mask_sum).squeeze(1) # [N] + + return lpips_scores \ No newline at end of file diff --git a/ArtiAgent - DefectFill/engine/DefectFill/visualize_results.py b/ArtiAgent - DefectFill/engine/DefectFill/visualize_results.py new file mode 100644 index 0000000000000000000000000000000000000000..d4e9ebddee278f20d9cd018d8cf41a7ec8f7944d --- /dev/null +++ b/ArtiAgent - DefectFill/engine/DefectFill/visualize_results.py @@ -0,0 +1,256 @@ +""" +DefectFill Experiment Results Visualization Module + +Generates the following visualization charts: +1. Heatmaps - Shows config vs class performance matrix +2. Scatter Plot - KID vs IC-LPIPS quality-diversity tradeoff analysis +3. Grouped Bar Charts - Comparison of each config across different classes + +Usage: + python visualize_results.py --csv_path evaluation_results.csv --output_dir ./figures +""" + +import os +import argparse +import pandas as pd +import numpy as np + +# Set non-interactive backend for headless environments (e.g., servers/clusters) +import matplotlib +matplotlib.use('Agg') + +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +from matplotlib.lines import Line2D +import seaborn as sns + +# Global Plotting Configuration +plt.rcParams['axes.unicode_minus'] = False +sns.set_style("whitegrid") +sns.set_context("paper", font_scale=1.2) + +# Visualization Color Scheme +CONFIG_COLORS = { + 'base': '#3498db', # Blue + 'tex': '#e74c3c', # Red + 'obj': '#2ecc71' # Green +} + +# Shapes to distinguish between Object and Texture datasets +CATEGORY_MARKERS = { + 'object': 'o', # Circle + 'texture': 's' # Square +} + +# MVTec AD Dataset Groupings +OBJECT_CLASSES = ['bottle', 'cable', 'hazelnut', 'metal_nut', 'toothbrush'] +TEXTURE_CLASSES = ['carpet', 'grid', 'leather', 'tile', 'wood'] + + +def load_and_validate_data(csv_path): + """Load and validate the evaluation CSV data.""" + if not os.path.exists(csv_path): + raise FileNotFoundError(f"CSV file not found: {csv_path}") + + df = pd.read_csv(csv_path) + + # Required metric columns + required_columns = ['class', 'config', 'category_type', 'KID_mean', 'KID_std', + 'IC_LPIPS_mean', 'IC_LPIPS_std'] + missing_columns = [col for col in required_columns if col not in df.columns] + + if missing_columns: + raise ValueError(f"CSV missing required columns: {missing_columns}") + + print(f"Successfully loaded data: {len(df)} records") + print(f"Classes found: {df['class'].unique().tolist()}") + print(f"Configs found: {df['config'].unique().tolist()}") + + return df + + +def create_heatmaps(df, output_dir): + """ + Generate performance matrices. + 1. KID Heatmap: Quality evaluation (lower is better). + 2. IC-LPIPS Heatmap: Diversity evaluation (higher is better). + """ + print("\nGenerating performance heatmaps...") + + fig, axes = plt.subplots(1, 2, figsize=(16, 8)) + + # === KID Heatmap === + pivot_kid = df.pivot_table(index='class', columns='config', values='KID_mean', aggfunc='mean') + + # Sort by Category (Object classes first, then Texture) + class_order = [c for c in OBJECT_CLASSES if c in pivot_kid.index] + \ + [c for c in TEXTURE_CLASSES if c in pivot_kid.index] + pivot_kid = pivot_kid.reindex(class_order) + + config_order = ['base', 'tex', 'obj'] + pivot_kid = pivot_kid.reindex(columns=[c for c in config_order if c in pivot_kid.columns]) + + ax1 = axes[0] + sns.heatmap(pivot_kid, annot=True, fmt='.4f', cmap='RdYlGn_r', + linewidths=0.5, ax=ax1, + cbar_kws={'label': 'KID (lower is better)', 'shrink': 0.8}, + annot_kws={'size': 11, 'weight': 'bold'}) + ax1.set_title('KID Quality Evaluation Matrix', fontsize=14, fontweight='bold', pad=15) + + # Visual separators for dataset types + n_object = len([c for c in OBJECT_CLASSES if c in class_order]) + if 0 < n_object < len(class_order): + ax1.axhline(y=n_object, color='black', linewidth=2) + + # === IC-LPIPS Heatmap === + pivot_lpips = df.pivot_table(index='class', columns='config', values='IC_LPIPS_mean', aggfunc='mean') + pivot_lpips = pivot_lpips.reindex(class_order) + pivot_lpips = pivot_lpips.reindex(columns=[c for c in config_order if c in pivot_lpips.columns]) + + ax2 = axes[1] + sns.heatmap(pivot_lpips, annot=True, fmt='.4f', cmap='RdYlGn', + linewidths=0.5, ax=ax2, + cbar_kws={'label': 'IC-LPIPS (higher is better)', 'shrink': 0.8}, + annot_kws={'size': 11, 'weight': 'bold'}) + ax2.set_title('IC-LPIPS Diversity Evaluation Matrix', fontsize=14, fontweight='bold', pad=15) + + if 0 < n_object < len(class_order): + ax2.axhline(y=n_object, color='black', linewidth=2) + + plt.tight_layout() + output_path = os.path.join(output_dir, 'heatmaps.png') + plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white') + print(f" Heatmaps saved to: {output_path}") + plt.close() + + +def create_scatter_plot(df, output_dir): + """ + Generates a Quality vs. Diversity Trade-off scatter plot. + Top-left corner represents the "Ideal Region" (Low KID, High Diversity). + """ + print("\nGenerating Quality vs Diversity scatter plot...") + + fig, ax = plt.subplots(figsize=(12, 9)) + + for _, row in df.iterrows(): + color = CONFIG_COLORS.get(row['config'], '#95a5a6') + marker = CATEGORY_MARKERS.get(row['category_type'], 'o') + + ax.scatter(row['KID_mean'], row['IC_LPIPS_mean'], + c=color, marker=marker, s=200, alpha=0.8, + edgecolors='black', linewidth=1.5, zorder=3) + + # Annotate class names + label_text = row['class'][:4] if len(row['class']) > 4 else row['class'] + ax.annotate(label_text, (row['KID_mean'], row['IC_LPIPS_mean']), + fontsize=8, ha='center', va='bottom', xytext=(0, 8), + textcoords='offset points', fontweight='bold') + + # Reference Median Lines + ax.axhline(y=df['IC_LPIPS_mean'].median(), color='gray', linestyle='--', alpha=0.5) + ax.axvline(x=df['KID_mean'].median(), color='gray', linestyle='--', alpha=0.5) + + # Highlight the Ideal Region + kid_min, lpips_max = df['KID_mean'].min(), df['IC_LPIPS_mean'].max() + ax.annotate('Ideal Region\n(High Quality + High Diversity)', + xy=(kid_min, lpips_max), fontsize=11, color='#27ae60', fontweight='bold', + ha='left', va='top', bbox=dict(boxstyle='round,pad=0.3', facecolor='#d5f4e6', edgecolor='#27ae60')) + + ax.set_xlabel('KID (Lower = Higher Fidelity)', fontsize=13, fontweight='bold') + ax.set_ylabel('IC-LPIPS (Higher = More Diverse)', fontsize=13, fontweight='bold') + ax.set_title('Quality vs Diversity Trade-off Analysis', fontsize=15, fontweight='bold', pad=15) + + # Dynamic Legend Generation + legend_elements = [Line2D([0], [0], marker='o', color='w', markerfacecolor=c, markersize=12, label=f'Config {k.upper()}', markeredgecolor='black') + for k, c in CONFIG_COLORS.items() if k in df['config'].values] + + ax.legend(handles=legend_elements, loc='upper right', framealpha=0.95) + ax.grid(True, alpha=0.3) + + plt.tight_layout() + output_path = os.path.join(output_dir, 'scatter_tradeoff.png') + plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white') + print(f" Scatter plot saved to: {output_path}") + plt.close() + + +def create_grouped_bar_charts(df, output_dir): + """Generates comparison bar charts across all classes and configs.""" + print("\nGenerating comparison bar charts...") + + fig, axes = plt.subplots(2, 1, figsize=(16, 12)) + all_classes = df['class'].unique().tolist() + class_order = [c for c in OBJECT_CLASSES if c in all_classes] + \ + [c for c in TEXTURE_CLASSES if c in all_classes] + + config_order = ['base', 'tex', 'obj'] + configs = [c for c in config_order if c in df['config'].values] + + bar_width = 0.25 + x = np.arange(len(class_order)) + + for idx, (metric, ax_title, ylabel) in enumerate([ + ('KID', 'KID Quality Evaluation - Config Comparison', 'KID (lower is better)'), + ('IC_LPIPS', 'IC-LPIPS Diversity Evaluation - Config Comparison', 'IC-LPIPS (higher is better)') + ]): + ax = axes[idx] + for i, config in enumerate(configs): + config_data = df[df['config'] == config].set_index('class') + values = [config_data.loc[c, f'{metric}_mean'] if c in config_data.index else 0 for c in class_order] + errors = [config_data.loc[c, f'{metric}_std'] if c in config_data.index else 0 for c in class_order] + + ax.bar(x + i * bar_width, values, bar_width, label=f'Config {config.upper()}', + color=CONFIG_COLORS.get(config, '#95a5a6'), yerr=errors, capsize=3, + alpha=0.85, edgecolor='black', linewidth=0.5) + + ax.set_title(ax_title, fontsize=14, fontweight='bold') + ax.set_ylabel(ylabel, fontweight='bold') + ax.set_xticks(x + bar_width * (len(configs) - 1) / 2) + ax.set_xticklabels(class_order, rotation=45, ha='right') + ax.legend() + + plt.tight_layout() + output_path = os.path.join(output_dir, 'grouped_bar_charts.png') + plt.savefig(output_path, dpi=300, bbox_inches='tight') + print(f" Bar charts saved to: {output_path}") + plt.close() + + +def create_summary_table(df, output_dir): + """Calculates and saves grouped summary statistics.""" + print("\nCalculating summary statistics...") + + summary = df.groupby(['category_type', 'config']).agg({ + 'KID_mean': ['mean', 'std', 'min', 'max'], + 'IC_LPIPS_mean': ['mean', 'std', 'min', 'max'] + }).round(4) + + summary.columns = ['_'.join(col).strip() for col in summary.columns.values] + summary_path = os.path.join(output_dir, 'summary_statistics.csv') + summary.to_csv(summary_path) + + print("\n" + "="*80 + "\nSummary Statistics:\n" + "="*80) + print(summary.to_string()) + return summary + + +def main(): + parser = argparse.ArgumentParser(description="Visualize DefectFill Experiment Results") + parser.add_argument("--csv_path", type=str, required=True, help="Path to evaluation results CSV") + parser.add_argument("--output_dir", type=str, default="./figures", help="Output directory for plots") + + args = parser.parse_args() + os.makedirs(args.output_dir, exist_ok=True) + + # Process all visualizations + df = load_and_validate_data(args.csv_path) + create_heatmaps(df, args.output_dir) + create_scatter_plot(df, args.output_dir) + create_grouped_bar_charts(df, args.output_dir) + create_summary_table(df, args.output_dir) + + print(f"\nAll visualization charts generated in: {args.output_dir}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/LICENSE b/ArtiAgent - DefectFill/src/GroundingDINO/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..b1395e94b016dd1b95b4c7e3ed493e1d0b342917 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2020 - present, Facebook, Inc + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/README.md b/ArtiAgent - DefectFill/src/GroundingDINO/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b6610df03d409633e572ef49d67a445d35a63967 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/README.md @@ -0,0 +1,163 @@ +# Grounding DINO + +--- + +[![arXiv](https://img.shields.io/badge/arXiv-2303.05499-b31b1b.svg)](https://arxiv.org/abs/2303.05499) +[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/wxWDt5UiwY8) +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/zero-shot-object-detection-with-grounding-dino.ipynb) +[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/cMa77r3YrDk) +[![HuggingFace space](https://img.shields.io/badge/๐Ÿค—-HuggingFace%20Space-cyan.svg)](https://huggingface.co/spaces/ShilongLiu/Grounding_DINO_demo) + +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/zero-shot-object-detection-on-mscoco)](https://paperswithcode.com/sota/zero-shot-object-detection-on-mscoco?p=grounding-dino-marrying-dino-with-grounded) \ +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/zero-shot-object-detection-on-odinw)](https://paperswithcode.com/sota/zero-shot-object-detection-on-odinw?p=grounding-dino-marrying-dino-with-grounded) \ +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/object-detection-on-coco-minival)](https://paperswithcode.com/sota/object-detection-on-coco-minival?p=grounding-dino-marrying-dino-with-grounded) \ +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/grounding-dino-marrying-dino-with-grounded/object-detection-on-coco)](https://paperswithcode.com/sota/object-detection-on-coco?p=grounding-dino-marrying-dino-with-grounded) + + + +Official PyTorch implementation of [Grounding DINO](https://arxiv.org/abs/2303.05499), a stronger open-set object detector. Code is available now! + + +## Highlight + +- **Open-Set Detection.** Detect **everything** with language! +- **High Performancce.** COCO zero-shot **52.5 AP** (training without COCO data!). COCO fine-tune **63.0 AP**. +- **Flexible.** Collaboration with Stable Diffusion for Image Editting. + +## News +[2023/03/28] A YouTube [video](https://youtu.be/cMa77r3YrDk) about Grounding DINO and basic object detection prompt engineering. [[SkalskiP](https://github.com/SkalskiP)] \ +[2023/03/28] Add a [demo](https://huggingface.co/spaces/ShilongLiu/Grounding_DINO_demo) on Hugging Face Space! \ +[2023/03/27] Support CPU-only mode. Now the model can run on machines without GPUs.\ +[2023/03/25] A [demo](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/zero-shot-object-detection-with-grounding-dino.ipynb) for Grounding DINO is available at Colab. [[SkalskiP](https://github.com/SkalskiP)] \ +[2023/03/22] Code is available Now! + +
+ +Description + +ODinW +
+ + + +## TODO + +- [x] Release inference code and demo. +- [x] Release checkpoints. +- [ ] Grounding DINO with Stable Diffusion and GLIGEN demos. +- [ ] Release training codes. + +## Install + +If you have a CUDA environment, please make sure the environment variable `CUDA_HOME` is set. It will be compiled under CPU-only mode if no CUDA available. + +```bash +pip install -e . +``` + +## Demo + +```bash +CUDA_VISIBLE_DEVICES=6 python demo/inference_on_a_image.py \ + -c /path/to/config \ + -p /path/to/checkpoint \ + -i .asset/cats.png \ + -o "outputs/0" \ + -t "cat ear." \ + [--cpu-only] # open it for cpu mode +``` +See the `demo/inference_on_a_image.py` for more details. + +**Web UI** + +We also provide a demo code to integrate Grounding DINO with Gradio Web UI. See the file `demo/gradio_app.py` for more details. + +## Checkpoints + + + + + + + + + + + + + + + + + + + + + + + + + +
namebackboneDatabox AP on COCOCheckpointConfig
1GroundingDINO-TSwin-TO365,GoldG,Cap4M48.4 (zero-shot) / 57.2 (fine-tune)Github link | HF linklink
+ +## Results + +
+ +COCO Object Detection Results + +COCO +
+ +
+ +ODinW Object Detection Results + +ODinW +
+ +
+ +Marrying Grounding DINO with Stable Diffusion for Image Editing + +GD_SD +
+ +
+ +Marrying Grounding DINO with GLIGEN for more Detailed Image Editing + +GD_GLIGEN +
+ +## Model + +Includes: a text backbone, an image backbone, a feature enhancer, a language-guided query selection, and a cross-modality decoder. + +![arch](.asset/arch.png) + + +## Acknowledgement + +Our model is related to [DINO](https://github.com/IDEA-Research/DINO) and [GLIP](https://github.com/microsoft/GLIP). Thanks for their great work! + +We also thank great previous work including DETR, Deformable DETR, SMCA, Conditional DETR, Anchor DETR, Dynamic DETR, DAB-DETR, DN-DETR, etc. More related work are available at [Awesome Detection Transformer](https://github.com/IDEACVR/awesome-detection-transformer). A new toolbox [detrex](https://github.com/IDEA-Research/detrex) is available as well. + +Thanks [Stable Diffusion](https://github.com/Stability-AI/StableDiffusion) and [GLIGEN](https://github.com/gligen/GLIGEN) for their awesome models. + + +## Citation + +If you find our work helpful for your research, please consider citing the following BibTeX entry. + +```bibtex +@inproceedings{ShilongLiu2023GroundingDM, + title={Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection}, + author={Shilong Liu and Zhaoyang Zeng and Tianhe Ren and Feng Li and Hao Zhang and Jie Yang and Chunyuan Li and Jianwei Yang and Hang Su and Jun Zhu and Lei Zhang}, + year={2023} +} +``` + + + + diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/__init__.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5963a5e6797a32718001851e7a31bc82f911ff0c Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinB.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinB.py new file mode 100644 index 0000000000000000000000000000000000000000..f490c4bbd598a35de43d36ceafcbd769e7ff21bf --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinB.py @@ -0,0 +1,43 @@ +batch_size = 1 +modelname = "groundingdino" +backbone = "swin_B_384_22k" +position_embedding = "sine" +pe_temperatureH = 20 +pe_temperatureW = 20 +return_interm_indices = [1, 2, 3] +backbone_freeze_keywords = None +enc_layers = 6 +dec_layers = 6 +pre_norm = False +dim_feedforward = 2048 +hidden_dim = 256 +dropout = 0.0 +nheads = 8 +num_queries = 900 +query_dim = 4 +num_patterns = 0 +num_feature_levels = 4 +enc_n_points = 4 +dec_n_points = 4 +two_stage_type = "standard" +two_stage_bbox_embed_share = False +two_stage_class_embed_share = False +transformer_activation = "relu" +dec_pred_bbox_embed_share = True +dn_box_noise_scale = 1.0 +dn_label_noise_ratio = 0.5 +dn_label_coef = 1.0 +dn_bbox_coef = 1.0 +embed_init_tgt = True +dn_labelbook_size = 2000 +max_text_len = 256 +text_encoder_type = "bert-base-uncased" +use_text_enhancer = True +use_fusion_layer = True +use_checkpoint = True +use_transformer_ckpt = True +use_text_cross_attention = True +text_dropout = 0.0 +fusion_dropout = 0.0 +fusion_droppath = 0.1 +sub_sentence_present = True diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py new file mode 100644 index 0000000000000000000000000000000000000000..9158d5f6260ec74bded95377d382387430d7cd70 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py @@ -0,0 +1,43 @@ +batch_size = 1 +modelname = "groundingdino" +backbone = "swin_T_224_1k" +position_embedding = "sine" +pe_temperatureH = 20 +pe_temperatureW = 20 +return_interm_indices = [1, 2, 3] +backbone_freeze_keywords = None +enc_layers = 6 +dec_layers = 6 +pre_norm = False +dim_feedforward = 2048 +hidden_dim = 256 +dropout = 0.0 +nheads = 8 +num_queries = 900 +query_dim = 4 +num_patterns = 0 +num_feature_levels = 4 +enc_n_points = 4 +dec_n_points = 4 +two_stage_type = "standard" +two_stage_bbox_embed_share = False +two_stage_class_embed_share = False +transformer_activation = "relu" +dec_pred_bbox_embed_share = True +dn_box_noise_scale = 1.0 +dn_label_noise_ratio = 0.5 +dn_label_coef = 1.0 +dn_bbox_coef = 1.0 +embed_init_tgt = True +dn_labelbook_size = 2000 +max_text_len = 256 +text_encoder_type = "bert-base-uncased" +use_text_enhancer = True +use_fusion_layer = True +use_checkpoint = True +use_transformer_ckpt = True +use_text_cross_attention = True +text_dropout = 0.0 +fusion_dropout = 0.0 +fusion_droppath = 0.1 +sub_sentence_present = True diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__init__.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..217e2897203cb3ae56638ee6c9df19512e2cf0c3 Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__pycache__/transforms.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__pycache__/transforms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..292ecd04e8c8a5a73e2248e2cf999de15cc3c092 Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/__pycache__/transforms.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/transforms.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..91cf9269e4b31008a3ddca34a19b038a9b399991 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/datasets/transforms.py @@ -0,0 +1,311 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Transforms and data augmentation for both image + bbox. +""" +import os +import random + +import PIL +import torch +import torchvision.transforms as T +import torchvision.transforms.functional as F + +from groundingdino.util.box_ops import box_xyxy_to_cxcywh +from groundingdino.util.misc import interpolate + + +def crop(image, target, region): + cropped_image = F.crop(image, *region) + + target = target.copy() + i, j, h, w = region + + # should we do something wrt the original size? + target["size"] = torch.tensor([h, w]) + + fields = ["labels", "area", "iscrowd", "positive_map"] + + if "boxes" in target: + boxes = target["boxes"] + max_size = torch.as_tensor([w, h], dtype=torch.float32) + cropped_boxes = boxes - torch.as_tensor([j, i, j, i]) + cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) + cropped_boxes = cropped_boxes.clamp(min=0) + area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) + target["boxes"] = cropped_boxes.reshape(-1, 4) + target["area"] = area + fields.append("boxes") + + if "masks" in target: + # FIXME should we update the area here if there are no boxes? + target["masks"] = target["masks"][:, i : i + h, j : j + w] + fields.append("masks") + + # remove elements for which the boxes or masks that have zero area + if "boxes" in target or "masks" in target: + # favor boxes selection when defining which elements to keep + # this is compatible with previous implementation + if "boxes" in target: + cropped_boxes = target["boxes"].reshape(-1, 2, 2) + keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) + else: + keep = target["masks"].flatten(1).any(1) + + for field in fields: + if field in target: + target[field] = target[field][keep] + + if os.environ.get("IPDB_SHILONG_DEBUG", None) == "INFO": + # for debug and visualization only. + if "strings_positive" in target: + target["strings_positive"] = [ + _i for _i, _j in zip(target["strings_positive"], keep) if _j + ] + + return cropped_image, target + + +def hflip(image, target): + flipped_image = F.hflip(image) + + w, h = image.size + + target = target.copy() + if "boxes" in target: + boxes = target["boxes"] + boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor([-1, 1, -1, 1]) + torch.as_tensor( + [w, 0, w, 0] + ) + target["boxes"] = boxes + + if "masks" in target: + target["masks"] = target["masks"].flip(-1) + + return flipped_image, target + + +def resize(image, target, size, max_size=None): + # size can be min_size (scalar) or (w, h) tuple + + def get_size_with_aspect_ratio(image_size, size, max_size=None): + w, h = image_size + if max_size is not None: + min_original_size = float(min((w, h))) + max_original_size = float(max((w, h))) + if max_original_size / min_original_size * size > max_size: + size = int(round(max_size * min_original_size / max_original_size)) + + if (w <= h and w == size) or (h <= w and h == size): + return (h, w) + + if w < h: + ow = size + oh = int(size * h / w) + else: + oh = size + ow = int(size * w / h) + + return (oh, ow) + + def get_size(image_size, size, max_size=None): + if isinstance(size, (list, tuple)): + return size[::-1] + else: + return get_size_with_aspect_ratio(image_size, size, max_size) + + size = get_size(image.size, size, max_size) + rescaled_image = F.resize(image, size) + + if target is None: + return rescaled_image, None + + ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size)) + ratio_width, ratio_height = ratios + + target = target.copy() + if "boxes" in target: + boxes = target["boxes"] + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height] + ) + target["boxes"] = scaled_boxes + + if "area" in target: + area = target["area"] + scaled_area = area * (ratio_width * ratio_height) + target["area"] = scaled_area + + h, w = size + target["size"] = torch.tensor([h, w]) + + if "masks" in target: + target["masks"] = ( + interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0] > 0.5 + ) + + return rescaled_image, target + + +def pad(image, target, padding): + # assumes that we only pad on the bottom right corners + padded_image = F.pad(image, (0, 0, padding[0], padding[1])) + if target is None: + return padded_image, None + target = target.copy() + # should we do something wrt the original size? + target["size"] = torch.tensor(padded_image.size[::-1]) + if "masks" in target: + target["masks"] = torch.nn.functional.pad(target["masks"], (0, padding[0], 0, padding[1])) + return padded_image, target + + +class ResizeDebug(object): + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + return resize(img, target, self.size) + + +class RandomCrop(object): + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + region = T.RandomCrop.get_params(img, self.size) + return crop(img, target, region) + + +class RandomSizeCrop(object): + def __init__(self, min_size: int, max_size: int, respect_boxes: bool = False): + # respect_boxes: True to keep all boxes + # False to tolerence box filter + self.min_size = min_size + self.max_size = max_size + self.respect_boxes = respect_boxes + + def __call__(self, img: PIL.Image.Image, target: dict): + init_boxes = len(target["boxes"]) + max_patience = 10 + for i in range(max_patience): + w = random.randint(self.min_size, min(img.width, self.max_size)) + h = random.randint(self.min_size, min(img.height, self.max_size)) + region = T.RandomCrop.get_params(img, [h, w]) + result_img, result_target = crop(img, target, region) + if ( + not self.respect_boxes + or len(result_target["boxes"]) == init_boxes + or i == max_patience - 1 + ): + return result_img, result_target + return result_img, result_target + + +class CenterCrop(object): + def __init__(self, size): + self.size = size + + def __call__(self, img, target): + image_width, image_height = img.size + crop_height, crop_width = self.size + crop_top = int(round((image_height - crop_height) / 2.0)) + crop_left = int(round((image_width - crop_width) / 2.0)) + return crop(img, target, (crop_top, crop_left, crop_height, crop_width)) + + +class RandomHorizontalFlip(object): + def __init__(self, p=0.5): + self.p = p + + def __call__(self, img, target): + if random.random() < self.p: + return hflip(img, target) + return img, target + + +class RandomResize(object): + def __init__(self, sizes, max_size=None): + assert isinstance(sizes, (list, tuple)) + self.sizes = sizes + self.max_size = max_size + + def __call__(self, img, target=None): + size = random.choice(self.sizes) + return resize(img, target, size, self.max_size) + + +class RandomPad(object): + def __init__(self, max_pad): + self.max_pad = max_pad + + def __call__(self, img, target): + pad_x = random.randint(0, self.max_pad) + pad_y = random.randint(0, self.max_pad) + return pad(img, target, (pad_x, pad_y)) + + +class RandomSelect(object): + """ + Randomly selects between transforms1 and transforms2, + with probability p for transforms1 and (1 - p) for transforms2 + """ + + def __init__(self, transforms1, transforms2, p=0.5): + self.transforms1 = transforms1 + self.transforms2 = transforms2 + self.p = p + + def __call__(self, img, target): + if random.random() < self.p: + return self.transforms1(img, target) + return self.transforms2(img, target) + + +class ToTensor(object): + def __call__(self, img, target): + return F.to_tensor(img), target + + +class RandomErasing(object): + def __init__(self, *args, **kwargs): + self.eraser = T.RandomErasing(*args, **kwargs) + + def __call__(self, img, target): + return self.eraser(img), target + + +class Normalize(object): + def __init__(self, mean, std): + self.mean = mean + self.std = std + + def __call__(self, image, target=None): + image = F.normalize(image, mean=self.mean, std=self.std) + if target is None: + return image, None + target = target.copy() + h, w = image.shape[-2:] + if "boxes" in target: + boxes = target["boxes"] + boxes = box_xyxy_to_cxcywh(boxes) + boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) + target["boxes"] = boxes + return image, target + + +class Compose(object): + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, image, target): + for t in self.transforms: + image, target = t(image, target) + return image, target + + def __repr__(self): + format_string = self.__class__.__name__ + "(" + for t in self.transforms: + format_string += "\n" + format_string += " {0}".format(t) + format_string += "\n)" + return format_string diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__init__.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2af819d61d589cfec2e0ca46612a7456f42b831a --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__init__.py @@ -0,0 +1,15 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copied from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +from .groundingdino import build_groundingdino diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0140700cc85de677d5fad29bc7bfa8a5e31f9aa4 Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/groundingdino.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/groundingdino.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78b1c51c7da0ecea65b2a8fa4e6bb736d25cb1ea Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/__pycache__/groundingdino.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/__init__.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..76e4b272b479a26c63d120c818c140870cd8c287 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/__init__.py @@ -0,0 +1 @@ +from .backbone import build_backbone diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/backbone.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..c8340c723fad8e07e2fc62daaa3912487498814b --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/backbone.py @@ -0,0 +1,221 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copied from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +""" +Backbone modules. +""" + +from typing import Dict, List + +import torch +import torch.nn.functional as F +import torchvision +from torch import nn +from torchvision.models._utils import IntermediateLayerGetter + +from groundingdino.util.misc import NestedTensor, clean_state_dict, is_main_process + +from .position_encoding import build_position_encoding +from .swin_transformer import build_swin_transformer + + +class FrozenBatchNorm2d(torch.nn.Module): + """ + BatchNorm2d where the batch statistics and the affine parameters are fixed. + + Copy-paste from torchvision.misc.ops with added eps before rqsrt, + without which any other models than torchvision.models.resnet[18,34,50,101] + produce nans. + """ + + def __init__(self, n): + super(FrozenBatchNorm2d, self).__init__() + self.register_buffer("weight", torch.ones(n)) + self.register_buffer("bias", torch.zeros(n)) + self.register_buffer("running_mean", torch.zeros(n)) + self.register_buffer("running_var", torch.ones(n)) + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + num_batches_tracked_key = prefix + "num_batches_tracked" + if num_batches_tracked_key in state_dict: + del state_dict[num_batches_tracked_key] + + super(FrozenBatchNorm2d, self)._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + def forward(self, x): + # move reshapes to the beginning + # to make it fuser-friendly + w = self.weight.reshape(1, -1, 1, 1) + b = self.bias.reshape(1, -1, 1, 1) + rv = self.running_var.reshape(1, -1, 1, 1) + rm = self.running_mean.reshape(1, -1, 1, 1) + eps = 1e-5 + scale = w * (rv + eps).rsqrt() + bias = b - rm * scale + return x * scale + bias + + +class BackboneBase(nn.Module): + def __init__( + self, + backbone: nn.Module, + train_backbone: bool, + num_channels: int, + return_interm_indices: list, + ): + super().__init__() + for name, parameter in backbone.named_parameters(): + if ( + not train_backbone + or "layer2" not in name + and "layer3" not in name + and "layer4" not in name + ): + parameter.requires_grad_(False) + + return_layers = {} + for idx, layer_index in enumerate(return_interm_indices): + return_layers.update( + {"layer{}".format(5 - len(return_interm_indices) + idx): "{}".format(layer_index)} + ) + + # if len: + # if use_stage1_feature: + # return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} + # else: + # return_layers = {"layer2": "0", "layer3": "1", "layer4": "2"} + # else: + # return_layers = {'layer4': "0"} + self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) + self.num_channels = num_channels + + def forward(self, tensor_list: NestedTensor): + xs = self.body(tensor_list.tensors) + out: Dict[str, NestedTensor] = {} + for name, x in xs.items(): + m = tensor_list.mask + assert m is not None + mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] + out[name] = NestedTensor(x, mask) + # import ipdb; ipdb.set_trace() + return out + + +class Backbone(BackboneBase): + """ResNet backbone with frozen BatchNorm.""" + + def __init__( + self, + name: str, + train_backbone: bool, + dilation: bool, + return_interm_indices: list, + batch_norm=FrozenBatchNorm2d, + ): + if name in ["resnet18", "resnet34", "resnet50", "resnet101"]: + backbone = getattr(torchvision.models, name)( + replace_stride_with_dilation=[False, False, dilation], + pretrained=is_main_process(), + norm_layer=batch_norm, + ) + else: + raise NotImplementedError("Why you can get here with name {}".format(name)) + # num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 + assert name not in ("resnet18", "resnet34"), "Only resnet50 and resnet101 are available." + assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]] + num_channels_all = [256, 512, 1024, 2048] + num_channels = num_channels_all[4 - len(return_interm_indices) :] + super().__init__(backbone, train_backbone, num_channels, return_interm_indices) + + +class Joiner(nn.Sequential): + def __init__(self, backbone, position_embedding): + super().__init__(backbone, position_embedding) + + def forward(self, tensor_list: NestedTensor): + xs = self[0](tensor_list) + out: List[NestedTensor] = [] + pos = [] + for name, x in xs.items(): + out.append(x) + # position encoding + pos.append(self[1](x).to(x.tensors.dtype)) + + return out, pos + + +def build_backbone(args): + """ + Useful args: + - backbone: backbone name + - lr_backbone: + - dilation + - return_interm_indices: available: [0,1,2,3], [1,2,3], [3] + - backbone_freeze_keywords: + - use_checkpoint: for swin only for now + + """ + position_embedding = build_position_encoding(args) + train_backbone = True + if not train_backbone: + raise ValueError("Please set lr_backbone > 0") + return_interm_indices = args.return_interm_indices + assert return_interm_indices in [[0, 1, 2, 3], [1, 2, 3], [3]] + args.backbone_freeze_keywords + use_checkpoint = getattr(args, "use_checkpoint", False) + + if args.backbone in ["resnet50", "resnet101"]: + backbone = Backbone( + args.backbone, + train_backbone, + args.dilation, + return_interm_indices, + batch_norm=FrozenBatchNorm2d, + ) + bb_num_channels = backbone.num_channels + elif args.backbone in [ + "swin_T_224_1k", + "swin_B_224_22k", + "swin_B_384_22k", + "swin_L_224_22k", + "swin_L_384_22k", + ]: + pretrain_img_size = int(args.backbone.split("_")[-2]) + backbone = build_swin_transformer( + args.backbone, + pretrain_img_size=pretrain_img_size, + out_indices=tuple(return_interm_indices), + dilation=False, + use_checkpoint=use_checkpoint, + ) + + bb_num_channels = backbone.num_features[4 - len(return_interm_indices) :] + else: + raise NotImplementedError("Unknown backbone {}".format(args.backbone)) + + assert len(bb_num_channels) == len( + return_interm_indices + ), f"len(bb_num_channels) {len(bb_num_channels)} != len(return_interm_indices) {len(return_interm_indices)}" + + model = Joiner(backbone, position_embedding) + model.num_channels = bb_num_channels + assert isinstance( + bb_num_channels, List + ), "bb_num_channels is expected to be a List but {}".format(type(bb_num_channels)) + # import ipdb; ipdb.set_trace() + return model diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/position_encoding.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/position_encoding.py new file mode 100644 index 0000000000000000000000000000000000000000..eac7e896bbe85a670824bfe8ef487d0535d5bd99 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/position_encoding.py @@ -0,0 +1,186 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copied from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +""" +Various positional encodings for the transformer. +""" +import math + +import torch +from torch import nn + +from groundingdino.util.misc import NestedTensor + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + mask = tensor_list.mask + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + # if os.environ.get("SHILONG_AMP", None) == '1': + # eps = 1e-4 + # else: + # eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos + + +class PositionEmbeddingSineHW(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__( + self, num_pos_feats=64, temperatureH=10000, temperatureW=10000, normalize=False, scale=None + ): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperatureH = temperatureH + self.temperatureW = temperatureW + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + mask = tensor_list.mask + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + + # import ipdb; ipdb.set_trace() + + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_tx = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_tx = self.temperatureW ** (2 * (torch.div(dim_tx, 2, rounding_mode='floor')) / self.num_pos_feats) + pos_x = x_embed[:, :, :, None] / dim_tx + + dim_ty = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_ty = self.temperatureH ** (2 * (torch.div(dim_ty, 2, rounding_mode='floor')) / self.num_pos_feats) + pos_y = y_embed[:, :, :, None] / dim_ty + + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + + # import ipdb; ipdb.set_trace() + + return pos + + +class PositionEmbeddingLearned(nn.Module): + """ + Absolute pos embedding, learned. + """ + + def __init__(self, num_pos_feats=256): + super().__init__() + self.row_embed = nn.Embedding(50, num_pos_feats) + self.col_embed = nn.Embedding(50, num_pos_feats) + self.reset_parameters() + + def reset_parameters(self): + nn.init.uniform_(self.row_embed.weight) + nn.init.uniform_(self.col_embed.weight) + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + h, w = x.shape[-2:] + i = torch.arange(w, device=x.device) + j = torch.arange(h, device=x.device) + x_emb = self.col_embed(i) + y_emb = self.row_embed(j) + pos = ( + torch.cat( + [ + x_emb.unsqueeze(0).repeat(h, 1, 1), + y_emb.unsqueeze(1).repeat(1, w, 1), + ], + dim=-1, + ) + .permute(2, 0, 1) + .unsqueeze(0) + .repeat(x.shape[0], 1, 1, 1) + ) + return pos + + +def build_position_encoding(args): + N_steps = args.hidden_dim // 2 + if args.position_embedding in ("v2", "sine"): + # TODO find a better way of exposing other arguments + position_embedding = PositionEmbeddingSineHW( + N_steps, + temperatureH=args.pe_temperatureH, + temperatureW=args.pe_temperatureW, + normalize=True, + ) + elif args.position_embedding in ("v3", "learned"): + position_embedding = PositionEmbeddingLearned(N_steps) + else: + raise ValueError(f"not supported {args.position_embedding}") + + return position_embedding diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..fa8837e4001e41dfed6af99e6619f8b03b989824 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/backbone/swin_transformer.py @@ -0,0 +1,802 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# -------------------------------------------------------- +# modified from https://github.com/SwinTransformer/Swin-Transformer-Object-Detection/blob/master/mmdet/models/backbones/swin_transformer.py +# -------------------------------------------------------- + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +from groundingdino.util.misc import NestedTensor + + +class Mlp(nn.Module): + """Multilayer perceptron.""" + + def __init__( + self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0 + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=0.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """Forward function. + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """Swin Transformer Block. + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__( + self, + dim, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, + window_size=to_2tuple(self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop + ) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """Patch Merging Layer + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """A basic Swin Transformer layer for one stage. + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__( + self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + ): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + Hp = int(np.ceil(H / self.window_size)) * self.window_size + Wp = int(np.ceil(W / self.window_size)) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device, dtype=x.dtype) # 1 Hp Wp 1 + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( + attn_mask == 0, float(0.0) + ) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding + Args: + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + dilation (bool): if True, the output size if 16x downsample, ow 32x downsample. + """ + + def __init__( + self, + pretrain_img_size=224, + patch_size=4, + in_chans=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + dilation=False, + use_checkpoint=False, + ): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + self.dilation = dilation + + # if use_checkpoint: + # print("use_checkpoint!!!!!!!!!!!!!!!!!!!!!!!!") + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [ + pretrain_img_size[0] // patch_size[0], + pretrain_img_size[1] // patch_size[1], + ] + + self.absolute_pos_embed = nn.Parameter( + torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]) + ) + trunc_normal_(self.absolute_pos_embed, std=0.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + # prepare downsample list + downsamplelist = [PatchMerging for i in range(self.num_layers)] + downsamplelist[-1] = None + num_features = [int(embed_dim * 2**i) for i in range(self.num_layers)] + if self.dilation: + downsamplelist[-2] = None + num_features[-1] = int(embed_dim * 2 ** (self.num_layers - 1)) // 2 + for i_layer in range(self.num_layers): + layer = BasicLayer( + # dim=int(embed_dim * 2 ** i_layer), + dim=num_features[i_layer], + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + # downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + downsample=downsamplelist[i_layer], + use_checkpoint=use_checkpoint, + ) + self.layers.append(layer) + + # num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f"norm{i_layer}" + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + # def init_weights(self, pretrained=None): + # """Initialize the weights in backbone. + # Args: + # pretrained (str, optional): Path to pre-trained weights. + # Defaults to None. + # """ + + # def _init_weights(m): + # if isinstance(m, nn.Linear): + # trunc_normal_(m.weight, std=.02) + # if isinstance(m, nn.Linear) and m.bias is not None: + # nn.init.constant_(m.bias, 0) + # elif isinstance(m, nn.LayerNorm): + # nn.init.constant_(m.bias, 0) + # nn.init.constant_(m.weight, 1.0) + + # if isinstance(pretrained, str): + # self.apply(_init_weights) + # logger = get_root_logger() + # load_checkpoint(self, pretrained, strict=False, logger=logger) + # elif pretrained is None: + # self.apply(_init_weights) + # else: + # raise TypeError('pretrained must be a str or None') + + def forward_raw(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate( + self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic" + ) + x = (x + absolute_pos_embed).flatten(2).transpose(1, 2) # B Wh*Ww C + else: + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + + outs = [] + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + # import ipdb; ipdb.set_trace() + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + # in: + # torch.Size([2, 3, 1024, 1024]) + # outs: + # [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \ + # torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])] + return tuple(outs) + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate( + self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic" + ) + x = (x + absolute_pos_embed).flatten(2).transpose(1, 2) # B Wh*Ww C + else: + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + + outs = [] + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + # in: + # torch.Size([2, 3, 1024, 1024]) + # out: + # [torch.Size([2, 192, 256, 256]), torch.Size([2, 384, 128, 128]), \ + # torch.Size([2, 768, 64, 64]), torch.Size([2, 1536, 32, 32])] + + # collect for nesttensors + outs_dict = {} + for idx, out_i in enumerate(outs): + m = tensor_list.mask + assert m is not None + mask = F.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0] + outs_dict[idx] = NestedTensor(out_i, mask) + + return outs_dict + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + + +def build_swin_transformer(modelname, pretrain_img_size, **kw): + assert modelname in [ + "swin_T_224_1k", + "swin_B_224_22k", + "swin_B_384_22k", + "swin_L_224_22k", + "swin_L_384_22k", + ] + + model_para_dict = { + "swin_T_224_1k": dict( + embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7 + ), + "swin_B_224_22k": dict( + embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=7 + ), + "swin_B_384_22k": dict( + embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12 + ), + "swin_L_224_22k": dict( + embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=7 + ), + "swin_L_384_22k": dict( + embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12 + ), + } + kw_cgf = model_para_dict[modelname] + kw_cgf.update(kw) + model = SwinTransformer(pretrain_img_size=pretrain_img_size, **kw_cgf) + return model + + +if __name__ == "__main__": + model = build_swin_transformer("swin_L_384_22k", 384, dilation=True) + x = torch.rand(2, 3, 1024, 1024) + y = model.forward_raw(x) + import ipdb + + ipdb.set_trace() + x = torch.rand(2, 3, 384, 384) + y = model.forward_raw(x) diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/bertwarper.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/bertwarper.py new file mode 100644 index 0000000000000000000000000000000000000000..f0cf9779b270e1aead32845006f8b881fcba37ad --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/bertwarper.py @@ -0,0 +1,273 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from torch import Tensor, nn +from torchvision.ops.boxes import nms +from transformers import BertConfig, BertModel, BertPreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions + + +class BertModelWarper(nn.Module): + def __init__(self, bert_model): + super().__init__() + # self.bert = bert_modelc + + self.config = bert_model.config + self.embeddings = bert_model.embeddings + self.encoder = bert_model.encoder + self.pooler = bert_model.pooler + + self.get_extended_attention_mask = bert_model.get_extended_attention_mask + self.invert_attention_mask = bert_model.invert_attention_mask + self.get_head_mask = bert_model.get_head_mask + + def forward( + self, + input_ids=None, + attention_mask=None, + token_type_ids=None, + position_ids=None, + head_mask=None, + inputs_embeds=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + r""" + encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). + """ + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + batch_size, seq_length = input_shape + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + batch_size, seq_length = input_shape + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + # past_key_values_length + past_key_values_length = ( + past_key_values[0][0].shape[2] if past_key_values is not None else 0 + ) + + if attention_mask is None: + attention_mask = torch.ones( + ((batch_size, seq_length + past_key_values_length)), device=device + ) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask( + attention_mask, input_shape, device + ) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + +class TextEncoderShell(nn.Module): + def __init__(self, text_encoder): + super().__init__() + self.text_encoder = text_encoder + self.config = self.text_encoder.config + + def forward(self, **kw): + # feed into text encoder + return self.text_encoder(**kw) + + +def generate_masks_with_special_tokens(tokenized, special_tokens_list, tokenizer): + """Generate attention mask between each pair of special tokens + Args: + input_ids (torch.Tensor): input ids. Shape: [bs, num_token] + special_tokens_mask (list): special tokens mask. + Returns: + torch.Tensor: attention mask between each special tokens. + """ + input_ids = tokenized["input_ids"] + bs, num_token = input_ids.shape + # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens + special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool() + for special_token in special_tokens_list: + special_tokens_mask |= input_ids == special_token + + # idxs: each row is a list of indices of special tokens + idxs = torch.nonzero(special_tokens_mask) + + # generate attention mask and positional ids + attention_mask = ( + torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1) + ) + position_ids = torch.zeros((bs, num_token), device=input_ids.device) + previous_col = 0 + for i in range(idxs.shape[0]): + row, col = idxs[i] + if (col == 0) or (col == num_token - 1): + attention_mask[row, col, col] = True + position_ids[row, col] = 0 + else: + attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True + position_ids[row, previous_col + 1 : col + 1] = torch.arange( + 0, col - previous_col, device=input_ids.device + ) + + previous_col = col + + # # padding mask + # padding_mask = tokenized['attention_mask'] + # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool() + + return attention_mask, position_ids.to(torch.long) + + +def generate_masks_with_special_tokens_and_transfer_map(tokenized, special_tokens_list, tokenizer): + """Generate attention mask between each pair of special tokens + Args: + input_ids (torch.Tensor): input ids. Shape: [bs, num_token] + special_tokens_mask (list): special tokens mask. + Returns: + torch.Tensor: attention mask between each special tokens. + """ + input_ids = tokenized["input_ids"] + bs, num_token = input_ids.shape + # special_tokens_mask: bs, num_token. 1 for special tokens. 0 for normal tokens + special_tokens_mask = torch.zeros((bs, num_token), device=input_ids.device).bool() + for special_token in special_tokens_list: + special_tokens_mask |= input_ids == special_token + + # idxs: each row is a list of indices of special tokens + idxs = torch.nonzero(special_tokens_mask) + + # generate attention mask and positional ids + attention_mask = ( + torch.eye(num_token, device=input_ids.device).bool().unsqueeze(0).repeat(bs, 1, 1) + ) + position_ids = torch.zeros((bs, num_token), device=input_ids.device) + cate_to_token_mask_list = [[] for _ in range(bs)] + previous_col = 0 + for i in range(idxs.shape[0]): + row, col = idxs[i] + if (col == 0) or (col == num_token - 1): + attention_mask[row, col, col] = True + position_ids[row, col] = 0 + else: + attention_mask[row, previous_col + 1 : col + 1, previous_col + 1 : col + 1] = True + position_ids[row, previous_col + 1 : col + 1] = torch.arange( + 0, col - previous_col, device=input_ids.device + ) + c2t_maski = torch.zeros((num_token), device=input_ids.device).bool() + c2t_maski[previous_col + 1 : col] = True + cate_to_token_mask_list[row].append(c2t_maski) + previous_col = col + + cate_to_token_mask_list = [ + torch.stack(cate_to_token_mask_listi, dim=0) + for cate_to_token_mask_listi in cate_to_token_mask_list + ] + + # # padding mask + # padding_mask = tokenized['attention_mask'] + # attention_mask = attention_mask & padding_mask.unsqueeze(1).bool() & padding_mask.unsqueeze(2).bool() + + return attention_mask, position_ids.to(torch.long), cate_to_token_mask_list diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn.h b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn.h new file mode 100644 index 0000000000000000000000000000000000000000..c7408eba007b424194618baa63726657e36875e3 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn.h @@ -0,0 +1,64 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once + +#include "ms_deform_attn_cpu.h" + +#ifdef WITH_CUDA +#include "ms_deform_attn_cuda.h" +#endif + +namespace groundingdino { + +at::Tensor +ms_deform_attn_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + if (value.type().is_cuda()) + { +#ifdef WITH_CUDA + return ms_deform_attn_cuda_forward( + value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step); +#else + AT_ERROR("Not compiled with GPU support"); +#endif + } + AT_ERROR("Not implemented on the CPU"); +} + +std::vector +ms_deform_attn_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + if (value.type().is_cuda()) + { +#ifdef WITH_CUDA + return ms_deform_attn_cuda_backward( + value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step); +#else + AT_ERROR("Not compiled with GPU support"); +#endif + } + AT_ERROR("Not implemented on the CPU"); +} + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.cpp b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.cpp new file mode 100644 index 0000000000000000000000000000000000000000..551243fdadfd1682b5dc6628623b67a79b3f6c74 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.cpp @@ -0,0 +1,43 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include + +#include +#include + +namespace groundingdino { + +at::Tensor +ms_deform_attn_cpu_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ERROR("Not implement on cpu"); +} + +std::vector +ms_deform_attn_cpu_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + AT_ERROR("Not implement on cpu"); +} + +} // namespace groundingdino diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.h b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.h new file mode 100644 index 0000000000000000000000000000000000000000..b2b88e8c46f19b6db0933163e57ccdb51180f517 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cpu.h @@ -0,0 +1,35 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once +#include + +namespace groundingdino { + +at::Tensor +ms_deform_attn_cpu_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step); + +std::vector +ms_deform_attn_cpu_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step); + +} // namespace groundingdino diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..d04fae8a9a45c11e4e74f3035e94762796da4096 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.cu @@ -0,0 +1,156 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#include +#include "ms_deform_im2col_cuda.cuh" + +#include +#include +#include +#include + +namespace groundingdino { + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step) +{ + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto output = at::zeros({batch, num_query, num_heads, channels}, value.options()); + + const int batch_n = im2col_step_; + auto output_n = output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto columns = output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] { + ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + columns.data()); + + })); + } + + output = output.view({batch, num_query, num_heads*channels}); + + return output; +} + + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step) +{ + + AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous"); + AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous"); + AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous"); + AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous"); + + AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor"); + AT_ASSERTM(spatial_shapes.type().is_cuda(), "spatial_shapes must be a CUDA tensor"); + AT_ASSERTM(level_start_index.type().is_cuda(), "level_start_index must be a CUDA tensor"); + AT_ASSERTM(sampling_loc.type().is_cuda(), "sampling_loc must be a CUDA tensor"); + AT_ASSERTM(attn_weight.type().is_cuda(), "attn_weight must be a CUDA tensor"); + AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, im2col_step); + + AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); + + auto grad_value = at::zeros_like(value); + auto grad_sampling_loc = at::zeros_like(sampling_loc); + auto grad_attn_weight = at::zeros_like(attn_weight); + + const int batch_n = im2col_step_; + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); + + for (int n = 0; n < batch/im2col_step_; ++n) + { + auto grad_output_g = grad_output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] { + ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(), + grad_output_g.data(), + value.data() + n * im2col_step_ * per_value_size, + spatial_shapes.data(), + level_start_index.data(), + sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + attn_weight.data() + n * im2col_step_ * per_attn_weight_size, + batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point, + grad_value.data() + n * im2col_step_ * per_value_size, + grad_sampling_loc.data() + n * im2col_step_ * per_sample_loc_size, + grad_attn_weight.data() + n * im2col_step_ * per_attn_weight_size); + + })); + } + + return { + grad_value, grad_sampling_loc, grad_attn_weight + }; +} + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.h b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.h new file mode 100644 index 0000000000000000000000000000000000000000..ad1311a78f61303616504eb991aaa9c4a93d9948 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_attn_cuda.h @@ -0,0 +1,33 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +#pragma once +#include + +namespace groundingdino { + +at::Tensor ms_deform_attn_cuda_forward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const int im2col_step); + +std::vector ms_deform_attn_cuda_backward( + const at::Tensor &value, + const at::Tensor &spatial_shapes, + const at::Tensor &level_start_index, + const at::Tensor &sampling_loc, + const at::Tensor &attn_weight, + const at::Tensor &grad_output, + const int im2col_step); + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_im2col_cuda.cuh b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_im2col_cuda.cuh new file mode 100644 index 0000000000000000000000000000000000000000..6bc2acb7aea0eab2e9e91e769a16861e1652c284 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/MsDeformAttn/ms_deform_im2col_cuda.cuh @@ -0,0 +1,1327 @@ +/*! +************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************** +* Modified from DCN (https://github.com/msracver/Deformable-ConvNets) +* Copyright (c) 2018 Microsoft +************************************************************************** +*/ + +#include +#include +#include + +#include +#include + +#include + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +const int CUDA_NUM_THREADS = 1024; +inline int GET_BLOCKS(const int N, const int num_threads) +{ + return (N + num_threads - 1) / num_threads; +} + + +template +__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + } + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + return val; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + *grad_attn_weight = top_grad * val; + *grad_sampling_loc = width * grad_w_weight * top_grad_value; + *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + atomicAdd(grad_attn_weight, top_grad * val); + atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value); + atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value); +} + + +template +__global__ void ms_deformable_im2col_gpu_kernel(const int n, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *data_col) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + scalar_t *data_col_ptr = data_col + index; + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + scalar_t col = 0; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride); + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight; + } + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + } + } + *data_col_ptr = col; + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockSize; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockSize/2; s>0; s>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockDim.x; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]); + atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]); + atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]); + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear_gm( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + grad_sampling_loc, grad_attn_weight); + } + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +void ms_deformable_im2col_cuda(cudaStream_t stream, + const scalar_t* data_value, + const int64_t* data_spatial_shapes, + const int64_t* data_level_start_index, + const scalar_t* data_sampling_loc, + const scalar_t* data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* data_col) +{ + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + const int num_threads = CUDA_NUM_THREADS; + ms_deformable_im2col_gpu_kernel + <<>>( + num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, + batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +template +void ms_deformable_col2im_cuda(cudaStream_t stream, + const scalar_t* grad_col, + const scalar_t* data_value, + const int64_t * data_spatial_shapes, + const int64_t * data_level_start_index, + const scalar_t * data_sampling_loc, + const scalar_t * data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels; + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + if (channels > 1024) + { + if ((channels & 1023) == 0) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_gm + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + else{ + switch(channels) + { + case 1: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 2: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 4: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 8: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 16: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 32: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 64: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 128: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 256: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 512: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 1024: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + default: + if (channels < 64) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err)); + } + +} \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/cuda_version.cu b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/cuda_version.cu new file mode 100644 index 0000000000000000000000000000000000000000..64569e34ffb250964de27e33e7a53f3822270b9e --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/cuda_version.cu @@ -0,0 +1,7 @@ +#include + +namespace groundingdino { +int get_cudart_version() { + return CUDART_VERSION; +} +} // namespace groundingdino diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/vision.cpp b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/vision.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c1f2c50c82909bbd5492c163d634af77a3ba1781 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/csrc/vision.cpp @@ -0,0 +1,58 @@ +// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved + +#include "MsDeformAttn/ms_deform_attn.h" + +namespace groundingdino { + +#ifdef WITH_CUDA +extern int get_cudart_version(); +#endif + +std::string get_cuda_version() { +#ifdef WITH_CUDA + std::ostringstream oss; + + // copied from + // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/detail/CUDAHooks.cpp#L231 + auto printCudaStyleVersion = [&](int v) { + oss << (v / 1000) << "." << (v / 10 % 100); + if (v % 10 != 0) { + oss << "." << (v % 10); + } + }; + printCudaStyleVersion(get_cudart_version()); + return oss.str(); +#else + return std::string("not available"); +#endif +} + +// similar to +// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Version.cpp +std::string get_compiler_version() { + std::ostringstream ss; +#if defined(__GNUC__) +#ifndef __clang__ + { ss << "GCC " << __GNUC__ << "." << __GNUC_MINOR__; } +#endif +#endif + +#if defined(__clang_major__) + { + ss << "clang " << __clang_major__ << "." << __clang_minor__ << "." + << __clang_patchlevel__; + } +#endif + +#if defined(_MSC_VER) + { ss << "MSVC " << _MSC_FULL_VER; } +#endif + return ss.str(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward"); + m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward"); +} + +} // namespace groundingdino \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..2753b3ddee43c7a9fe28d1824db5d786e7e1ad59 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/fuse_modules.py @@ -0,0 +1,297 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.models.layers import DropPath + + +class FeatureResizer(nn.Module): + """ + This class takes as input a set of embeddings of dimension C1 and outputs a set of + embedding of dimension C2, after a linear transformation, dropout and normalization (LN). + """ + + def __init__(self, input_feat_size, output_feat_size, dropout, do_ln=True): + super().__init__() + self.do_ln = do_ln + # Object feature encoding + self.fc = nn.Linear(input_feat_size, output_feat_size, bias=True) + self.layer_norm = nn.LayerNorm(output_feat_size, eps=1e-12) + self.dropout = nn.Dropout(dropout) + + def forward(self, encoder_features): + x = self.fc(encoder_features) + if self.do_ln: + x = self.layer_norm(x) + output = self.dropout(x) + return output + + +def l1norm(X, dim, eps=1e-8): + """L1-normalize columns of X""" + norm = torch.abs(X).sum(dim=dim, keepdim=True) + eps + X = torch.div(X, norm) + return X + + +def l2norm(X, dim, eps=1e-8): + """L2-normalize columns of X""" + norm = torch.pow(X, 2).sum(dim=dim, keepdim=True).sqrt() + eps + X = torch.div(X, norm) + return X + + +def func_attention(query, context, smooth=1, raw_feature_norm="softmax", eps=1e-8): + """ + query: (n_context, queryL, d) + context: (n_context, sourceL, d) + """ + batch_size_q, queryL = query.size(0), query.size(1) + batch_size, sourceL = context.size(0), context.size(1) + + # Get attention + # --> (batch, d, queryL) + queryT = torch.transpose(query, 1, 2) + + # (batch, sourceL, d)(batch, d, queryL) + # --> (batch, sourceL, queryL) + attn = torch.bmm(context, queryT) + if raw_feature_norm == "softmax": + # --> (batch*sourceL, queryL) + attn = attn.view(batch_size * sourceL, queryL) + attn = nn.Softmax()(attn) + # --> (batch, sourceL, queryL) + attn = attn.view(batch_size, sourceL, queryL) + elif raw_feature_norm == "l2norm": + attn = l2norm(attn, 2) + elif raw_feature_norm == "clipped_l2norm": + attn = nn.LeakyReLU(0.1)(attn) + attn = l2norm(attn, 2) + else: + raise ValueError("unknown first norm type:", raw_feature_norm) + # --> (batch, queryL, sourceL) + attn = torch.transpose(attn, 1, 2).contiguous() + # --> (batch*queryL, sourceL) + attn = attn.view(batch_size * queryL, sourceL) + attn = nn.Softmax()(attn * smooth) + # --> (batch, queryL, sourceL) + attn = attn.view(batch_size, queryL, sourceL) + # --> (batch, sourceL, queryL) + attnT = torch.transpose(attn, 1, 2).contiguous() + + # --> (batch, d, sourceL) + contextT = torch.transpose(context, 1, 2) + # (batch x d x sourceL)(batch x sourceL x queryL) + # --> (batch, d, queryL) + weightedContext = torch.bmm(contextT, attnT) + # --> (batch, queryL, d) + weightedContext = torch.transpose(weightedContext, 1, 2) + + return weightedContext, attnT + + +class BiMultiHeadAttention(nn.Module): + def __init__(self, v_dim, l_dim, embed_dim, num_heads, dropout=0.1, cfg=None): + super(BiMultiHeadAttention, self).__init__() + + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.v_dim = v_dim + self.l_dim = l_dim + + assert ( + self.head_dim * self.num_heads == self.embed_dim + ), f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})." + self.scale = self.head_dim ** (-0.5) + self.dropout = dropout + + self.v_proj = nn.Linear(self.v_dim, self.embed_dim) + self.l_proj = nn.Linear(self.l_dim, self.embed_dim) + self.values_v_proj = nn.Linear(self.v_dim, self.embed_dim) + self.values_l_proj = nn.Linear(self.l_dim, self.embed_dim) + + self.out_v_proj = nn.Linear(self.embed_dim, self.v_dim) + self.out_l_proj = nn.Linear(self.embed_dim, self.l_dim) + + self.stable_softmax_2d = True + self.clamp_min_for_underflow = True + self.clamp_max_for_overflow = True + + self._reset_parameters() + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def _reset_parameters(self): + nn.init.xavier_uniform_(self.v_proj.weight) + self.v_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.l_proj.weight) + self.l_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.values_v_proj.weight) + self.values_v_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.values_l_proj.weight) + self.values_l_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.out_v_proj.weight) + self.out_v_proj.bias.data.fill_(0) + nn.init.xavier_uniform_(self.out_l_proj.weight) + self.out_l_proj.bias.data.fill_(0) + + def forward(self, v, l, attention_mask_v=None, attention_mask_l=None): + """_summary_ + + Args: + v (_type_): bs, n_img, dim + l (_type_): bs, n_text, dim + attention_mask_v (_type_, optional): _description_. bs, n_img + attention_mask_l (_type_, optional): _description_. bs, n_text + + Returns: + _type_: _description_ + """ + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + bsz, tgt_len, _ = v.size() + + query_states = self.v_proj(v) * self.scale + key_states = self._shape(self.l_proj(l), -1, bsz) + value_v_states = self._shape(self.values_v_proj(v), -1, bsz) + value_l_states = self._shape(self.values_l_proj(l), -1, bsz) + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + key_states = key_states.view(*proj_shape) + value_v_states = value_v_states.view(*proj_shape) + value_l_states = value_l_states.view(*proj_shape) + + src_len = key_states.size(1) + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) # bs*nhead, nimg, ntxt + + if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is {attn_weights.size()}" + ) + + if self.stable_softmax_2d: + attn_weights = attn_weights - attn_weights.max() + + if self.clamp_min_for_underflow: + attn_weights = torch.clamp( + attn_weights, min=-50000 + ) # Do not increase -50000, data type half has quite limited range + if self.clamp_max_for_overflow: + attn_weights = torch.clamp( + attn_weights, max=50000 + ) # Do not increase 50000, data type half has quite limited range + + attn_weights_T = attn_weights.transpose(1, 2) + attn_weights_l = attn_weights_T - torch.max(attn_weights_T, dim=-1, keepdim=True)[0] + if self.clamp_min_for_underflow: + attn_weights_l = torch.clamp( + attn_weights_l, min=-50000 + ) # Do not increase -50000, data type half has quite limited range + if self.clamp_max_for_overflow: + attn_weights_l = torch.clamp( + attn_weights_l, max=50000 + ) # Do not increase 50000, data type half has quite limited range + + # mask vison for language + if attention_mask_v is not None: + attention_mask_v = ( + attention_mask_v[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1) + ) + attn_weights_l.masked_fill_(attention_mask_v, float("-inf")) + + attn_weights_l = attn_weights_l.softmax(dim=-1) + + # mask language for vision + if attention_mask_l is not None: + attention_mask_l = ( + attention_mask_l[:, None, None, :].repeat(1, self.num_heads, 1, 1).flatten(0, 1) + ) + attn_weights.masked_fill_(attention_mask_l, float("-inf")) + attn_weights_v = attn_weights.softmax(dim=-1) + + attn_probs_v = F.dropout(attn_weights_v, p=self.dropout, training=self.training) + attn_probs_l = F.dropout(attn_weights_l, p=self.dropout, training=self.training) + + attn_output_v = torch.bmm(attn_probs_v, value_l_states) + attn_output_l = torch.bmm(attn_probs_l, value_v_states) + + if attn_output_v.size() != (bsz * self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output_v` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is {attn_output_v.size()}" + ) + + if attn_output_l.size() != (bsz * self.num_heads, src_len, self.head_dim): + raise ValueError( + f"`attn_output_l` should be of size {(bsz, self.num_heads, src_len, self.head_dim)}, but is {attn_output_l.size()}" + ) + + attn_output_v = attn_output_v.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output_v = attn_output_v.transpose(1, 2) + attn_output_v = attn_output_v.reshape(bsz, tgt_len, self.embed_dim) + + attn_output_l = attn_output_l.view(bsz, self.num_heads, src_len, self.head_dim) + attn_output_l = attn_output_l.transpose(1, 2) + attn_output_l = attn_output_l.reshape(bsz, src_len, self.embed_dim) + + attn_output_v = self.out_v_proj(attn_output_v) + attn_output_l = self.out_l_proj(attn_output_l) + + return attn_output_v, attn_output_l + + +# Bi-Direction MHA (text->image, image->text) +class BiAttentionBlock(nn.Module): + def __init__( + self, + v_dim, + l_dim, + embed_dim, + num_heads, + dropout=0.1, + drop_path=0.0, + init_values=1e-4, + cfg=None, + ): + """ + Inputs: + embed_dim - Dimensionality of input and attention feature vectors + hidden_dim - Dimensionality of hidden layer in feed-forward network + (usually 2-4x larger than embed_dim) + num_heads - Number of heads to use in the Multi-Head Attention block + dropout - Amount of dropout to apply in the feed-forward network + """ + super(BiAttentionBlock, self).__init__() + + # pre layer norm + self.layer_norm_v = nn.LayerNorm(v_dim) + self.layer_norm_l = nn.LayerNorm(l_dim) + self.attn = BiMultiHeadAttention( + v_dim=v_dim, l_dim=l_dim, embed_dim=embed_dim, num_heads=num_heads, dropout=dropout + ) + + # add layer scale for training stability + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.gamma_v = nn.Parameter(init_values * torch.ones((v_dim)), requires_grad=True) + self.gamma_l = nn.Parameter(init_values * torch.ones((l_dim)), requires_grad=True) + + def forward(self, v, l, attention_mask_v=None, attention_mask_l=None): + v = self.layer_norm_v(v) + l = self.layer_norm_l(l) + delta_v, delta_l = self.attn( + v, l, attention_mask_v=attention_mask_v, attention_mask_l=attention_mask_l + ) + # v, l = v + delta_v, l + delta_l + v = v + self.drop_path(self.gamma_v * delta_v) + l = l + self.drop_path(self.gamma_l * delta_l) + return v, l + + # def forward(self, v:List[torch.Tensor], l, attention_mask_v=None, attention_mask_l=None) diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/groundingdino.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/groundingdino.py new file mode 100644 index 0000000000000000000000000000000000000000..134cadaac010a2b23fe407fa304c71a08d6da206 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/groundingdino.py @@ -0,0 +1,398 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR model and criterion classes. +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Modified from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ +# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR) +# Copyright (c) 2020 SenseTime. All Rights Reserved. +# ------------------------------------------------------------------------ +import copy +from typing import List + +import torch +import torch.nn.functional as F +from torch import nn +from torchvision.ops.boxes import nms +from transformers import AutoTokenizer, BertModel, BertTokenizer, RobertaModel, RobertaTokenizerFast + +from groundingdino.util import box_ops, get_tokenlizer +from groundingdino.util.misc import ( + NestedTensor, + accuracy, + get_world_size, + interpolate, + inverse_sigmoid, + is_dist_avail_and_initialized, + nested_tensor_from_tensor_list, +) +from groundingdino.util.utils import get_phrases_from_posmap +from groundingdino.util.visualizer import COCOVisualizer +from groundingdino.util.vl_utils import create_positive_map_from_span + +from ..registry import MODULE_BUILD_FUNCS +from .backbone import build_backbone +from .bertwarper import ( + BertModelWarper, + generate_masks_with_special_tokens, + generate_masks_with_special_tokens_and_transfer_map, +) +from .transformer import build_transformer +from .utils import MLP, ContrastiveEmbed, sigmoid_focal_loss + + +class GroundingDINO(nn.Module): + """This is the Cross-Attention Detector module that performs object detection""" + + def __init__( + self, + backbone, + transformer, + num_queries, + bert_base_uncased_path, + aux_loss=False, + iter_update=False, + query_dim=2, + num_feature_levels=1, + nheads=8, + # two stage + two_stage_type="no", # ['no', 'standard'] + dec_pred_bbox_embed_share=True, + two_stage_class_embed_share=True, + two_stage_bbox_embed_share=True, + num_patterns=0, + dn_number=100, + dn_box_noise_scale=0.4, + dn_label_noise_ratio=0.5, + dn_labelbook_size=100, + text_encoder_type="bert-base-uncased", + sub_sentence_present=True, + max_text_len=256, + ): + """Initializes the model. + Parameters: + backbone: torch module of the backbone to be used. See backbone.py + transformer: torch module of the transformer architecture. See transformer.py + num_queries: number of object queries, ie detection slot. This is the maximal number of objects + Conditional DETR can detect in a single image. For COCO, we recommend 100 queries. + aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. + """ + super().__init__() + self.num_queries = num_queries + self.transformer = transformer + self.hidden_dim = hidden_dim = transformer.d_model + self.num_feature_levels = num_feature_levels + self.nheads = nheads + self.max_text_len = 256 + self.sub_sentence_present = sub_sentence_present + + # setting query dim + self.query_dim = query_dim + assert query_dim == 4 + + # for dn training + self.num_patterns = num_patterns + self.dn_number = dn_number + self.dn_box_noise_scale = dn_box_noise_scale + self.dn_label_noise_ratio = dn_label_noise_ratio + self.dn_labelbook_size = dn_labelbook_size + + # bert + self.tokenizer = get_tokenlizer.get_tokenlizer(text_encoder_type, bert_base_uncased_path) + self.bert = get_tokenlizer.get_pretrained_language_model(text_encoder_type, bert_base_uncased_path) + self.bert.pooler.dense.weight.requires_grad_(False) + self.bert.pooler.dense.bias.requires_grad_(False) + self.bert = BertModelWarper(bert_model=self.bert) + + self.feat_map = nn.Linear(self.bert.config.hidden_size, self.hidden_dim, bias=True) + nn.init.constant_(self.feat_map.bias.data, 0) + nn.init.xavier_uniform_(self.feat_map.weight.data) + # freeze + + # special tokens + self.specical_tokens = self.tokenizer.convert_tokens_to_ids(["[CLS]", "[SEP]", ".", "?"]) + + # prepare input projection layers + if num_feature_levels > 1: + num_backbone_outs = len(backbone.num_channels) + input_proj_list = [] + for _ in range(num_backbone_outs): + in_channels = backbone.num_channels[_] + input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, hidden_dim, kernel_size=1), + nn.GroupNorm(32, hidden_dim), + ) + ) + for _ in range(num_feature_levels - num_backbone_outs): + input_proj_list.append( + nn.Sequential( + nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1), + nn.GroupNorm(32, hidden_dim), + ) + ) + in_channels = hidden_dim + self.input_proj = nn.ModuleList(input_proj_list) + else: + assert two_stage_type == "no", "two_stage_type should be no if num_feature_levels=1 !!!" + self.input_proj = nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d(backbone.num_channels[-1], hidden_dim, kernel_size=1), + nn.GroupNorm(32, hidden_dim), + ) + ] + ) + + self.backbone = backbone + self.aux_loss = aux_loss + self.box_pred_damping = box_pred_damping = None + + self.iter_update = iter_update + assert iter_update, "Why not iter_update?" + + # prepare pred layers + self.dec_pred_bbox_embed_share = dec_pred_bbox_embed_share + # prepare class & box embed + _class_embed = ContrastiveEmbed() + + _bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3) + nn.init.constant_(_bbox_embed.layers[-1].weight.data, 0) + nn.init.constant_(_bbox_embed.layers[-1].bias.data, 0) + + if dec_pred_bbox_embed_share: + box_embed_layerlist = [_bbox_embed for i in range(transformer.num_decoder_layers)] + else: + box_embed_layerlist = [ + copy.deepcopy(_bbox_embed) for i in range(transformer.num_decoder_layers) + ] + class_embed_layerlist = [_class_embed for i in range(transformer.num_decoder_layers)] + self.bbox_embed = nn.ModuleList(box_embed_layerlist) + self.class_embed = nn.ModuleList(class_embed_layerlist) + self.transformer.decoder.bbox_embed = self.bbox_embed + self.transformer.decoder.class_embed = self.class_embed + + # two stage + self.two_stage_type = two_stage_type + assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format( + two_stage_type + ) + if two_stage_type != "no": + if two_stage_bbox_embed_share: + assert dec_pred_bbox_embed_share + self.transformer.enc_out_bbox_embed = _bbox_embed + else: + self.transformer.enc_out_bbox_embed = copy.deepcopy(_bbox_embed) + + if two_stage_class_embed_share: + assert dec_pred_bbox_embed_share + self.transformer.enc_out_class_embed = _class_embed + else: + self.transformer.enc_out_class_embed = copy.deepcopy(_class_embed) + + self.refpoint_embed = None + + self._reset_parameters() + + def _reset_parameters(self): + # init input_proj + for proj in self.input_proj: + nn.init.xavier_uniform_(proj[0].weight, gain=1) + nn.init.constant_(proj[0].bias, 0) + + def init_ref_points(self, use_num_queries): + self.refpoint_embed = nn.Embedding(use_num_queries, self.query_dim) + + def forward(self, samples: NestedTensor, targets: List = None, **kw): + """The forward expects a NestedTensor, which consists of: + - samples.tensor: batched images, of shape [batch_size x 3 x H x W] + - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels + + It returns a dict with the following elements: + - "pred_logits": the classification logits (including no-object) for all queries. + Shape= [batch_size x num_queries x num_classes] + - "pred_boxes": The normalized boxes coordinates for all queries, represented as + (center_x, center_y, width, height). These values are normalized in [0, 1], + relative to the size of each individual image (disregarding possible padding). + See PostProcess for information on how to retrieve the unnormalized bounding box. + - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of + dictionnaries containing the two above keys for each decoder layer. + """ + if targets is None: + captions = kw["captions"] + else: + captions = [t["caption"] for t in targets] + len(captions) + + # encoder texts + tokenized = self.tokenizer(captions, padding="longest", return_tensors="pt").to( + samples.device + ) + ( + text_self_attention_masks, + position_ids, + cate_to_token_mask_list, + ) = generate_masks_with_special_tokens_and_transfer_map( + tokenized, self.specical_tokens, self.tokenizer + ) + + if text_self_attention_masks.shape[1] > self.max_text_len: + text_self_attention_masks = text_self_attention_masks[ + :, : self.max_text_len, : self.max_text_len + ] + position_ids = position_ids[:, : self.max_text_len] + tokenized["input_ids"] = tokenized["input_ids"][:, : self.max_text_len] + tokenized["attention_mask"] = tokenized["attention_mask"][:, : self.max_text_len] + tokenized["token_type_ids"] = tokenized["token_type_ids"][:, : self.max_text_len] + + # extract text embeddings + if self.sub_sentence_present: + tokenized_for_encoder = {k: v for k, v in tokenized.items() if k != "attention_mask"} + tokenized_for_encoder["attention_mask"] = text_self_attention_masks + tokenized_for_encoder["position_ids"] = position_ids + else: + # import ipdb; ipdb.set_trace() + tokenized_for_encoder = tokenized + + bert_output = self.bert(**tokenized_for_encoder) # bs, 195, 768 + + encoded_text = self.feat_map(bert_output["last_hidden_state"]) # bs, 195, d_model + text_token_mask = tokenized.attention_mask.bool() # bs, 195 + # text_token_mask: True for nomask, False for mask + # text_self_attention_masks: True for nomask, False for mask + + if encoded_text.shape[1] > self.max_text_len: + encoded_text = encoded_text[:, : self.max_text_len, :] + text_token_mask = text_token_mask[:, : self.max_text_len] + position_ids = position_ids[:, : self.max_text_len] + text_self_attention_masks = text_self_attention_masks[ + :, : self.max_text_len, : self.max_text_len + ] + + text_dict = { + "encoded_text": encoded_text, # bs, 195, d_model + "text_token_mask": text_token_mask, # bs, 195 + "position_ids": position_ids, # bs, 195 + "text_self_attention_masks": text_self_attention_masks, # bs, 195,195 + } + + # import ipdb; ipdb.set_trace() + + if isinstance(samples, (list, torch.Tensor)): + samples = nested_tensor_from_tensor_list(samples) + features, poss = self.backbone(samples) + + srcs = [] + masks = [] + for l, feat in enumerate(features): + src, mask = feat.decompose() + srcs.append(self.input_proj[l](src)) + masks.append(mask) + assert mask is not None + if self.num_feature_levels > len(srcs): + _len_srcs = len(srcs) + for l in range(_len_srcs, self.num_feature_levels): + if l == _len_srcs: + src = self.input_proj[l](features[-1].tensors) + else: + src = self.input_proj[l](srcs[-1]) + m = samples.mask + mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0] + pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype) + srcs.append(src) + masks.append(mask) + poss.append(pos_l) + + input_query_bbox = input_query_label = attn_mask = dn_meta = None + hs, reference, hs_enc, ref_enc, init_box_proposal = self.transformer( + srcs, masks, input_query_bbox, poss, input_query_label, attn_mask, text_dict + ) + + # deformable-detr-like anchor update + outputs_coord_list = [] + for dec_lid, (layer_ref_sig, layer_bbox_embed, layer_hs) in enumerate( + zip(reference[:-1], self.bbox_embed, hs) + ): + layer_delta_unsig = layer_bbox_embed(layer_hs) + layer_outputs_unsig = layer_delta_unsig + inverse_sigmoid(layer_ref_sig) + layer_outputs_unsig = layer_outputs_unsig.sigmoid() + outputs_coord_list.append(layer_outputs_unsig) + outputs_coord_list = torch.stack(outputs_coord_list) + + # output + outputs_class = torch.stack( + [ + layer_cls_embed(layer_hs, text_dict) + for layer_cls_embed, layer_hs in zip(self.class_embed, hs) + ] + ) + out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord_list[-1]} + + # # for intermediate outputs + # if self.aux_loss: + # out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord_list) + + # # for encoder output + # if hs_enc is not None: + # # prepare intermediate outputs + # interm_coord = ref_enc[-1] + # interm_class = self.transformer.enc_out_class_embed(hs_enc[-1], text_dict) + # out['interm_outputs'] = {'pred_logits': interm_class, 'pred_boxes': interm_coord} + # out['interm_outputs_for_matching_pre'] = {'pred_logits': interm_class, 'pred_boxes': init_box_proposal} + + return out + + @torch.jit.unused + def _set_aux_loss(self, outputs_class, outputs_coord): + # this is a workaround to make torchscript happy, as torchscript + # doesn't support dictionary with non-homogeneous values, such + # as a dict having both a Tensor and a list. + return [ + {"pred_logits": a, "pred_boxes": b} + for a, b in zip(outputs_class[:-1], outputs_coord[:-1]) + ] + + +@MODULE_BUILD_FUNCS.registe_with_name(module_name="groundingdino") +def build_groundingdino(args): + + backbone = build_backbone(args) + transformer = build_transformer(args) + + dn_labelbook_size = args.dn_labelbook_size + dec_pred_bbox_embed_share = args.dec_pred_bbox_embed_share + sub_sentence_present = args.sub_sentence_present + bert_base_uncased_path = args.bert_base_uncased_path if 'bert_base_uncased_path' in args else None + + model = GroundingDINO( + backbone, + transformer, + num_queries=args.num_queries, + bert_base_uncased_path=bert_base_uncased_path, + aux_loss=True, + iter_update=True, + query_dim=4, + num_feature_levels=args.num_feature_levels, + nheads=args.nheads, + dec_pred_bbox_embed_share=dec_pred_bbox_embed_share, + two_stage_type=args.two_stage_type, + two_stage_bbox_embed_share=args.two_stage_bbox_embed_share, + two_stage_class_embed_share=args.two_stage_class_embed_share, + num_patterns=args.num_patterns, + dn_number=0, + dn_box_noise_scale=args.dn_box_noise_scale, + dn_label_noise_ratio=args.dn_label_noise_ratio, + dn_labelbook_size=dn_labelbook_size, + text_encoder_type=args.text_encoder_type, + sub_sentence_present=sub_sentence_present, + max_text_len=args.max_text_len, + ) + + return model diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/ms_deform_attn.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/ms_deform_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..489d501bef364020212306d81e9b85c8daa27491 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/ms_deform_attn.py @@ -0,0 +1,413 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Deformable DETR +# Copyright (c) 2020 SenseTime. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------------------------------ +# Modified from: +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/functions/ms_deform_attn_func.py +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/modules/ms_deform_attn.py +# https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/multi_scale_deform_attn.py +# ------------------------------------------------------------------------------------------------ + +import math +import warnings +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Function +from torch.autograd.function import once_differentiable +from torch.nn.init import constant_, xavier_uniform_ + +try: + from groundingdino import _C +except: + warnings.warn("Failed to load custom C++ ops. Running on CPU mode Only!") + + +# helpers +def _is_power_of_2(n): + if (not isinstance(n, int)) or (n < 0): + raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n))) + return (n & (n - 1) == 0) and n != 0 + + +class MultiScaleDeformableAttnFunction(Function): + @staticmethod + def forward( + ctx, + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + im2col_step, + ): + ctx.im2col_step = im2col_step + output = _C.ms_deform_attn_forward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ctx.im2col_step, + ) + ctx.save_for_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ) + return output + + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + ( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + ) = ctx.saved_tensors + grad_value, grad_sampling_loc, grad_attn_weight = _C.ms_deform_attn_backward( + value, + value_spatial_shapes, + value_level_start_index, + sampling_locations, + attention_weights, + grad_output, + ctx.im2col_step, + ) + + return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None + + +def multi_scale_deformable_attn_pytorch( + value: torch.Tensor, + value_spatial_shapes: torch.Tensor, + sampling_locations: torch.Tensor, + attention_weights: torch.Tensor, +) -> torch.Tensor: + + bs, _, num_heads, embed_dims = value.shape + _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape + value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1) + sampling_grids = 2 * sampling_locations - 1 + sampling_value_list = [] + for level, (H_, W_) in enumerate(value_spatial_shapes): + # bs, H_*W_, num_heads, embed_dims -> + # bs, H_*W_, num_heads*embed_dims -> + # bs, num_heads*embed_dims, H_*W_ -> + # bs*num_heads, embed_dims, H_, W_ + value_l_ = ( + value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_) + ) + # bs, num_queries, num_heads, num_points, 2 -> + # bs, num_heads, num_queries, num_points, 2 -> + # bs*num_heads, num_queries, num_points, 2 + sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1) + # bs*num_heads, embed_dims, num_queries, num_points + sampling_value_l_ = F.grid_sample( + value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False + ) + sampling_value_list.append(sampling_value_l_) + # (bs, num_queries, num_heads, num_levels, num_points) -> + # (bs, num_heads, num_queries, num_levels, num_points) -> + # (bs, num_heads, 1, num_queries, num_levels*num_points) + attention_weights = attention_weights.transpose(1, 2).reshape( + bs * num_heads, 1, num_queries, num_levels * num_points + ) + output = ( + (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights) + .sum(-1) + .view(bs, num_heads * embed_dims, num_queries) + ) + return output.transpose(1, 2).contiguous() + + +class MultiScaleDeformableAttention(nn.Module): + """Multi-Scale Deformable Attention Module used in Deformable-DETR + + `Deformable DETR: Deformable Transformers for End-to-End Object Detection. + `_. + + Args: + embed_dim (int): The embedding dimension of Attention. Default: 256. + num_heads (int): The number of attention heads. Default: 8. + num_levels (int): The number of feature map used in Attention. Default: 4. + num_points (int): The number of sampling points for each query + in each head. Default: 4. + img2col_steps (int): The step used in image_to_column. Defualt: 64. + dropout (float): Dropout layer used in output. Default: 0.1. + batch_first (bool): if ``True``, then the input and output tensor will be + provided as `(bs, n, embed_dim)`. Default: False. `(n, bs, embed_dim)` + """ + + def __init__( + self, + embed_dim: int = 256, + num_heads: int = 8, + num_levels: int = 4, + num_points: int = 4, + img2col_step: int = 64, + batch_first: bool = False, + ): + super().__init__() + if embed_dim % num_heads != 0: + raise ValueError( + "embed_dim must be divisible by num_heads, but got {} and {}".format( + embed_dim, num_heads + ) + ) + head_dim = embed_dim // num_heads + + self.batch_first = batch_first + + if not _is_power_of_2(head_dim): + warnings.warn( + """ + You'd better set d_model in MSDeformAttn to make sure that + each dim of the attention head a power of 2, which is more efficient. + """ + ) + + self.im2col_step = img2col_step + self.embed_dim = embed_dim + self.num_heads = num_heads + self.num_levels = num_levels + self.num_points = num_points + self.sampling_offsets = nn.Linear(embed_dim, num_heads * num_levels * num_points * 2) + self.attention_weights = nn.Linear(embed_dim, num_heads * num_levels * num_points) + self.value_proj = nn.Linear(embed_dim, embed_dim) + self.output_proj = nn.Linear(embed_dim, embed_dim) + + self.init_weights() + + def _reset_parameters(self): + return self.init_weights() + + def init_weights(self): + """ + Default initialization for Parameters of Module. + """ + constant_(self.sampling_offsets.weight.data, 0.0) + thetas = torch.arange(self.num_heads, dtype=torch.float32) * ( + 2.0 * math.pi / self.num_heads + ) + grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) + grid_init = ( + (grid_init / grid_init.abs().max(-1, keepdim=True)[0]) + .view(self.num_heads, 1, 1, 2) + .repeat(1, self.num_levels, self.num_points, 1) + ) + for i in range(self.num_points): + grid_init[:, :, i, :] *= i + 1 + with torch.no_grad(): + self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1)) + constant_(self.attention_weights.weight.data, 0.0) + constant_(self.attention_weights.bias.data, 0.0) + xavier_uniform_(self.value_proj.weight.data) + constant_(self.value_proj.bias.data, 0.0) + xavier_uniform_(self.output_proj.weight.data) + constant_(self.output_proj.bias.data, 0.0) + + def freeze_sampling_offsets(self): + print("Freeze sampling offsets") + self.sampling_offsets.weight.requires_grad = False + self.sampling_offsets.bias.requires_grad = False + + def freeze_attention_weights(self): + print("Freeze attention weights") + self.attention_weights.weight.requires_grad = False + self.attention_weights.bias.requires_grad = False + + def forward( + self, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + value: Optional[torch.Tensor] = None, + query_pos: Optional[torch.Tensor] = None, + key_padding_mask: Optional[torch.Tensor] = None, + reference_points: Optional[torch.Tensor] = None, + spatial_shapes: Optional[torch.Tensor] = None, + level_start_index: Optional[torch.Tensor] = None, + **kwargs + ) -> torch.Tensor: + + """Forward Function of MultiScaleDeformableAttention + + Args: + query (torch.Tensor): Query embeddings with shape + `(num_query, bs, embed_dim)` + key (torch.Tensor): Key embeddings with shape + `(num_key, bs, embed_dim)` + value (torch.Tensor): Value embeddings with shape + `(num_key, bs, embed_dim)` + query_pos (torch.Tensor): The position embedding for `query`. Default: None. + key_padding_mask (torch.Tensor): ByteTensor for `query`, with shape `(bs, num_key)`, + indicating which elements within `key` to be ignored in attention. + reference_points (torch.Tensor): The normalized reference points + with shape `(bs, num_query, num_levels, 2)`, + all elements is range in [0, 1], top-left (0, 0), + bottom-right (1, 1), including padding are. + or `(N, Length_{query}, num_levels, 4)`, add additional + two dimensions `(h, w)` to form reference boxes. + spatial_shapes (torch.Tensor): Spatial shape of features in different levels. + With shape `(num_levels, 2)`, last dimension represents `(h, w)`. + level_start_index (torch.Tensor): The start index of each level. A tensor with + shape `(num_levels, )` which can be represented as + `[0, h_0 * w_0, h_0 * w_0 + h_1 * w_1, ...]`. + + Returns: + torch.Tensor: forward results with shape `(num_query, bs, embed_dim)` + """ + + if value is None: + value = query + + if query_pos is not None: + query = query + query_pos + + if not self.batch_first: + # change to (bs, num_query ,embed_dims) + query = query.permute(1, 0, 2) + value = value.permute(1, 0, 2) + + bs, num_query, _ = query.shape + bs, num_value, _ = value.shape + + assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value + + value = self.value_proj(value) + if key_padding_mask is not None: + value = value.masked_fill(key_padding_mask[..., None], float(0)) + value = value.view(bs, num_value, self.num_heads, -1) + sampling_offsets = self.sampling_offsets(query).view( + bs, num_query, self.num_heads, self.num_levels, self.num_points, 2 + ) + attention_weights = self.attention_weights(query).view( + bs, num_query, self.num_heads, self.num_levels * self.num_points + ) + attention_weights = attention_weights.softmax(-1) + attention_weights = attention_weights.view( + bs, + num_query, + self.num_heads, + self.num_levels, + self.num_points, + ) + + # bs, num_query, num_heads, num_levels, num_points, 2 + if reference_points.shape[-1] == 2: + offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) + sampling_locations = ( + reference_points[:, :, None, :, None, :] + + sampling_offsets / offset_normalizer[None, None, None, :, None, :] + ) + elif reference_points.shape[-1] == 4: + sampling_locations = ( + reference_points[:, :, None, :, None, :2] + + sampling_offsets + / self.num_points + * reference_points[:, :, None, :, None, 2:] + * 0.5 + ) + else: + raise ValueError( + "Last dim of reference_points must be 2 or 4, but get {} instead.".format( + reference_points.shape[-1] + ) + ) + + if torch.cuda.is_available() and value.is_cuda: + halffloat = False + if value.dtype == torch.float16: + halffloat = True + value = value.float() + sampling_locations = sampling_locations.float() + attention_weights = attention_weights.float() + + output = MultiScaleDeformableAttnFunction.apply( + value, + spatial_shapes, + level_start_index, + sampling_locations, + attention_weights, + self.im2col_step, + ) + + if halffloat: + output = output.half() + else: + output = multi_scale_deformable_attn_pytorch( + value, spatial_shapes, sampling_locations, attention_weights + ) + + output = self.output_proj(output) + + if not self.batch_first: + output = output.permute(1, 0, 2) + + return output + + +def create_dummy_class(klass, dependency, message=""): + """ + When a dependency of a class is not available, create a dummy class which throws ImportError + when used. + + Args: + klass (str): name of the class. + dependency (str): name of the dependency. + message: extra message to print + Returns: + class: a class object + """ + err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, klass) + if message: + err = err + " " + message + + class _DummyMetaClass(type): + # throw error on class attribute access + def __getattr__(_, __): # noqa: B902 + raise ImportError(err) + + class _Dummy(object, metaclass=_DummyMetaClass): + # throw error on constructor + def __init__(self, *args, **kwargs): + raise ImportError(err) + + return _Dummy + + +def create_dummy_func(func, dependency, message=""): + """ + When a dependency of a function is not available, create a dummy function which throws + ImportError when used. + + Args: + func (str): name of the function. + dependency (str or list[str]): name(s) of the dependency. + message: extra message to print + Returns: + function: a function object + """ + err = "Cannot import '{}', therefore '{}' is not available.".format(dependency, func) + if message: + err = err + " " + message + + if isinstance(dependency, (list, tuple)): + dependency = ",".join(dependency) + + def _dummy(*args, **kwargs): + raise ImportError(err) + + return _dummy diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..d554215ecfaa7ad5a7661fa50757e5de713f0b32 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer.py @@ -0,0 +1,960 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# DINO +# Copyright (c) 2022 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Conditional DETR Transformer class. +# Copyright (c) 2021 Microsoft. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Modified from DETR (https://github.com/facebookresearch/detr) +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. +# ------------------------------------------------------------------------ + +from typing import Optional + +import torch +import torch.utils.checkpoint as checkpoint +from torch import Tensor, nn + +from groundingdino.util.misc import inverse_sigmoid + +from .fuse_modules import BiAttentionBlock +from .ms_deform_attn import MultiScaleDeformableAttention as MSDeformAttn +from .transformer_vanilla import TransformerEncoderLayer +from .utils import ( + MLP, + _get_activation_fn, + _get_clones, + gen_encoder_output_proposals, + gen_sineembed_for_position, + get_sine_pos_embed, +) + + +class Transformer(nn.Module): + def __init__( + self, + d_model=256, + nhead=8, + num_queries=300, + num_encoder_layers=6, + num_unicoder_layers=0, + num_decoder_layers=6, + dim_feedforward=2048, + dropout=0.0, + activation="relu", + normalize_before=False, + return_intermediate_dec=False, + query_dim=4, + num_patterns=0, + # for deformable encoder + num_feature_levels=1, + enc_n_points=4, + dec_n_points=4, + # init query + learnable_tgt_init=False, + # two stage + two_stage_type="no", # ['no', 'standard', 'early', 'combine', 'enceachlayer', 'enclayer1'] + embed_init_tgt=False, + # for text + use_text_enhancer=False, + use_fusion_layer=False, + use_checkpoint=False, + use_transformer_ckpt=False, + use_text_cross_attention=False, + text_dropout=0.1, + fusion_dropout=0.1, + fusion_droppath=0.0, + ): + super().__init__() + self.num_feature_levels = num_feature_levels + self.num_encoder_layers = num_encoder_layers + self.num_unicoder_layers = num_unicoder_layers + self.num_decoder_layers = num_decoder_layers + self.num_queries = num_queries + assert query_dim == 4 + + # choose encoder layer type + encoder_layer = DeformableTransformerEncoderLayer( + d_model, dim_feedforward, dropout, activation, num_feature_levels, nhead, enc_n_points + ) + + if use_text_enhancer: + text_enhance_layer = TransformerEncoderLayer( + d_model=d_model, + nhead=nhead // 2, + dim_feedforward=dim_feedforward // 2, + dropout=text_dropout, + ) + else: + text_enhance_layer = None + + if use_fusion_layer: + feature_fusion_layer = BiAttentionBlock( + v_dim=d_model, + l_dim=d_model, + embed_dim=dim_feedforward // 2, + num_heads=nhead // 2, + dropout=fusion_dropout, + drop_path=fusion_droppath, + ) + else: + feature_fusion_layer = None + + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + assert encoder_norm is None + self.encoder = TransformerEncoder( + encoder_layer, + num_encoder_layers, + d_model=d_model, + num_queries=num_queries, + text_enhance_layer=text_enhance_layer, + feature_fusion_layer=feature_fusion_layer, + use_checkpoint=use_checkpoint, + use_transformer_ckpt=use_transformer_ckpt, + ) + + # choose decoder layer type + decoder_layer = DeformableTransformerDecoderLayer( + d_model, + dim_feedforward, + dropout, + activation, + num_feature_levels, + nhead, + dec_n_points, + use_text_cross_attention=use_text_cross_attention, + ) + + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder( + decoder_layer, + num_decoder_layers, + decoder_norm, + return_intermediate=return_intermediate_dec, + d_model=d_model, + query_dim=query_dim, + num_feature_levels=num_feature_levels, + ) + + self.d_model = d_model + self.nhead = nhead + self.dec_layers = num_decoder_layers + self.num_queries = num_queries # useful for single stage model only + self.num_patterns = num_patterns + if not isinstance(num_patterns, int): + Warning("num_patterns should be int but {}".format(type(num_patterns))) + self.num_patterns = 0 + + if num_feature_levels > 1: + if self.num_encoder_layers > 0: + self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model)) + else: + self.level_embed = None + + self.learnable_tgt_init = learnable_tgt_init + assert learnable_tgt_init, "why not learnable_tgt_init" + self.embed_init_tgt = embed_init_tgt + if (two_stage_type != "no" and embed_init_tgt) or (two_stage_type == "no"): + self.tgt_embed = nn.Embedding(self.num_queries, d_model) + nn.init.normal_(self.tgt_embed.weight.data) + else: + self.tgt_embed = None + + # for two stage + self.two_stage_type = two_stage_type + assert two_stage_type in ["no", "standard"], "unknown param {} of two_stage_type".format( + two_stage_type + ) + if two_stage_type == "standard": + # anchor selection at the output of encoder + self.enc_output = nn.Linear(d_model, d_model) + self.enc_output_norm = nn.LayerNorm(d_model) + self.two_stage_wh_embedding = None + + if two_stage_type == "no": + self.init_ref_points(num_queries) # init self.refpoint_embed + + self.enc_out_class_embed = None + self.enc_out_bbox_embed = None + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + for m in self.modules(): + if isinstance(m, MSDeformAttn): + m._reset_parameters() + if self.num_feature_levels > 1 and self.level_embed is not None: + nn.init.normal_(self.level_embed) + + def get_valid_ratio(self, mask): + _, H, W = mask.shape + valid_H = torch.sum(~mask[:, :, 0], 1) + valid_W = torch.sum(~mask[:, 0, :], 1) + valid_ratio_h = valid_H.float() / H + valid_ratio_w = valid_W.float() / W + valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1) + return valid_ratio + + def init_ref_points(self, use_num_queries): + self.refpoint_embed = nn.Embedding(use_num_queries, 4) + + def forward(self, srcs, masks, refpoint_embed, pos_embeds, tgt, attn_mask=None, text_dict=None): + """ + Input: + - srcs: List of multi features [bs, ci, hi, wi] + - masks: List of multi masks [bs, hi, wi] + - refpoint_embed: [bs, num_dn, 4]. None in infer + - pos_embeds: List of multi pos embeds [bs, ci, hi, wi] + - tgt: [bs, num_dn, d_model]. None in infer + + """ + # prepare input for encoder + src_flatten = [] + mask_flatten = [] + lvl_pos_embed_flatten = [] + spatial_shapes = [] + for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)): + bs, c, h, w = src.shape + spatial_shape = (h, w) + spatial_shapes.append(spatial_shape) + + src = src.flatten(2).transpose(1, 2) # bs, hw, c + mask = mask.flatten(1) # bs, hw + pos_embed = pos_embed.flatten(2).transpose(1, 2) # bs, hw, c + if self.num_feature_levels > 1 and self.level_embed is not None: + lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1) + else: + lvl_pos_embed = pos_embed + lvl_pos_embed_flatten.append(lvl_pos_embed) + src_flatten.append(src) + mask_flatten.append(mask) + src_flatten = torch.cat(src_flatten, 1) # bs, \sum{hxw}, c + mask_flatten = torch.cat(mask_flatten, 1) # bs, \sum{hxw} + lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) # bs, \sum{hxw}, c + spatial_shapes = torch.as_tensor( + spatial_shapes, dtype=torch.long, device=src_flatten.device + ) + level_start_index = torch.cat( + (spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]) + ) + valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1).to(src.dtype) + + # two stage + enc_topk_proposals = enc_refpoint_embed = None + + ######################################################### + # Begin Encoder + ######################################################### + memory, memory_text = self.encoder( + src_flatten, + pos=lvl_pos_embed_flatten, + level_start_index=level_start_index, + spatial_shapes=spatial_shapes, + valid_ratios=valid_ratios, + key_padding_mask=mask_flatten, + memory_text=text_dict["encoded_text"], + text_attention_mask=~text_dict["text_token_mask"], + # we ~ the mask . False means use the token; True means pad the token + position_ids=text_dict["position_ids"], + text_self_attention_masks=text_dict["text_self_attention_masks"], + ) + ######################################################### + # End Encoder + # - memory: bs, \sum{hw}, c + # - mask_flatten: bs, \sum{hw} + # - lvl_pos_embed_flatten: bs, \sum{hw}, c + # - enc_intermediate_output: None or (nenc+1, bs, nq, c) or (nenc, bs, nq, c) + # - enc_intermediate_refpoints: None or (nenc+1, bs, nq, c) or (nenc, bs, nq, c) + ######################################################### + text_dict["encoded_text"] = memory_text + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # if memory.isnan().any() | memory.isinf().any(): + # import ipdb; ipdb.set_trace() + + if self.two_stage_type == "standard": + output_memory, output_proposals = gen_encoder_output_proposals( + memory, mask_flatten, spatial_shapes + ) + output_memory = self.enc_output_norm(self.enc_output(output_memory)) + + if text_dict is not None: + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory, text_dict) + else: + enc_outputs_class_unselected = self.enc_out_class_embed(output_memory) + + topk_logits = enc_outputs_class_unselected.max(-1)[0] + enc_outputs_coord_unselected = ( + self.enc_out_bbox_embed(output_memory) + output_proposals + ) # (bs, \sum{hw}, 4) unsigmoid + topk = self.num_queries + + topk_proposals = torch.topk(topk_logits, topk, dim=1)[1] # bs, nq + + # gather boxes + refpoint_embed_undetach = torch.gather( + enc_outputs_coord_unselected, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4) + ) # unsigmoid + refpoint_embed_ = refpoint_embed_undetach.detach() + init_box_proposal = torch.gather( + output_proposals, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4) + ).sigmoid() # sigmoid + + # gather tgt + tgt_undetach = torch.gather( + output_memory, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, self.d_model) + ) + if self.embed_init_tgt: + tgt_ = ( + self.tgt_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, d_model + else: + tgt_ = tgt_undetach.detach() + + if refpoint_embed is not None: + refpoint_embed = torch.cat([refpoint_embed, refpoint_embed_], dim=1) + tgt = torch.cat([tgt, tgt_], dim=1) + else: + refpoint_embed, tgt = refpoint_embed_, tgt_ + + elif self.two_stage_type == "no": + tgt_ = ( + self.tgt_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, d_model + refpoint_embed_ = ( + self.refpoint_embed.weight[:, None, :].repeat(1, bs, 1).transpose(0, 1) + ) # nq, bs, 4 + + if refpoint_embed is not None: + refpoint_embed = torch.cat([refpoint_embed, refpoint_embed_], dim=1) + tgt = torch.cat([tgt, tgt_], dim=1) + else: + refpoint_embed, tgt = refpoint_embed_, tgt_ + + if self.num_patterns > 0: + tgt_embed = tgt.repeat(1, self.num_patterns, 1) + refpoint_embed = refpoint_embed.repeat(1, self.num_patterns, 1) + tgt_pat = self.patterns.weight[None, :, :].repeat_interleave( + self.num_queries, 1 + ) # 1, n_q*n_pat, d_model + tgt = tgt_embed + tgt_pat + + init_box_proposal = refpoint_embed_.sigmoid() + + else: + raise NotImplementedError("unknown two_stage_type {}".format(self.two_stage_type)) + ######################################################### + # End preparing tgt + # - tgt: bs, NQ, d_model + # - refpoint_embed(unsigmoid): bs, NQ, d_model + ######################################################### + + ######################################################### + # Begin Decoder + ######################################################### + hs, references = self.decoder( + tgt=tgt.transpose(0, 1), + memory=memory.transpose(0, 1), + memory_key_padding_mask=mask_flatten, + pos=lvl_pos_embed_flatten.transpose(0, 1), + refpoints_unsigmoid=refpoint_embed.transpose(0, 1), + level_start_index=level_start_index, + spatial_shapes=spatial_shapes, + valid_ratios=valid_ratios, + tgt_mask=attn_mask, + memory_text=text_dict["encoded_text"], + text_attention_mask=~text_dict["text_token_mask"], + # we ~ the mask . False means use the token; True means pad the token + ) + ######################################################### + # End Decoder + # hs: n_dec, bs, nq, d_model + # references: n_dec+1, bs, nq, query_dim + ######################################################### + + ######################################################### + # Begin postprocess + ######################################################### + if self.two_stage_type == "standard": + hs_enc = tgt_undetach.unsqueeze(0) + ref_enc = refpoint_embed_undetach.sigmoid().unsqueeze(0) + else: + hs_enc = ref_enc = None + ######################################################### + # End postprocess + # hs_enc: (n_enc+1, bs, nq, d_model) or (1, bs, nq, d_model) or (n_enc, bs, nq, d_model) or None + # ref_enc: (n_enc+1, bs, nq, query_dim) or (1, bs, nq, query_dim) or (n_enc, bs, nq, d_model) or None + ######################################################### + + return hs, references, hs_enc, ref_enc, init_box_proposal + # hs: (n_dec, bs, nq, d_model) + # references: sigmoid coordinates. (n_dec+1, bs, bq, 4) + # hs_enc: (n_enc+1, bs, nq, d_model) or (1, bs, nq, d_model) or None + # ref_enc: sigmoid coordinates. \ + # (n_enc+1, bs, nq, query_dim) or (1, bs, nq, query_dim) or None + + +class TransformerEncoder(nn.Module): + def __init__( + self, + encoder_layer, + num_layers, + d_model=256, + num_queries=300, + enc_layer_share=False, + text_enhance_layer=None, + feature_fusion_layer=None, + use_checkpoint=False, + use_transformer_ckpt=False, + ): + """_summary_ + + Args: + encoder_layer (_type_): _description_ + num_layers (_type_): _description_ + norm (_type_, optional): _description_. Defaults to None. + d_model (int, optional): _description_. Defaults to 256. + num_queries (int, optional): _description_. Defaults to 300. + enc_layer_share (bool, optional): _description_. Defaults to False. + + """ + super().__init__() + # prepare layers + self.layers = [] + self.text_layers = [] + self.fusion_layers = [] + if num_layers > 0: + self.layers = _get_clones(encoder_layer, num_layers, layer_share=enc_layer_share) + + if text_enhance_layer is not None: + self.text_layers = _get_clones( + text_enhance_layer, num_layers, layer_share=enc_layer_share + ) + if feature_fusion_layer is not None: + self.fusion_layers = _get_clones( + feature_fusion_layer, num_layers, layer_share=enc_layer_share + ) + else: + self.layers = [] + del encoder_layer + + if text_enhance_layer is not None: + self.text_layers = [] + del text_enhance_layer + if feature_fusion_layer is not None: + self.fusion_layers = [] + del feature_fusion_layer + + self.query_scale = None + self.num_queries = num_queries + self.num_layers = num_layers + self.d_model = d_model + + self.use_checkpoint = use_checkpoint + self.use_transformer_ckpt = use_transformer_ckpt + + @staticmethod + def get_reference_points(spatial_shapes, valid_ratios, device): + reference_points_list = [] + for lvl, (H_, W_) in enumerate(spatial_shapes): + + ref_y, ref_x = torch.meshgrid( + torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device), + torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device), + ) + ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_) + ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_) + ref = torch.stack((ref_x, ref_y), -1) + reference_points_list.append(ref) + reference_points = torch.cat(reference_points_list, 1) + reference_points = reference_points[:, :, None] * valid_ratios[:, None] + return reference_points + + def forward( + self, + # for images + src: Tensor, + pos: Tensor, + spatial_shapes: Tensor, + level_start_index: Tensor, + valid_ratios: Tensor, + key_padding_mask: Tensor, + # for texts + memory_text: Tensor = None, + text_attention_mask: Tensor = None, + pos_text: Tensor = None, + text_self_attention_masks: Tensor = None, + position_ids: Tensor = None, + ): + """ + Input: + - src: [bs, sum(hi*wi), 256] + - pos: pos embed for src. [bs, sum(hi*wi), 256] + - spatial_shapes: h,w of each level [num_level, 2] + - level_start_index: [num_level] start point of level in sum(hi*wi). + - valid_ratios: [bs, num_level, 2] + - key_padding_mask: [bs, sum(hi*wi)] + + - memory_text: bs, n_text, 256 + - text_attention_mask: bs, n_text + False for no padding; True for padding + - pos_text: bs, n_text, 256 + + - position_ids: bs, n_text + Intermedia: + - reference_points: [bs, sum(hi*wi), num_level, 2] + Outpus: + - output: [bs, sum(hi*wi), 256] + """ + + output = src + + # preparation and reshape + if self.num_layers > 0: + reference_points = self.get_reference_points( + spatial_shapes, valid_ratios, device=src.device + ) + + if self.text_layers: + # generate pos_text + bs, n_text, text_dim = memory_text.shape + if pos_text is None and position_ids is None: + pos_text = ( + torch.arange(n_text, device=memory_text.device) + .float() + .unsqueeze(0) + .unsqueeze(-1) + .repeat(bs, 1, 1) + ) + pos_text = get_sine_pos_embed(pos_text, num_pos_feats=256, exchange_xy=False) + if position_ids is not None: + pos_text = get_sine_pos_embed( + position_ids[..., None], num_pos_feats=256, exchange_xy=False + ) + pos_text = pos_text.to(src.dtype) + + # main process + for layer_id, layer in enumerate(self.layers): + # if output.isnan().any() or memory_text.isnan().any(): + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + if self.fusion_layers: + if self.use_checkpoint: + output, memory_text = checkpoint.checkpoint( + self.fusion_layers[layer_id], + output, + memory_text, + key_padding_mask, + text_attention_mask, + ) + else: + output, memory_text = self.fusion_layers[layer_id]( + v=output, + l=memory_text, + attention_mask_v=key_padding_mask, + attention_mask_l=text_attention_mask, + ) + + if self.text_layers: + memory_text = self.text_layers[layer_id]( + src=memory_text.transpose(0, 1), + src_mask=~text_self_attention_masks, # note we use ~ for mask here + src_key_padding_mask=text_attention_mask, + pos=(pos_text.transpose(0, 1) if pos_text is not None else None), + ).transpose(0, 1) + + # main process + if self.use_transformer_ckpt: + output = checkpoint.checkpoint( + layer, + output, + pos, + reference_points, + spatial_shapes, + level_start_index, + key_padding_mask, + ) + else: + output = layer( + src=output, + pos=pos, + reference_points=reference_points, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + key_padding_mask=key_padding_mask, + ) + + return output, memory_text + + +class TransformerDecoder(nn.Module): + def __init__( + self, + decoder_layer, + num_layers, + norm=None, + return_intermediate=False, + d_model=256, + query_dim=4, + num_feature_levels=1, + ): + super().__init__() + if num_layers > 0: + self.layers = _get_clones(decoder_layer, num_layers) + else: + self.layers = [] + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + assert return_intermediate, "support return_intermediate only" + self.query_dim = query_dim + assert query_dim in [2, 4], "query_dim should be 2/4 but {}".format(query_dim) + self.num_feature_levels = num_feature_levels + + self.ref_point_head = MLP(query_dim // 2 * d_model, d_model, d_model, 2) + self.query_pos_sine_scale = None + + self.query_scale = None + self.bbox_embed = None + self.class_embed = None + + self.d_model = d_model + + self.ref_anchor_head = None + + def forward( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + refpoints_unsigmoid: Optional[Tensor] = None, # num_queries, bs, 2 + # for memory + level_start_index: Optional[Tensor] = None, # num_levels + spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + valid_ratios: Optional[Tensor] = None, + # for text + memory_text: Optional[Tensor] = None, + text_attention_mask: Optional[Tensor] = None, + ): + """ + Input: + - tgt: nq, bs, d_model + - memory: hw, bs, d_model + - pos: hw, bs, d_model + - refpoints_unsigmoid: nq, bs, 2/4 + - valid_ratios/spatial_shapes: bs, nlevel, 2 + """ + output = tgt + + intermediate = [] + reference_points = refpoints_unsigmoid.sigmoid() + ref_points = [reference_points] + + for layer_id, layer in enumerate(self.layers): + + if reference_points.shape[-1] == 4: + reference_points_input = ( + reference_points[:, :, None] + * torch.cat([valid_ratios, valid_ratios], -1)[None, :] + ) # nq, bs, nlevel, 4 + else: + assert reference_points.shape[-1] == 2 + reference_points_input = reference_points[:, :, None] * valid_ratios[None, :] + query_sine_embed = gen_sineembed_for_position( + reference_points_input[:, :, 0, :] + ) # nq, bs, 256*2 + + # conditional query + raw_query_pos = self.ref_point_head(query_sine_embed) # nq, bs, 256 + pos_scale = self.query_scale(output) if self.query_scale is not None else 1 + query_pos = pos_scale * raw_query_pos + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # if query_pos.isnan().any() | query_pos.isinf().any(): + # import ipdb; ipdb.set_trace() + + # main process + output = layer( + tgt=output, + tgt_query_pos=query_pos, + tgt_query_sine_embed=query_sine_embed, + tgt_key_padding_mask=tgt_key_padding_mask, + tgt_reference_points=reference_points_input, + memory_text=memory_text, + text_attention_mask=text_attention_mask, + memory=memory, + memory_key_padding_mask=memory_key_padding_mask, + memory_level_start_index=level_start_index, + memory_spatial_shapes=spatial_shapes, + memory_pos=pos, + self_attn_mask=tgt_mask, + cross_attn_mask=memory_mask, + ) + if output.isnan().any() | output.isinf().any(): + print(f"output layer_id {layer_id} is nan") + try: + num_nan = output.isnan().sum().item() + num_inf = output.isinf().sum().item() + print(f"num_nan {num_nan}, num_inf {num_inf}") + except Exception as e: + print(e) + # if os.environ.get("SHILONG_AMP_INFNAN_DEBUG") == '1': + # import ipdb; ipdb.set_trace() + + # iter update + if self.bbox_embed is not None: + # box_holder = self.bbox_embed(output) + # box_holder[..., :self.query_dim] += inverse_sigmoid(reference_points) + # new_reference_points = box_holder[..., :self.query_dim].sigmoid() + + reference_before_sigmoid = inverse_sigmoid(reference_points) + delta_unsig = self.bbox_embed[layer_id](output) + outputs_unsig = delta_unsig + reference_before_sigmoid + new_reference_points = outputs_unsig.sigmoid() + + reference_points = new_reference_points.detach() + # if layer_id != self.num_layers - 1: + ref_points.append(new_reference_points) + + intermediate.append(self.norm(output)) + + return [ + [itm_out.transpose(0, 1) for itm_out in intermediate], + [itm_refpoint.transpose(0, 1) for itm_refpoint in ref_points], + ] + + +class DeformableTransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + ): + super().__init__() + + # self attention + self.self_attn = MSDeformAttn( + embed_dim=d_model, + num_levels=n_levels, + num_heads=n_heads, + num_points=n_points, + batch_first=True, + ) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn) + self.dropout2 = nn.Dropout(dropout) + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout3 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(d_model) + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, src): + src2 = self.linear2(self.dropout2(self.activation(self.linear1(src)))) + src = src + self.dropout3(src2) + src = self.norm2(src) + return src + + def forward( + self, src, pos, reference_points, spatial_shapes, level_start_index, key_padding_mask=None + ): + # self attention + # import ipdb; ipdb.set_trace() + src2 = self.self_attn( + query=self.with_pos_embed(src, pos), + reference_points=reference_points, + value=src, + spatial_shapes=spatial_shapes, + level_start_index=level_start_index, + key_padding_mask=key_padding_mask, + ) + src = src + self.dropout1(src2) + src = self.norm1(src) + + # ffn + src = self.forward_ffn(src) + + return src + + +class DeformableTransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model=256, + d_ffn=1024, + dropout=0.1, + activation="relu", + n_levels=4, + n_heads=8, + n_points=4, + use_text_feat_guide=False, + use_text_cross_attention=False, + ): + super().__init__() + + # cross attention + self.cross_attn = MSDeformAttn( + embed_dim=d_model, + num_levels=n_levels, + num_heads=n_heads, + num_points=n_points, + batch_first=True, + ) + self.dropout1 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm1 = nn.LayerNorm(d_model) + + # cross attention text + if use_text_cross_attention: + self.ca_text = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.catext_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.catext_norm = nn.LayerNorm(d_model) + + # self attention + self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout) + self.dropout2 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm2 = nn.LayerNorm(d_model) + + # ffn + self.linear1 = nn.Linear(d_model, d_ffn) + self.activation = _get_activation_fn(activation, d_model=d_ffn, batch_dim=1) + self.dropout3 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.linear2 = nn.Linear(d_ffn, d_model) + self.dropout4 = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + self.norm3 = nn.LayerNorm(d_model) + + self.key_aware_proj = None + self.use_text_feat_guide = use_text_feat_guide + assert not use_text_feat_guide + self.use_text_cross_attention = use_text_cross_attention + + def rm_self_attn_modules(self): + self.self_attn = None + self.dropout2 = None + self.norm2 = None + + @staticmethod + def with_pos_embed(tensor, pos): + return tensor if pos is None else tensor + pos + + def forward_ffn(self, tgt): + with torch.cuda.amp.autocast(enabled=False): + tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout4(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward( + self, + # for tgt + tgt: Optional[Tensor], # nq, bs, d_model + tgt_query_pos: Optional[Tensor] = None, # pos for query. MLP(Sine(pos)) + tgt_query_sine_embed: Optional[Tensor] = None, # pos for query. Sine(pos) + tgt_key_padding_mask: Optional[Tensor] = None, + tgt_reference_points: Optional[Tensor] = None, # nq, bs, 4 + memory_text: Optional[Tensor] = None, # bs, num_token, d_model + text_attention_mask: Optional[Tensor] = None, # bs, num_token + # for memory + memory: Optional[Tensor] = None, # hw, bs, d_model + memory_key_padding_mask: Optional[Tensor] = None, + memory_level_start_index: Optional[Tensor] = None, # num_levels + memory_spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2 + memory_pos: Optional[Tensor] = None, # pos for memory + # sa + self_attn_mask: Optional[Tensor] = None, # mask used for self-attention + cross_attn_mask: Optional[Tensor] = None, # mask used for cross-attention + ): + """ + Input: + - tgt/tgt_query_pos: nq, bs, d_model + - + """ + assert cross_attn_mask is None + + # self attention + if self.self_attn is not None: + # import ipdb; ipdb.set_trace() + q = k = self.with_pos_embed(tgt, tgt_query_pos) + tgt2 = self.self_attn(q, k, tgt, attn_mask=self_attn_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + + if self.use_text_cross_attention: + tgt2 = self.ca_text( + self.with_pos_embed(tgt, tgt_query_pos), + memory_text.transpose(0, 1), + memory_text.transpose(0, 1), + key_padding_mask=text_attention_mask, + )[0] + tgt = tgt + self.catext_dropout(tgt2) + tgt = self.catext_norm(tgt) + + tgt2 = self.cross_attn( + query=self.with_pos_embed(tgt, tgt_query_pos).transpose(0, 1), + reference_points=tgt_reference_points.transpose(0, 1).contiguous(), + value=memory.transpose(0, 1), + spatial_shapes=memory_spatial_shapes, + level_start_index=memory_level_start_index, + key_padding_mask=memory_key_padding_mask, + ).transpose(0, 1) + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + + # ffn + tgt = self.forward_ffn(tgt) + + return tgt + + +def build_transformer(args): + return Transformer( + d_model=args.hidden_dim, + dropout=args.dropout, + nhead=args.nheads, + num_queries=args.num_queries, + dim_feedforward=args.dim_feedforward, + num_encoder_layers=args.enc_layers, + num_decoder_layers=args.dec_layers, + normalize_before=args.pre_norm, + return_intermediate_dec=True, + query_dim=args.query_dim, + activation=args.transformer_activation, + num_patterns=args.num_patterns, + num_feature_levels=args.num_feature_levels, + enc_n_points=args.enc_n_points, + dec_n_points=args.dec_n_points, + learnable_tgt_init=True, + # two stage + two_stage_type=args.two_stage_type, # ['no', 'standard', 'early'] + embed_init_tgt=args.embed_init_tgt, + use_text_enhancer=args.use_text_enhancer, + use_fusion_layer=args.use_fusion_layer, + use_checkpoint=args.use_checkpoint, + use_transformer_ckpt=args.use_transformer_ckpt, + use_text_cross_attention=args.use_text_cross_attention, + text_dropout=args.text_dropout, + fusion_dropout=args.fusion_dropout, + fusion_droppath=args.fusion_droppath, + ) diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer_vanilla.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer_vanilla.py new file mode 100644 index 0000000000000000000000000000000000000000..10c0920c1a217af5bb3e1b13077568035ab3b7b5 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/transformer_vanilla.py @@ -0,0 +1,123 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copyright (c) Aishwarya Kamath & Nicolas Carion. Licensed under the Apache License 2.0. All Rights Reserved +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +DETR Transformer class. + +Copy-paste from torch.nn.Transformer with modifications: + * positional encodings are passed in MHattention + * extra LN at the end of encoder is removed + * decoder returns a stack of activations from all decoding layers +""" +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from .utils import ( + MLP, + _get_activation_fn, + _get_clones, + gen_encoder_output_proposals, + gen_sineembed_for_position, + sigmoid_focal_loss, +) + + +class TextTransformer(nn.Module): + def __init__(self, num_layers, d_model=256, nheads=8, dim_feedforward=2048, dropout=0.1): + super().__init__() + self.num_layers = num_layers + self.d_model = d_model + self.nheads = nheads + self.dim_feedforward = dim_feedforward + self.norm = None + + single_encoder_layer = TransformerEncoderLayer( + d_model=d_model, nhead=nheads, dim_feedforward=dim_feedforward, dropout=dropout + ) + self.layers = _get_clones(single_encoder_layer, num_layers) + + def forward(self, memory_text: torch.Tensor, text_attention_mask: torch.Tensor): + """ + + Args: + text_attention_mask: bs, num_token + memory_text: bs, num_token, d_model + + Raises: + RuntimeError: _description_ + + Returns: + output: bs, num_token, d_model + """ + + output = memory_text.transpose(0, 1) + + for layer in self.layers: + output = layer(output, src_key_padding_mask=text_attention_mask) + + if self.norm is not None: + output = self.norm(output) + + return output.transpose(0, 1) + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + self.nhead = nhead + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward( + self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + # repeat attn mask + if src_mask.dim() == 3 and src_mask.shape[0] == src.shape[1]: + # bs, num_q, num_k + src_mask = src_mask.repeat(self.nhead, 1, 1) + + q = k = self.with_pos_embed(src, pos) + + src2 = self.self_attn(q, k, value=src, attn_mask=src_mask)[0] + + # src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/utils.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..27da9bbd88843598238467951c8339d5f92c95a4 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/GroundingDINO/utils.py @@ -0,0 +1,270 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +import copy +import math + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +def _get_clones(module, N, layer_share=False): + # import ipdb; ipdb.set_trace() + if layer_share: + return nn.ModuleList([module for i in range(N)]) + else: + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def get_sine_pos_embed( + pos_tensor: torch.Tensor, + num_pos_feats: int = 128, + temperature: int = 10000, + exchange_xy: bool = True, +): + """generate sine position embedding from a position tensor + Args: + pos_tensor (torch.Tensor): shape: [..., n]. + num_pos_feats (int): projected shape for each float in the tensor. + temperature (int): temperature in the sine/cosine function. + exchange_xy (bool, optional): exchange pos x and pos y. \ + For example, input tensor is [x,y], the results will be [pos(y), pos(x)]. Defaults to True. + Returns: + pos_embed (torch.Tensor): shape: [..., n*num_pos_feats]. + """ + scale = 2 * math.pi + dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos_tensor.device) + dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats) + + def sine_func(x: torch.Tensor): + sin_x = x * scale / dim_t + sin_x = torch.stack((sin_x[..., 0::2].sin(), sin_x[..., 1::2].cos()), dim=3).flatten(2) + return sin_x + + pos_res = [sine_func(x) for x in pos_tensor.split([1] * pos_tensor.shape[-1], dim=-1)] + if exchange_xy: + pos_res[0], pos_res[1] = pos_res[1], pos_res[0] + pos_res = torch.cat(pos_res, dim=-1) + return pos_res + + +def gen_encoder_output_proposals( + memory: Tensor, memory_padding_mask: Tensor, spatial_shapes: Tensor, learnedwh=None +): + """ + Input: + - memory: bs, \sum{hw}, d_model + - memory_padding_mask: bs, \sum{hw} + - spatial_shapes: nlevel, 2 + - learnedwh: 2 + Output: + - output_memory: bs, \sum{hw}, d_model + - output_proposals: bs, \sum{hw}, 4 + """ + N_, S_, C_ = memory.shape + proposals = [] + _cur = 0 + for lvl, (H_, W_) in enumerate(spatial_shapes): + mask_flatten_ = memory_padding_mask[:, _cur : (_cur + H_ * W_)].view(N_, H_, W_, 1) + valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1) + valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1) + + # import ipdb; ipdb.set_trace() + + grid_y, grid_x = torch.meshgrid( + torch.linspace(0, H_ - 1, H_, dtype=torch.float32, device=memory.device), + torch.linspace(0, W_ - 1, W_, dtype=torch.float32, device=memory.device), + ) + grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) # H_, W_, 2 + + scale = torch.cat([valid_W.unsqueeze(-1), valid_H.unsqueeze(-1)], 1).view(N_, 1, 1, 2) + grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale + + if learnedwh is not None: + # import ipdb; ipdb.set_trace() + wh = torch.ones_like(grid) * learnedwh.sigmoid() * (2.0**lvl) + else: + wh = torch.ones_like(grid) * 0.05 * (2.0**lvl) + + # scale = torch.cat([W_[None].unsqueeze(-1), H_[None].unsqueeze(-1)], 1).view(1, 1, 1, 2).repeat(N_, 1, 1, 1) + # grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale + # wh = torch.ones_like(grid) / scale + proposal = torch.cat((grid, wh), -1).view(N_, -1, 4) + proposals.append(proposal) + _cur += H_ * W_ + # import ipdb; ipdb.set_trace() + output_proposals = torch.cat(proposals, 1) + output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all( + -1, keepdim=True + ) + output_proposals = torch.log(output_proposals / (1 - output_proposals)) # unsigmoid + output_proposals = output_proposals.masked_fill(memory_padding_mask.unsqueeze(-1), float("inf")) + output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf")) + + output_memory = memory + output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float(0)) + output_memory = output_memory.masked_fill(~output_proposals_valid, float(0)) + + # output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float('inf')) + # output_memory = output_memory.masked_fill(~output_proposals_valid, float('inf')) + + output_proposals = output_proposals.to(output_memory.dtype) + return output_memory, output_proposals + + +class RandomBoxPerturber: + def __init__( + self, x_noise_scale=0.2, y_noise_scale=0.2, w_noise_scale=0.2, h_noise_scale=0.2 + ) -> None: + self.noise_scale = torch.Tensor( + [x_noise_scale, y_noise_scale, w_noise_scale, h_noise_scale] + ) + + def __call__(self, refanchors: Tensor) -> Tensor: + nq, bs, query_dim = refanchors.shape + device = refanchors.device + + noise_raw = torch.rand_like(refanchors) + noise_scale = self.noise_scale.to(device)[:query_dim] + + new_refanchors = refanchors * (1 + (noise_raw - 0.5) * noise_scale) + return new_refanchors.clamp_(0, 1) + + +def sigmoid_focal_loss( + inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2, no_reduction=False +): + """ + Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. + Args: + inputs: A float tensor of arbitrary shape. + The predictions for each example. + targets: A float tensor with the same shape as inputs. Stores the binary + classification label for each element in inputs + (0 for the negative class and 1 for the positive class). + alpha: (optional) Weighting factor in range (0,1) to balance + positive vs negative examples. Default = -1 (no weighting). + gamma: Exponent of the modulating factor (1 - p_t) to + balance easy vs hard examples. + Returns: + Loss tensor + """ + prob = inputs.sigmoid() + ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") + p_t = prob * targets + (1 - prob) * (1 - targets) + loss = ce_loss * ((1 - p_t) ** gamma) + + if alpha >= 0: + alpha_t = alpha * targets + (1 - alpha) * (1 - targets) + loss = alpha_t * loss + + if no_reduction: + return loss + + return loss.mean(1).sum() / num_boxes + + +class MLP(nn.Module): + """Very simple multi-layer perceptron (also called FFN)""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +def _get_activation_fn(activation, d_model=256, batch_dim=0): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + if activation == "prelu": + return nn.PReLU() + if activation == "selu": + return F.selu + + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def gen_sineembed_for_position(pos_tensor): + # n_query, bs, _ = pos_tensor.size() + # sineembed_tensor = torch.zeros(n_query, bs, 256) + scale = 2 * math.pi + dim_t = torch.arange(128, dtype=torch.float32, device=pos_tensor.device) + dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode='floor')) / 128) + x_embed = pos_tensor[:, :, 0] * scale + y_embed = pos_tensor[:, :, 1] * scale + pos_x = x_embed[:, :, None] / dim_t + pos_y = y_embed[:, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) + pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) + if pos_tensor.size(-1) == 2: + pos = torch.cat((pos_y, pos_x), dim=2) + elif pos_tensor.size(-1) == 4: + w_embed = pos_tensor[:, :, 2] * scale + pos_w = w_embed[:, :, None] / dim_t + pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2) + + h_embed = pos_tensor[:, :, 3] * scale + pos_h = h_embed[:, :, None] / dim_t + pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2) + + pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) + else: + raise ValueError("Unknown pos_tensor shape(-1):{}".format(pos_tensor.size(-1))) + pos = pos.to(pos_tensor.dtype) + return pos + + +class ContrastiveEmbed(nn.Module): + def __init__(self, max_text_len=256): + """ + Args: + max_text_len: max length of text. + """ + super().__init__() + self.max_text_len = max_text_len + + def forward(self, x, text_dict): + """_summary_ + + Args: + x (_type_): _description_ + text_dict (_type_): _description_ + { + 'encoded_text': encoded_text, # bs, 195, d_model + 'text_token_mask': text_token_mask, # bs, 195 + # True for used tokens. False for padding tokens + } + Returns: + _type_: _description_ + """ + assert isinstance(text_dict, dict) + + y = text_dict["encoded_text"] + text_token_mask = text_dict["text_token_mask"] + + res = x @ y.transpose(-1, -2) + res.masked_fill_(~text_token_mask[:, None, :], float("-inf")) + + # padding to max_text_len + new_res = torch.full((*res.shape[:-1], self.max_text_len), float("-inf"), device=res.device, dtype=res.dtype) + new_res[..., : res.shape[-1]] = res + + return new_res diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/__init__.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e3413961d1d184b99835eb1e919b052d70298bc6 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/__init__.py @@ -0,0 +1,18 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +from .GroundingDINO import build_groundingdino + + +def build_model(args): + # we use register to maintain models from catdet6 on. + from .registry import MODULE_BUILD_FUNCS + + assert args.modelname in MODULE_BUILD_FUNCS._module_dict + build_func = MODULE_BUILD_FUNCS.get(args.modelname) + model = build_func(args) + return model diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87900c238577a114de73990d868beb39462cfd3f Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/registry.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..2d22a59eec79a2a19b83fa1779f2adaf5753aec6 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/models/registry.py @@ -0,0 +1,66 @@ +# ------------------------------------------------------------------------ +# Grounding DINO +# url: https://github.com/IDEA-Research/GroundingDINO +# Copyright (c) 2023 IDEA. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +# -*- coding: utf-8 -*- +# @Author: Yihao Chen +# @Date: 2021-08-16 16:03:17 +# @Last Modified by: Shilong Liu +# @Last Modified time: 2022-01-23 15:26 +# modified from mmcv + +import inspect +from functools import partial + + +class Registry(object): + def __init__(self, name): + self._name = name + self._module_dict = dict() + + def __repr__(self): + format_str = self.__class__.__name__ + "(name={}, items={})".format( + self._name, list(self._module_dict.keys()) + ) + return format_str + + def __len__(self): + return len(self._module_dict) + + @property + def name(self): + return self._name + + @property + def module_dict(self): + return self._module_dict + + def get(self, key): + return self._module_dict.get(key, None) + + def registe_with_name(self, module_name=None, force=False): + return partial(self.register, module_name=module_name, force=force) + + def register(self, module_build_function, module_name=None, force=False): + """Register a module build function. + Args: + module (:obj:`nn.Module`): Module to be registered. + """ + if not inspect.isfunction(module_build_function): + raise TypeError( + "module_build_function must be a function, but got {}".format( + type(module_build_function) + ) + ) + if module_name is None: + module_name = module_build_function.__name__ + if not force and module_name in self._module_dict: + raise KeyError("{} is already registered in {}".format(module_name, self.name)) + self._module_dict[module_name] = module_build_function + + return module_build_function + + +MODULE_BUILD_FUNCS = Registry("model build functions") diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__init__.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..168f9979a4623806934b0ff1102ac166704e7dec --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84df51d9a1a19906b2a987924e8b0abbc5972f54 Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/box_ops.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/box_ops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a10f34511df4ee7bbccf9bec819f36450442e3b Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/box_ops.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/get_tokenlizer.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/get_tokenlizer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c11731abbe3497787c4ee851dedf435720df20c Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/get_tokenlizer.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/inference.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/inference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2da245fa6fed294f378c6e24d8a4be294c7ebe3 Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/inference.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/misc.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/misc.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3930c4130c24e20cdda9df1cbc8cd6fc67f563be Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/misc.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/slconfig.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/slconfig.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dff34dbdf701d3ce47930726c8f0d41eaed8bbb Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/slconfig.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/utils.cpython-310.pyc b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1c171fd6ec49a7cdd25a58a33108002f926cbdd Binary files /dev/null and b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/__pycache__/utils.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/box_ops.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/box_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..781068d294e576954edb4bd07b6e0f30e4e1bcd9 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/box_ops.py @@ -0,0 +1,140 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Utilities for bounding box manipulation and GIoU. +""" +import torch +from torchvision.ops.boxes import box_area + + +def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.stack(b, dim=-1) + + +def box_xyxy_to_cxcywh(x): + x0, y0, x1, y1 = x.unbind(-1) + b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)] + return torch.stack(b, dim=-1) + + +# modified from torchvision to also return the union +def box_iou(boxes1, boxes2): + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + # import ipdb; ipdb.set_trace() + lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] + rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] + + wh = (rb - lt).clamp(min=0) # [N,M,2] + inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] + + union = area1[:, None] + area2 - inter + + iou = inter / (union + 1e-6) + return iou, union + + +def generalized_box_iou(boxes1, boxes2): + """ + Generalized IoU from https://giou.stanford.edu/ + + The boxes should be in [x0, y0, x1, y1] format + + Returns a [N, M] pairwise matrix, where N = len(boxes1) + and M = len(boxes2) + """ + # degenerate boxes gives inf / nan results + # so do an early check + assert (boxes1[:, 2:] >= boxes1[:, :2]).all() + assert (boxes2[:, 2:] >= boxes2[:, :2]).all() + # except: + # import ipdb; ipdb.set_trace() + iou, union = box_iou(boxes1, boxes2) + + lt = torch.min(boxes1[:, None, :2], boxes2[:, :2]) + rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:]) + + wh = (rb - lt).clamp(min=0) # [N,M,2] + area = wh[:, :, 0] * wh[:, :, 1] + + return iou - (area - union) / (area + 1e-6) + + +# modified from torchvision to also return the union +def box_iou_pairwise(boxes1, boxes2): + area1 = box_area(boxes1) + area2 = box_area(boxes2) + + lt = torch.max(boxes1[:, :2], boxes2[:, :2]) # [N,2] + rb = torch.min(boxes1[:, 2:], boxes2[:, 2:]) # [N,2] + + wh = (rb - lt).clamp(min=0) # [N,2] + inter = wh[:, 0] * wh[:, 1] # [N] + + union = area1 + area2 - inter + + iou = inter / union + return iou, union + + +def generalized_box_iou_pairwise(boxes1, boxes2): + """ + Generalized IoU from https://giou.stanford.edu/ + + Input: + - boxes1, boxes2: N,4 + Output: + - giou: N, 4 + """ + # degenerate boxes gives inf / nan results + # so do an early check + assert (boxes1[:, 2:] >= boxes1[:, :2]).all() + assert (boxes2[:, 2:] >= boxes2[:, :2]).all() + assert boxes1.shape == boxes2.shape + iou, union = box_iou_pairwise(boxes1, boxes2) # N, 4 + + lt = torch.min(boxes1[:, :2], boxes2[:, :2]) + rb = torch.max(boxes1[:, 2:], boxes2[:, 2:]) + + wh = (rb - lt).clamp(min=0) # [N,2] + area = wh[:, 0] * wh[:, 1] + + return iou - (area - union) / area + + +def masks_to_boxes(masks): + """Compute the bounding boxes around the provided masks + + The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. + + Returns a [N, 4] tensors, with the boxes in xyxy format + """ + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device) + + h, w = masks.shape[-2:] + + y = torch.arange(0, h, dtype=torch.float) + x = torch.arange(0, w, dtype=torch.float) + y, x = torch.meshgrid(y, x) + + x_mask = masks * x.unsqueeze(0) + x_max = x_mask.flatten(1).max(-1)[0] + x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + y_mask = masks * y.unsqueeze(0) + y_max = y_mask.flatten(1).max(-1)[0] + y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] + + return torch.stack([x_min, y_min, x_max, y_max], 1) + + +if __name__ == "__main__": + x = torch.rand(5, 4) + y = torch.rand(3, 4) + iou, union = box_iou(x, y) + import ipdb + + ipdb.set_trace() diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/get_tokenlizer.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/get_tokenlizer.py new file mode 100644 index 0000000000000000000000000000000000000000..b7b5d72aef873453361cc019427ba31b08c11798 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/get_tokenlizer.py @@ -0,0 +1,37 @@ +from transformers import AutoTokenizer, BertModel, BertTokenizer, RobertaModel, RobertaTokenizerFast + + +def get_tokenlizer(text_encoder_type, bert_base_uncased_path): + if not isinstance(text_encoder_type, str): + # print("text_encoder_type is not a str") + if hasattr(text_encoder_type, "text_encoder_type"): + text_encoder_type = text_encoder_type.text_encoder_type + elif text_encoder_type.get("text_encoder_type", False): + text_encoder_type = text_encoder_type.get("text_encoder_type") + else: + raise ValueError( + "Unknown type of text_encoder_type: {}".format(type(text_encoder_type)) + ) + + # solve huggingface connect issue + if is_bert_model_use_local_path(bert_base_uncased_path) and text_encoder_type == "bert-base-uncased": + print("use local bert model path: {}".format(bert_base_uncased_path)) + return AutoTokenizer.from_pretrained(bert_base_uncased_path) + + print("final text_encoder_type: {}".format(text_encoder_type)) + + tokenizer = AutoTokenizer.from_pretrained(text_encoder_type) + return tokenizer + + +def get_pretrained_language_model(text_encoder_type, bert_base_uncased_path): + if text_encoder_type == "bert-base-uncased": + if is_bert_model_use_local_path(bert_base_uncased_path): + return BertModel.from_pretrained(bert_base_uncased_path) + return BertModel.from_pretrained(text_encoder_type) + if text_encoder_type == "roberta-base": + return RobertaModel.from_pretrained(text_encoder_type) + raise ValueError("Unknown text_encoder_type {}".format(text_encoder_type)) + +def is_bert_model_use_local_path(bert_base_uncased_path): + return bert_base_uncased_path is not None and len(bert_base_uncased_path) > 0 diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/inference.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..fe8077e103e5e269ad175ab056699136cf1c74b1 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/inference.py @@ -0,0 +1,259 @@ +from typing import Tuple, List + +import re +import cv2 +import numpy as np +import supervision as sv +import torch +from PIL import Image +from torchvision.ops import box_convert + +import groundingdino.datasets.transforms as T +from groundingdino.models import build_model +from groundingdino.util.misc import clean_state_dict +from groundingdino.util.slconfig import SLConfig +from groundingdino.util.utils import get_phrases_from_posmap + +# ---------------------------------------------------------------------------------------------------------------------- +# OLD API +# ---------------------------------------------------------------------------------------------------------------------- + + +def preprocess_caption(caption: str) -> str: + result = caption.lower().strip() + if result.endswith("."): + return result + return result + "." + + +def load_model(model_config_path: str, model_checkpoint_path: str, device: str = "cuda"): + args = SLConfig.fromfile(model_config_path) + args.device = device + model = build_model(args) + checkpoint = torch.load(model_checkpoint_path, map_location="cpu") + model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) + model.eval() + return model + + +def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]: + transform = T.Compose( + [ + T.RandomResize([800], max_size=1333), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + image_source = Image.open(image_path).convert("RGB") + image = np.asarray(image_source) + image_transformed, _ = transform(image_source, None) + return image, image_transformed + + +def predict( + model, + image: torch.Tensor, + caption: str, + box_threshold: float, + text_threshold: float, + device: str = "cuda" +) -> Tuple[torch.Tensor, torch.Tensor, List[str]]: + caption = preprocess_caption(caption=caption) + + model = model.to(device) + image = image.to(device) + + with torch.no_grad(): + outputs = model(image[None], captions=[caption]) + + prediction_logits = outputs["pred_logits"].cpu().sigmoid()[0] # prediction_logits.shape = (nq, 256) + prediction_boxes = outputs["pred_boxes"].cpu()[0] # prediction_boxes.shape = (nq, 4) + + mask = prediction_logits.max(dim=1)[0] > box_threshold + logits = prediction_logits[mask] # logits.shape = (n, 256) + boxes = prediction_boxes[mask] # boxes.shape = (n, 4) + + tokenizer = model.tokenizer + tokenized = tokenizer(caption) + + phrases = [ + get_phrases_from_posmap(logit > text_threshold, tokenized, tokenizer).replace('.', '') + for logit + in logits + ] + + return boxes, logits.max(dim=1)[0], phrases + + +def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str]) -> np.ndarray: + h, w, _ = image_source.shape + boxes = boxes * torch.Tensor([w, h, w, h]) + xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() + detections = sv.Detections(xyxy=xyxy) + + labels = [ + f"{phrase} {logit:.2f}" + for phrase, logit + in zip(phrases, logits) + ] + + box_annotator = sv.BoxAnnotator() + # box_annotator = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX) + annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR) + annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels) + return annotated_frame + + +# ---------------------------------------------------------------------------------------------------------------------- +# NEW API +# ---------------------------------------------------------------------------------------------------------------------- + + +class Model: + + def __init__( + self, + model_config_path: str, + model_checkpoint_path: str, + device: str = "cuda" + ): + + self.model = load_model( + model_config_path=model_config_path, + model_checkpoint_path=model_checkpoint_path, + device=device + ).to(device) + self.device = device + + def predict_with_caption( + self, + image: np.ndarray, + caption: str, + box_threshold: float = 0.35, + text_threshold: float = 0.25 + ) -> Tuple[sv.Detections, List[str]]: + """ + import cv2 + + image = cv2.imread(IMAGE_PATH) + + model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH) + detections, labels = model.predict_with_caption( + image=image, + caption=caption, + box_threshold=BOX_THRESHOLD, + text_threshold=TEXT_THRESHOLD + ) + + import supervision as sv + + box_annotator = sv.BoxAnnotator() + annotated_image = box_annotator.annotate(scene=image, detections=detections, labels=labels) + """ + processed_image = Model.preprocess_image(image_bgr=image).to(self.device) + boxes, logits, phrases = predict( + model=self.model, + image=processed_image, + caption=caption, + box_threshold=box_threshold, + text_threshold=text_threshold, + device=self.device) + source_h, source_w, _ = image.shape + detections = Model.post_process_result( + source_h=source_h, + source_w=source_w, + boxes=boxes, + logits=logits) + return detections, phrases + + def predict_with_classes( + self, + image: np.ndarray, + classes: List[str], + box_threshold: float, + text_threshold: float + ) -> sv.Detections: + """ + import cv2 + + image = cv2.imread(IMAGE_PATH) + + model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH) + detections = model.predict_with_classes( + image=image, + classes=CLASSES, + box_threshold=BOX_THRESHOLD, + text_threshold=TEXT_THRESHOLD + ) + + + import supervision as sv + + box_annotator = sv.BoxAnnotator() + annotated_image = box_annotator.annotate(scene=image, detections=detections) + """ + caption = ". ".join(classes) + processed_image = Model.preprocess_image(image_bgr=image).to(self.device) + boxes, logits, phrases = predict( + model=self.model, + image=processed_image, + caption=caption, + box_threshold=box_threshold, + text_threshold=text_threshold, + device=self.device) + source_h, source_w, _ = image.shape + detections = Model.post_process_result( + source_h=source_h, + source_w=source_w, + boxes=boxes, + logits=logits) + class_id = Model.phrases2classes(phrases=phrases, classes=classes) + detections.class_id = class_id + return detections + + @staticmethod + def preprocess_image(image_bgr: np.ndarray) -> torch.Tensor: + transform = T.Compose( + [ + T.RandomResize([800], max_size=1333), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ] + ) + image_pillow = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)) + image_transformed, _ = transform(image_pillow, None) + return image_transformed + + @staticmethod + def post_process_result( + source_h: int, + source_w: int, + boxes: torch.Tensor, + logits: torch.Tensor + ) -> sv.Detections: + boxes = boxes * torch.Tensor([source_w, source_h, source_w, source_h]) + xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy() + confidence = logits.numpy() + return sv.Detections(xyxy=xyxy, confidence=confidence) + + @staticmethod + def phrases2classes(phrases: List[str], classes: List[str]) -> np.ndarray: + class_ids = [] + for phrase in phrases: + try: + # class_ids.append(classes.index(phrase)) + class_ids.append(Model.find_index(phrase, classes)) + except ValueError: + class_ids.append(None) + return np.array(class_ids) + + @staticmethod + def find_index(string, lst): + # if meet string like "lake river" will only keep "lake" + # this is an hack implementation for visualization which will be updated in the future + string = string.lower().split()[0] + for i, s in enumerate(lst): + if string in s.lower(): + return i + print("There's a wrong phrase happen, this is because of our post-process merged wrong tokens, which will be modified in the future. We will assign it with a random label at this time.") + return 0 \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/logger.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..18145f54c927abd59b95f3fa6e6da8002bc2ce97 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/logger.py @@ -0,0 +1,93 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +import functools +import logging +import os +import sys + +from termcolor import colored + + +class _ColorfulFormatter(logging.Formatter): + def __init__(self, *args, **kwargs): + self._root_name = kwargs.pop("root_name") + "." + self._abbrev_name = kwargs.pop("abbrev_name", "") + if len(self._abbrev_name): + self._abbrev_name = self._abbrev_name + "." + super(_ColorfulFormatter, self).__init__(*args, **kwargs) + + def formatMessage(self, record): + record.name = record.name.replace(self._root_name, self._abbrev_name) + log = super(_ColorfulFormatter, self).formatMessage(record) + if record.levelno == logging.WARNING: + prefix = colored("WARNING", "red", attrs=["blink"]) + elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL: + prefix = colored("ERROR", "red", attrs=["blink", "underline"]) + else: + return log + return prefix + " " + log + + +# so that calling setup_logger multiple times won't add many handlers +@functools.lru_cache() +def setup_logger(output=None, distributed_rank=0, *, color=True, name="imagenet", abbrev_name=None): + """ + Initialize the detectron2 logger and set its verbosity level to "INFO". + + Args: + output (str): a file name or a directory to save log. If None, will not save log file. + If ends with ".txt" or ".log", assumed to be a file name. + Otherwise, logs will be saved to `output/log.txt`. + name (str): the root module name of this logger + + Returns: + logging.Logger: a logger + """ + logger = logging.getLogger(name) + logger.setLevel(logging.DEBUG) + logger.propagate = False + + if abbrev_name is None: + abbrev_name = name + + plain_formatter = logging.Formatter( + "[%(asctime)s.%(msecs)03d]: %(message)s", datefmt="%m/%d %H:%M:%S" + ) + # stdout logging: master only + if distributed_rank == 0: + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(logging.DEBUG) + if color: + formatter = _ColorfulFormatter( + colored("[%(asctime)s.%(msecs)03d]: ", "green") + "%(message)s", + datefmt="%m/%d %H:%M:%S", + root_name=name, + abbrev_name=str(abbrev_name), + ) + else: + formatter = plain_formatter + ch.setFormatter(formatter) + logger.addHandler(ch) + + # file logging: all workers + if output is not None: + if output.endswith(".txt") or output.endswith(".log"): + filename = output + else: + filename = os.path.join(output, "log.txt") + if distributed_rank > 0: + filename = filename + f".rank{distributed_rank}" + os.makedirs(os.path.dirname(filename), exist_ok=True) + + fh = logging.StreamHandler(_cached_log_stream(filename)) + fh.setLevel(logging.DEBUG) + fh.setFormatter(plain_formatter) + logger.addHandler(fh) + + return logger + + +# cache the opened file object, so that different calls to `setup_logger` +# with the same file name can safely write to the same file. +@functools.lru_cache(maxsize=None) +def _cached_log_stream(filename): + return open(filename, "a") diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/misc.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..d64b84ef24bea0c98e76824feb1903f6bfebe7a5 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/misc.py @@ -0,0 +1,717 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Misc functions, including distributed helpers. + +Mostly copy-paste from torchvision references. +""" +import colorsys +import datetime +import functools +import io +import json +import os +import pickle +import subprocess +import time +from collections import OrderedDict, defaultdict, deque +from typing import List, Optional + +import numpy as np +import torch +import torch.distributed as dist + +# needed due to empty tensor bug in pytorch and torchvision 0.5 +import torchvision +from torch import Tensor + +__torchvision_need_compat_flag = float(torchvision.__version__.split(".")[1]) < 7 +if __torchvision_need_compat_flag: + from torchvision.ops import _new_empty_tensor + from torchvision.ops.misc import _output_size + + +class SmoothedValue(object): + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda") + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + d = torch.tensor(list(self.deque)) + if d.shape[0] == 0: + return 0 + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + if os.environ.get("SHILONG_AMP", None) == "1": + eps = 1e-4 + else: + eps = 1e-6 + return self.total / (self.count + eps) + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +@functools.lru_cache() +def _get_global_gloo_group(): + """ + Return a process group based on gloo backend, containing all the ranks + The result is cached. + """ + + if dist.get_backend() == "nccl": + return dist.new_group(backend="gloo") + + return dist.group.WORLD + + +def all_gather_cpu(data): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + + world_size = get_world_size() + if world_size == 1: + return [data] + + cpu_group = _get_global_gloo_group() + + buffer = io.BytesIO() + torch.save(data, buffer) + data_view = buffer.getbuffer() + device = "cuda" if cpu_group is None else "cpu" + tensor = torch.ByteTensor(data_view).to(device) + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device=device, dtype=torch.long) + size_list = [torch.tensor([0], device=device, dtype=torch.long) for _ in range(world_size)] + if cpu_group is None: + dist.all_gather(size_list, local_size) + else: + print("gathering on cpu") + dist.all_gather(size_list, local_size, group=cpu_group) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + assert isinstance(local_size.item(), int) + local_size = int(local_size.item()) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=device)) + if local_size != max_size: + padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device=device) + tensor = torch.cat((tensor, padding), dim=0) + if cpu_group is None: + dist.all_gather(tensor_list, tensor) + else: + dist.all_gather(tensor_list, tensor, group=cpu_group) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + tensor = torch.split(tensor, [size, max_size - size], dim=0)[0] + buffer = io.BytesIO(tensor.cpu().numpy()) + obj = torch.load(buffer) + data_list.append(obj) + + return data_list + + +def all_gather(data): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + + if os.getenv("CPU_REDUCE") == "1": + return all_gather_cpu(data) + + world_size = get_world_size() + if world_size == 1: + return [data] + + # serialized to a Tensor + buffer = pickle.dumps(data) + storage = torch.ByteStorage.from_buffer(buffer) + tensor = torch.ByteTensor(storage).to("cuda") + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device="cuda") + size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)] + dist.all_gather(size_list, local_size) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda")) + if local_size != max_size: + padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda") + tensor = torch.cat((tensor, padding), dim=0) + dist.all_gather(tensor_list, tensor) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + buffer = tensor.cpu().numpy().tobytes()[:size] + data_list.append(pickle.loads(buffer)) + + return data_list + + +def reduce_dict(input_dict, average=True): + """ + Args: + input_dict (dict): all the values will be reduced + average (bool): whether to do average or sum + Reduce the values in the dictionary from all processes so that all processes + have the averaged results. Returns a dict with the same fields as + input_dict, after reduction. + """ + world_size = get_world_size() + if world_size < 2: + return input_dict + with torch.no_grad(): + names = [] + values = [] + # sort the keys so that they are consistent across processes + for k in sorted(input_dict.keys()): + names.append(k) + values.append(input_dict[k]) + values = torch.stack(values, dim=0) + dist.all_reduce(values) + if average: + values /= world_size + reduced_dict = {k: v for k, v in zip(names, values)} + return reduced_dict + + +class MetricLogger(object): + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + # print(name, str(meter)) + # import ipdb;ipdb.set_trace() + if meter.count > 0: + loss_str.append("{}: {}".format(name, str(meter))) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None, logger=None): + if logger is None: + print_func = print + else: + print_func = logger.info + + i = 0 + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(len(iterable)))) + "d" + if torch.cuda.is_available(): + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + "max mem: {memory:.0f}", + ] + ) + else: + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + ) + MB = 1024.0 * 1024.0 + for obj in iterable: + data_time.update(time.time() - end) + yield obj + # import ipdb; ipdb.set_trace() + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == len(iterable) - 1: + eta_seconds = iter_time.global_avg * (len(iterable) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print_func( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + else: + print_func( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print_func( + "{} Total time: {} ({:.4f} s / it)".format( + header, total_time_str, total_time / len(iterable) + ) + ) + + +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + subprocess.check_output(["git", "diff"], cwd=cwd) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommited changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +def collate_fn(batch): + # import ipdb; ipdb.set_trace() + batch = list(zip(*batch)) + batch[0] = nested_tensor_from_tensor_list(batch[0]) + return tuple(batch) + + +def _max_by_axis(the_list): + # type: (List[List[int]]) -> List[int] + maxes = the_list[0] + for sublist in the_list[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + +class NestedTensor(object): + def __init__(self, tensors, mask: Optional[Tensor]): + self.tensors = tensors + self.mask = mask + if mask == "auto": + self.mask = torch.zeros_like(tensors).to(tensors.device) + if self.mask.dim() == 3: + self.mask = self.mask.sum(0).to(bool) + elif self.mask.dim() == 4: + self.mask = self.mask.sum(1).to(bool) + else: + raise ValueError( + "tensors dim must be 3 or 4 but {}({})".format( + self.tensors.dim(), self.tensors.shape + ) + ) + + def imgsize(self): + res = [] + for i in range(self.tensors.shape[0]): + mask = self.mask[i] + maxH = (~mask).sum(0).max() + maxW = (~mask).sum(1).max() + res.append(torch.Tensor([maxH, maxW])) + return res + + def to(self, device): + # type: (Device) -> NestedTensor # noqa + cast_tensor = self.tensors.to(device) + mask = self.mask + if mask is not None: + assert mask is not None + cast_mask = mask.to(device) + else: + cast_mask = None + return NestedTensor(cast_tensor, cast_mask) + + def to_img_list_single(self, tensor, mask): + assert tensor.dim() == 3, "dim of tensor should be 3 but {}".format(tensor.dim()) + maxH = (~mask).sum(0).max() + maxW = (~mask).sum(1).max() + img = tensor[:, :maxH, :maxW] + return img + + def to_img_list(self): + """remove the padding and convert to img list + + Returns: + [type]: [description] + """ + if self.tensors.dim() == 3: + return self.to_img_list_single(self.tensors, self.mask) + else: + res = [] + for i in range(self.tensors.shape[0]): + tensor_i = self.tensors[i] + mask_i = self.mask[i] + res.append(self.to_img_list_single(tensor_i, mask_i)) + return res + + @property + def device(self): + return self.tensors.device + + def decompose(self): + return self.tensors, self.mask + + def __repr__(self): + return str(self.tensors) + + @property + def shape(self): + return {"tensors.shape": self.tensors.shape, "mask.shape": self.mask.shape} + + +def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): + # TODO make this more general + if tensor_list[0].ndim == 3: + if torchvision._is_tracing(): + # nested_tensor_from_tensor_list() does not export well to ONNX + # call _onnx_nested_tensor_from_tensor_list() instead + return _onnx_nested_tensor_from_tensor_list(tensor_list) + + # TODO make it support different-sized images + max_size = _max_by_axis([list(img.shape) for img in tensor_list]) + # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) + batch_shape = [len(tensor_list)] + max_size + b, c, h, w = batch_shape + dtype = tensor_list[0].dtype + device = tensor_list[0].device + tensor = torch.zeros(batch_shape, dtype=dtype, device=device) + mask = torch.ones((b, h, w), dtype=torch.bool, device=device) + for img, pad_img, m in zip(tensor_list, tensor, mask): + pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + m[: img.shape[1], : img.shape[2]] = False + else: + raise ValueError("not supported") + return NestedTensor(tensor, mask) + + +# _onnx_nested_tensor_from_tensor_list() is an implementation of +# nested_tensor_from_tensor_list() that is supported by ONNX tracing. +@torch.jit.unused +def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor: + max_size = [] + for i in range(tensor_list[0].dim()): + max_size_i = torch.max( + torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32) + ).to(torch.int64) + max_size.append(max_size_i) + max_size = tuple(max_size) + + # work around for + # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + # m[: img.shape[1], :img.shape[2]] = False + # which is not yet supported in onnx + padded_imgs = [] + padded_masks = [] + for img in tensor_list: + padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))] + padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0])) + padded_imgs.append(padded_img) + + m = torch.zeros_like(img[0], dtype=torch.int, device=img.device) + padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1) + padded_masks.append(padded_mask.to(torch.bool)) + + tensor = torch.stack(padded_imgs) + mask = torch.stack(padded_masks) + + return NestedTensor(tensor, mask=mask) + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(*args, **kwargs): + if is_main_process(): + torch.save(*args, **kwargs) + + +def init_distributed_mode(args): + if "WORLD_SIZE" in os.environ and os.environ["WORLD_SIZE"] != "": # 'RANK' in os.environ and + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = args.local_rank = int(os.environ["LOCAL_RANK"]) + + # launch by torch.distributed.launch + # Single node + # python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 1 --rank 0 ... + # Multi nodes + # python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 2 --rank 0 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' ... + # python -m torch.distributed.launch --nproc_per_node=8 main.py --world-size 2 --rank 1 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' ... + # args.rank = int(os.environ.get('OMPI_COMM_WORLD_RANK')) + # local_world_size = int(os.environ['GPU_PER_NODE_COUNT']) + # args.world_size = args.world_size * local_world_size + # args.gpu = args.local_rank = int(os.environ['LOCAL_RANK']) + # args.rank = args.rank * local_world_size + args.local_rank + print( + "world size: {}, rank: {}, local rank: {}".format( + args.world_size, args.rank, args.local_rank + ) + ) + print(json.dumps(dict(os.environ), indent=2)) + elif "SLURM_PROCID" in os.environ: + args.rank = int(os.environ["SLURM_PROCID"]) + args.gpu = args.local_rank = int(os.environ["SLURM_LOCALID"]) + args.world_size = int(os.environ["SLURM_NPROCS"]) + + print( + "world size: {}, world rank: {}, local rank: {}, device_count: {}".format( + args.world_size, args.rank, args.local_rank, torch.cuda.device_count() + ) + ) + else: + print("Not using distributed mode") + args.distributed = False + args.world_size = 1 + args.rank = 0 + args.local_rank = 0 + return + + print("world_size:{} rank:{} local_rank:{}".format(args.world_size, args.rank, args.local_rank)) + args.distributed = True + torch.cuda.set_device(args.local_rank) + args.dist_backend = "nccl" + print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True) + + torch.distributed.init_process_group( + backend=args.dist_backend, + world_size=args.world_size, + rank=args.rank, + init_method=args.dist_url, + ) + + print("Before torch.distributed.barrier()") + torch.distributed.barrier() + print("End torch.distributed.barrier()") + setup_for_distributed(args.rank == 0) + + +@torch.no_grad() +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + if target.numel() == 0: + return [torch.zeros([], device=output.device)] + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +@torch.no_grad() +def accuracy_onehot(pred, gt): + """_summary_ + + Args: + pred (_type_): n, c + gt (_type_): n, c + """ + tp = ((pred - gt).abs().sum(-1) < 1e-4).float().sum() + acc = tp / gt.shape[0] * 100 + return acc + + +def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None): + # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor + """ + Equivalent to nn.functional.interpolate, but with support for empty batch sizes. + This will eventually be supported natively by PyTorch, and this + class can go away. + """ + if __torchvision_need_compat_flag < 0.7: + if input.numel() > 0: + return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners) + + output_shape = _output_size(2, input, size, scale_factor) + output_shape = list(input.shape[:-2]) + list(output_shape) + return _new_empty_tensor(input, output_shape) + else: + return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners) + + +class color_sys: + def __init__(self, num_colors) -> None: + self.num_colors = num_colors + colors = [] + for i in np.arange(0.0, 360.0, 360.0 / num_colors): + hue = i / 360.0 + lightness = (50 + np.random.rand() * 10) / 100.0 + saturation = (90 + np.random.rand() * 10) / 100.0 + colors.append( + tuple([int(j * 255) for j in colorsys.hls_to_rgb(hue, lightness, saturation)]) + ) + self.colors = colors + + def __call__(self, idx): + return self.colors[idx] + + +def inverse_sigmoid(x, eps=1e-3): + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + + +def clean_state_dict(state_dict): + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + if k[:7] == "module.": + k = k[7:] # remove `module.` + new_state_dict[k] = v + return new_state_dict diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/slconfig.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/slconfig.py new file mode 100644 index 0000000000000000000000000000000000000000..3f293e3aff215a3c7c2f7d21d27853493b6ebfbc --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/slconfig.py @@ -0,0 +1,427 @@ +# ========================================================== +# Modified from mmcv +# ========================================================== +import ast +import os.path as osp +import shutil +import sys +import tempfile +from argparse import Action +from importlib import import_module +import platform + +from addict import Dict +from yapf.yapflib.yapf_api import FormatCode + +BASE_KEY = "_base_" +DELETE_KEY = "_delete_" +RESERVED_KEYS = ["filename", "text", "pretty_text", "get", "dump", "merge_from_dict"] + + +def check_file_exist(filename, msg_tmpl='file "{}" does not exist'): + if not osp.isfile(filename): + raise FileNotFoundError(msg_tmpl.format(filename)) + + +class ConfigDict(Dict): + def __missing__(self, name): + raise KeyError(name) + + def __getattr__(self, name): + try: + value = super(ConfigDict, self).__getattr__(name) + except KeyError: + ex = AttributeError(f"'{self.__class__.__name__}' object has no " f"attribute '{name}'") + except Exception as e: + ex = e + else: + return value + raise ex + + +class SLConfig(object): + """ + config files. + only support .py file as config now. + + ref: mmcv.utils.config + + Example: + >>> cfg = Config(dict(a=1, b=dict(b1=[0, 1]))) + >>> cfg.a + 1 + >>> cfg.b + {'b1': [0, 1]} + >>> cfg.b.b1 + [0, 1] + >>> cfg = Config.fromfile('tests/data/config/a.py') + >>> cfg.filename + "/home/kchen/projects/mmcv/tests/data/config/a.py" + >>> cfg.item4 + 'test' + >>> cfg + "Config [path: /home/kchen/projects/mmcv/tests/data/config/a.py]: " + "{'item1': [1, 2], 'item2': {'a': 0}, 'item3': True, 'item4': 'test'}" + """ + + @staticmethod + def _validate_py_syntax(filename): + with open(filename) as f: + content = f.read() + try: + ast.parse(content) + except SyntaxError: + raise SyntaxError("There are syntax errors in config " f"file {filename}") + + @staticmethod + def _file2dict(filename): + filename = osp.abspath(osp.expanduser(filename)) + check_file_exist(filename) + if filename.lower().endswith(".py"): + with tempfile.TemporaryDirectory() as temp_config_dir: + temp_config_file = tempfile.NamedTemporaryFile(dir=temp_config_dir, suffix=".py") + temp_config_name = osp.basename(temp_config_file.name) + if platform.system() == 'Windows': + temp_config_file.close() + shutil.copyfile(filename, osp.join(temp_config_dir, temp_config_name)) + temp_module_name = osp.splitext(temp_config_name)[0] + sys.path.insert(0, temp_config_dir) + SLConfig._validate_py_syntax(filename) + mod = import_module(temp_module_name) + sys.path.pop(0) + cfg_dict = { + name: value for name, value in mod.__dict__.items() if not name.startswith("__") + } + # delete imported module + del sys.modules[temp_module_name] + # close temp file + temp_config_file.close() + elif filename.lower().endswith((".yml", ".yaml", ".json")): + from .slio import slload + + cfg_dict = slload(filename) + else: + raise IOError("Only py/yml/yaml/json type are supported now!") + + cfg_text = filename + "\n" + with open(filename, "r") as f: + cfg_text += f.read() + + # parse the base file + if BASE_KEY in cfg_dict: + cfg_dir = osp.dirname(filename) + base_filename = cfg_dict.pop(BASE_KEY) + base_filename = base_filename if isinstance(base_filename, list) else [base_filename] + + cfg_dict_list = list() + cfg_text_list = list() + for f in base_filename: + _cfg_dict, _cfg_text = SLConfig._file2dict(osp.join(cfg_dir, f)) + cfg_dict_list.append(_cfg_dict) + cfg_text_list.append(_cfg_text) + + base_cfg_dict = dict() + for c in cfg_dict_list: + if len(base_cfg_dict.keys() & c.keys()) > 0: + raise KeyError("Duplicate key is not allowed among bases") + # TODO Allow the duplicate key while warnning user + base_cfg_dict.update(c) + + base_cfg_dict = SLConfig._merge_a_into_b(cfg_dict, base_cfg_dict) + cfg_dict = base_cfg_dict + + # merge cfg_text + cfg_text_list.append(cfg_text) + cfg_text = "\n".join(cfg_text_list) + + return cfg_dict, cfg_text + + @staticmethod + def _merge_a_into_b(a, b): + """merge dict `a` into dict `b` (non-inplace). + values in `a` will overwrite `b`. + copy first to avoid inplace modification + + Args: + a ([type]): [description] + b ([type]): [description] + + Returns: + [dict]: [description] + """ + # import ipdb; ipdb.set_trace() + if not isinstance(a, dict): + return a + + b = b.copy() + for k, v in a.items(): + if isinstance(v, dict) and k in b and not v.pop(DELETE_KEY, False): + + if not isinstance(b[k], dict) and not isinstance(b[k], list): + # if : + # import ipdb; ipdb.set_trace() + raise TypeError( + f"{k}={v} in child config cannot inherit from base " + f"because {k} is a dict in the child config but is of " + f"type {type(b[k])} in base config. You may set " + f"`{DELETE_KEY}=True` to ignore the base config" + ) + b[k] = SLConfig._merge_a_into_b(v, b[k]) + elif isinstance(b, list): + try: + _ = int(k) + except: + raise TypeError( + f"b is a list, " f"index {k} should be an int when input but {type(k)}" + ) + b[int(k)] = SLConfig._merge_a_into_b(v, b[int(k)]) + else: + b[k] = v + + return b + + @staticmethod + def fromfile(filename): + cfg_dict, cfg_text = SLConfig._file2dict(filename) + return SLConfig(cfg_dict, cfg_text=cfg_text, filename=filename) + + def __init__(self, cfg_dict=None, cfg_text=None, filename=None): + if cfg_dict is None: + cfg_dict = dict() + elif not isinstance(cfg_dict, dict): + raise TypeError("cfg_dict must be a dict, but " f"got {type(cfg_dict)}") + for key in cfg_dict: + if key in RESERVED_KEYS: + raise KeyError(f"{key} is reserved for config file") + + super(SLConfig, self).__setattr__("_cfg_dict", ConfigDict(cfg_dict)) + super(SLConfig, self).__setattr__("_filename", filename) + if cfg_text: + text = cfg_text + elif filename: + with open(filename, "r") as f: + text = f.read() + else: + text = "" + super(SLConfig, self).__setattr__("_text", text) + + @property + def filename(self): + return self._filename + + @property + def text(self): + return self._text + + @property + def pretty_text(self): + + indent = 4 + + def _indent(s_, num_spaces): + s = s_.split("\n") + if len(s) == 1: + return s_ + first = s.pop(0) + s = [(num_spaces * " ") + line for line in s] + s = "\n".join(s) + s = first + "\n" + s + return s + + def _format_basic_types(k, v, use_mapping=False): + if isinstance(v, str): + v_str = f"'{v}'" + else: + v_str = str(v) + + if use_mapping: + k_str = f"'{k}'" if isinstance(k, str) else str(k) + attr_str = f"{k_str}: {v_str}" + else: + attr_str = f"{str(k)}={v_str}" + attr_str = _indent(attr_str, indent) + + return attr_str + + def _format_list(k, v, use_mapping=False): + # check if all items in the list are dict + if all(isinstance(_, dict) for _ in v): + v_str = "[\n" + v_str += "\n".join( + f"dict({_indent(_format_dict(v_), indent)})," for v_ in v + ).rstrip(",") + if use_mapping: + k_str = f"'{k}'" if isinstance(k, str) else str(k) + attr_str = f"{k_str}: {v_str}" + else: + attr_str = f"{str(k)}={v_str}" + attr_str = _indent(attr_str, indent) + "]" + else: + attr_str = _format_basic_types(k, v, use_mapping) + return attr_str + + def _contain_invalid_identifier(dict_str): + contain_invalid_identifier = False + for key_name in dict_str: + contain_invalid_identifier |= not str(key_name).isidentifier() + return contain_invalid_identifier + + def _format_dict(input_dict, outest_level=False): + r = "" + s = [] + + use_mapping = _contain_invalid_identifier(input_dict) + if use_mapping: + r += "{" + for idx, (k, v) in enumerate(input_dict.items()): + is_last = idx >= len(input_dict) - 1 + end = "" if outest_level or is_last else "," + if isinstance(v, dict): + v_str = "\n" + _format_dict(v) + if use_mapping: + k_str = f"'{k}'" if isinstance(k, str) else str(k) + attr_str = f"{k_str}: dict({v_str}" + else: + attr_str = f"{str(k)}=dict({v_str}" + attr_str = _indent(attr_str, indent) + ")" + end + elif isinstance(v, list): + attr_str = _format_list(k, v, use_mapping) + end + else: + attr_str = _format_basic_types(k, v, use_mapping) + end + + s.append(attr_str) + r += "\n".join(s) + if use_mapping: + r += "}" + return r + + cfg_dict = self._cfg_dict.to_dict() + text = _format_dict(cfg_dict, outest_level=True) + # copied from setup.cfg + yapf_style = dict( + based_on_style="pep8", + blank_line_before_nested_class_or_def=True, + split_before_expression_after_opening_paren=True, + ) + text, _ = FormatCode(text, style_config=yapf_style, verify=True) + + return text + + def __repr__(self): + return f"Config (path: {self.filename}): {self._cfg_dict.__repr__()}" + + def __len__(self): + return len(self._cfg_dict) + + def __getattr__(self, name): + # # debug + # print('+'*15) + # print('name=%s' % name) + # print("addr:", id(self)) + # # print('type(self):', type(self)) + # print(self.__dict__) + # print('+'*15) + # if self.__dict__ == {}: + # raise ValueError + + return getattr(self._cfg_dict, name) + + def __getitem__(self, name): + return self._cfg_dict.__getitem__(name) + + def __setattr__(self, name, value): + if isinstance(value, dict): + value = ConfigDict(value) + self._cfg_dict.__setattr__(name, value) + + def __setitem__(self, name, value): + if isinstance(value, dict): + value = ConfigDict(value) + self._cfg_dict.__setitem__(name, value) + + def __iter__(self): + return iter(self._cfg_dict) + + def dump(self, file=None): + # import ipdb; ipdb.set_trace() + if file is None: + return self.pretty_text + else: + with open(file, "w") as f: + f.write(self.pretty_text) + + def merge_from_dict(self, options): + """Merge list into cfg_dict + + Merge the dict parsed by MultipleKVAction into this cfg. + + Examples: + >>> options = {'model.backbone.depth': 50, + ... 'model.backbone.with_cp':True} + >>> cfg = Config(dict(model=dict(backbone=dict(type='ResNet')))) + >>> cfg.merge_from_dict(options) + >>> cfg_dict = super(Config, self).__getattribute__('_cfg_dict') + >>> assert cfg_dict == dict( + ... model=dict(backbone=dict(depth=50, with_cp=True))) + + Args: + options (dict): dict of configs to merge from. + """ + option_cfg_dict = {} + for full_key, v in options.items(): + d = option_cfg_dict + key_list = full_key.split(".") + for subkey in key_list[:-1]: + d.setdefault(subkey, ConfigDict()) + d = d[subkey] + subkey = key_list[-1] + d[subkey] = v + + cfg_dict = super(SLConfig, self).__getattribute__("_cfg_dict") + super(SLConfig, self).__setattr__( + "_cfg_dict", SLConfig._merge_a_into_b(option_cfg_dict, cfg_dict) + ) + + # for multiprocess + def __setstate__(self, state): + self.__init__(state) + + def copy(self): + return SLConfig(self._cfg_dict.copy()) + + def deepcopy(self): + return SLConfig(self._cfg_dict.deepcopy()) + + +class DictAction(Action): + """ + argparse action to split an argument into KEY=VALUE form + on the first = and append to a dictionary. List options should + be passed as comma separated values, i.e KEY=V1,V2,V3 + """ + + @staticmethod + def _parse_int_float_bool(val): + try: + return int(val) + except ValueError: + pass + try: + return float(val) + except ValueError: + pass + if val.lower() in ["true", "false"]: + return True if val.lower() == "true" else False + if val.lower() in ["none", "null"]: + return None + return val + + def __call__(self, parser, namespace, values, option_string=None): + options = {} + for kv in values: + key, val = kv.split("=", maxsplit=1) + val = [self._parse_int_float_bool(v) for v in val.split(",")] + if len(val) == 1: + val = val[0] + options[key] = val + setattr(namespace, self.dest, options) diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/slio.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/slio.py new file mode 100644 index 0000000000000000000000000000000000000000..72c1f0f7b82cdc931d381feef64fe15815ba657e --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/slio.py @@ -0,0 +1,177 @@ +# ========================================================== +# Modified from mmcv +# ========================================================== + +import json +import pickle +from abc import ABCMeta, abstractmethod +from pathlib import Path + +import yaml + +try: + from yaml import CLoader as Loader, CDumper as Dumper +except ImportError: + from yaml import Loader, Dumper + + +# =========================== +# Rigister handler +# =========================== + + +class BaseFileHandler(metaclass=ABCMeta): + @abstractmethod + def load_from_fileobj(self, file, **kwargs): + pass + + @abstractmethod + def dump_to_fileobj(self, obj, file, **kwargs): + pass + + @abstractmethod + def dump_to_str(self, obj, **kwargs): + pass + + def load_from_path(self, filepath, mode="r", **kwargs): + with open(filepath, mode) as f: + return self.load_from_fileobj(f, **kwargs) + + def dump_to_path(self, obj, filepath, mode="w", **kwargs): + with open(filepath, mode) as f: + self.dump_to_fileobj(obj, f, **kwargs) + + +class JsonHandler(BaseFileHandler): + def load_from_fileobj(self, file): + return json.load(file) + + def dump_to_fileobj(self, obj, file, **kwargs): + json.dump(obj, file, **kwargs) + + def dump_to_str(self, obj, **kwargs): + return json.dumps(obj, **kwargs) + + +class PickleHandler(BaseFileHandler): + def load_from_fileobj(self, file, **kwargs): + return pickle.load(file, **kwargs) + + def load_from_path(self, filepath, **kwargs): + return super(PickleHandler, self).load_from_path(filepath, mode="rb", **kwargs) + + def dump_to_str(self, obj, **kwargs): + kwargs.setdefault("protocol", 2) + return pickle.dumps(obj, **kwargs) + + def dump_to_fileobj(self, obj, file, **kwargs): + kwargs.setdefault("protocol", 2) + pickle.dump(obj, file, **kwargs) + + def dump_to_path(self, obj, filepath, **kwargs): + super(PickleHandler, self).dump_to_path(obj, filepath, mode="wb", **kwargs) + + +class YamlHandler(BaseFileHandler): + def load_from_fileobj(self, file, **kwargs): + kwargs.setdefault("Loader", Loader) + return yaml.load(file, **kwargs) + + def dump_to_fileobj(self, obj, file, **kwargs): + kwargs.setdefault("Dumper", Dumper) + yaml.dump(obj, file, **kwargs) + + def dump_to_str(self, obj, **kwargs): + kwargs.setdefault("Dumper", Dumper) + return yaml.dump(obj, **kwargs) + + +file_handlers = { + "json": JsonHandler(), + "yaml": YamlHandler(), + "yml": YamlHandler(), + "pickle": PickleHandler(), + "pkl": PickleHandler(), +} + +# =========================== +# load and dump +# =========================== + + +def is_str(x): + """Whether the input is an string instance. + + Note: This method is deprecated since python 2 is no longer supported. + """ + return isinstance(x, str) + + +def slload(file, file_format=None, **kwargs): + """Load data from json/yaml/pickle files. + + This method provides a unified api for loading data from serialized files. + + Args: + file (str or :obj:`Path` or file-like object): Filename or a file-like + object. + file_format (str, optional): If not specified, the file format will be + inferred from the file extension, otherwise use the specified one. + Currently supported formats include "json", "yaml/yml" and + "pickle/pkl". + + Returns: + The content from the file. + """ + if isinstance(file, Path): + file = str(file) + if file_format is None and is_str(file): + file_format = file.split(".")[-1] + if file_format not in file_handlers: + raise TypeError(f"Unsupported format: {file_format}") + + handler = file_handlers[file_format] + if is_str(file): + obj = handler.load_from_path(file, **kwargs) + elif hasattr(file, "read"): + obj = handler.load_from_fileobj(file, **kwargs) + else: + raise TypeError('"file" must be a filepath str or a file-object') + return obj + + +def sldump(obj, file=None, file_format=None, **kwargs): + """Dump data to json/yaml/pickle strings or files. + + This method provides a unified api for dumping data as strings or to files, + and also supports custom arguments for each file format. + + Args: + obj (any): The python object to be dumped. + file (str or :obj:`Path` or file-like object, optional): If not + specified, then the object is dump to a str, otherwise to a file + specified by the filename or file-like object. + file_format (str, optional): Same as :func:`load`. + + Returns: + bool: True for success, False otherwise. + """ + if isinstance(file, Path): + file = str(file) + if file_format is None: + if is_str(file): + file_format = file.split(".")[-1] + elif file is None: + raise ValueError("file_format must be specified since file is None") + if file_format not in file_handlers: + raise TypeError(f"Unsupported format: {file_format}") + + handler = file_handlers[file_format] + if file is None: + return handler.dump_to_str(obj, **kwargs) + elif is_str(file): + handler.dump_to_path(obj, file, **kwargs) + elif hasattr(file, "write"): + handler.dump_to_fileobj(obj, file, **kwargs) + else: + raise TypeError('"file" must be a filename str or a file-object') diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/time_counter.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/time_counter.py new file mode 100644 index 0000000000000000000000000000000000000000..0aedb2e4d61bfbe7571dca9d50053f0fedaa1359 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/time_counter.py @@ -0,0 +1,62 @@ +import json +import time + + +class TimeCounter: + def __init__(self) -> None: + pass + + def clear(self): + self.timedict = {} + self.basetime = time.perf_counter() + + def timeit(self, name): + nowtime = time.perf_counter() - self.basetime + self.timedict[name] = nowtime + self.basetime = time.perf_counter() + + +class TimeHolder: + def __init__(self) -> None: + self.timedict = {} + + def update(self, _timedict: dict): + for k, v in _timedict.items(): + if k not in self.timedict: + self.timedict[k] = AverageMeter(name=k, val_only=True) + self.timedict[k].update(val=v) + + def final_res(self): + return {k: v.avg for k, v in self.timedict.items()} + + def __str__(self): + return json.dumps(self.final_res(), indent=2) + + +class AverageMeter(object): + """Computes and stores the average and current value""" + + def __init__(self, name, fmt=":f", val_only=False): + self.name = name + self.fmt = fmt + self.val_only = val_only + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + def __str__(self): + if self.val_only: + fmtstr = "{name} {val" + self.fmt + "}" + else: + fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})" + return fmtstr.format(**self.__dict__) diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/utils.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e9f0318e306fa04bff0ada70486b41aaa69b07c8 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/utils.py @@ -0,0 +1,608 @@ +import argparse +import json +import warnings +from collections import OrderedDict +from copy import deepcopy +from typing import Any, Dict, List + +import numpy as np +import torch +from transformers import AutoTokenizer + +from groundingdino.util.slconfig import SLConfig + + +def slprint(x, name="x"): + if isinstance(x, (torch.Tensor, np.ndarray)): + print(f"{name}.shape:", x.shape) + elif isinstance(x, (tuple, list)): + print("type x:", type(x)) + for i in range(min(10, len(x))): + slprint(x[i], f"{name}[{i}]") + elif isinstance(x, dict): + for k, v in x.items(): + slprint(v, f"{name}[{k}]") + else: + print(f"{name}.type:", type(x)) + + +def clean_state_dict(state_dict): + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + if k[:7] == "module.": + k = k[7:] # remove `module.` + new_state_dict[k] = v + return new_state_dict + + +def renorm( + img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] +) -> torch.FloatTensor: + # img: tensor(3,H,W) or tensor(B,3,H,W) + # return: same as img + assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim() + if img.dim() == 3: + assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % ( + img.size(0), + str(img.size()), + ) + img_perm = img.permute(1, 2, 0) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(2, 0, 1) + else: # img.dim() == 4 + assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % ( + img.size(1), + str(img.size()), + ) + img_perm = img.permute(0, 2, 3, 1) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(0, 3, 1, 2) + + +class CocoClassMapper: + def __init__(self) -> None: + self.category_map_str = { + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "13": 12, + "14": 13, + "15": 14, + "16": 15, + "17": 16, + "18": 17, + "19": 18, + "20": 19, + "21": 20, + "22": 21, + "23": 22, + "24": 23, + "25": 24, + "27": 25, + "28": 26, + "31": 27, + "32": 28, + "33": 29, + "34": 30, + "35": 31, + "36": 32, + "37": 33, + "38": 34, + "39": 35, + "40": 36, + "41": 37, + "42": 38, + "43": 39, + "44": 40, + "46": 41, + "47": 42, + "48": 43, + "49": 44, + "50": 45, + "51": 46, + "52": 47, + "53": 48, + "54": 49, + "55": 50, + "56": 51, + "57": 52, + "58": 53, + "59": 54, + "60": 55, + "61": 56, + "62": 57, + "63": 58, + "64": 59, + "65": 60, + "67": 61, + "70": 62, + "72": 63, + "73": 64, + "74": 65, + "75": 66, + "76": 67, + "77": 68, + "78": 69, + "79": 70, + "80": 71, + "81": 72, + "82": 73, + "84": 74, + "85": 75, + "86": 76, + "87": 77, + "88": 78, + "89": 79, + "90": 80, + } + self.origin2compact_mapper = {int(k): v - 1 for k, v in self.category_map_str.items()} + self.compact2origin_mapper = {int(v - 1): int(k) for k, v in self.category_map_str.items()} + + def origin2compact(self, idx): + return self.origin2compact_mapper[int(idx)] + + def compact2origin(self, idx): + return self.compact2origin_mapper[int(idx)] + + +def to_device(item, device): + if isinstance(item, torch.Tensor): + return item.to(device) + elif isinstance(item, list): + return [to_device(i, device) for i in item] + elif isinstance(item, dict): + return {k: to_device(v, device) for k, v in item.items()} + else: + raise NotImplementedError( + "Call Shilong if you use other containers! type: {}".format(type(item)) + ) + + +# +def get_gaussian_mean(x, axis, other_axis, softmax=True): + """ + + Args: + x (float): Input images(BxCxHxW) + axis (int): The index for weighted mean + other_axis (int): The other index + + Returns: weighted index for axis, BxC + + """ + mat2line = torch.sum(x, axis=other_axis) + # mat2line = mat2line / mat2line.mean() * 10 + if softmax: + u = torch.softmax(mat2line, axis=2) + else: + u = mat2line / (mat2line.sum(2, keepdim=True) + 1e-6) + size = x.shape[axis] + ind = torch.linspace(0, 1, size).to(x.device) + batch = x.shape[0] + channel = x.shape[1] + index = ind.repeat([batch, channel, 1]) + mean_position = torch.sum(index * u, dim=2) + return mean_position + + +def get_expected_points_from_map(hm, softmax=True): + """get_gaussian_map_from_points + B,C,H,W -> B,N,2 float(0, 1) float(0, 1) + softargmax function + + Args: + hm (float): Input images(BxCxHxW) + + Returns: + weighted index for axis, BxCx2. float between 0 and 1. + + """ + # hm = 10*hm + B, C, H, W = hm.shape + y_mean = get_gaussian_mean(hm, 2, 3, softmax=softmax) # B,C + x_mean = get_gaussian_mean(hm, 3, 2, softmax=softmax) # B,C + # return torch.cat((x_mean.unsqueeze(-1), y_mean.unsqueeze(-1)), 2) + return torch.stack([x_mean, y_mean], dim=2) + + +# Positional encoding (section 5.1) +# borrow from nerf +class Embedder: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.create_embedding_fn() + + def create_embedding_fn(self): + embed_fns = [] + d = self.kwargs["input_dims"] + out_dim = 0 + if self.kwargs["include_input"]: + embed_fns.append(lambda x: x) + out_dim += d + + max_freq = self.kwargs["max_freq_log2"] + N_freqs = self.kwargs["num_freqs"] + + if self.kwargs["log_sampling"]: + freq_bands = 2.0 ** torch.linspace(0.0, max_freq, steps=N_freqs) + else: + freq_bands = torch.linspace(2.0**0.0, 2.0**max_freq, steps=N_freqs) + + for freq in freq_bands: + for p_fn in self.kwargs["periodic_fns"]: + embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq)) + out_dim += d + + self.embed_fns = embed_fns + self.out_dim = out_dim + + def embed(self, inputs): + return torch.cat([fn(inputs) for fn in self.embed_fns], -1) + + +def get_embedder(multires, i=0): + import torch.nn as nn + + if i == -1: + return nn.Identity(), 3 + + embed_kwargs = { + "include_input": True, + "input_dims": 3, + "max_freq_log2": multires - 1, + "num_freqs": multires, + "log_sampling": True, + "periodic_fns": [torch.sin, torch.cos], + } + + embedder_obj = Embedder(**embed_kwargs) + embed = lambda x, eo=embedder_obj: eo.embed(x) + return embed, embedder_obj.out_dim + + +class APOPMeter: + def __init__(self) -> None: + self.tp = 0 + self.fp = 0 + self.tn = 0 + self.fn = 0 + + def update(self, pred, gt): + """ + Input: + pred, gt: Tensor() + """ + assert pred.shape == gt.shape + self.tp += torch.logical_and(pred == 1, gt == 1).sum().item() + self.fp += torch.logical_and(pred == 1, gt == 0).sum().item() + self.tn += torch.logical_and(pred == 0, gt == 0).sum().item() + self.tn += torch.logical_and(pred == 1, gt == 0).sum().item() + + def update_cm(self, tp, fp, tn, fn): + self.tp += tp + self.fp += fp + self.tn += tn + self.tn += fn + + +def inverse_sigmoid(x, eps=1e-5): + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + + +def get_raw_dict(args): + """ + return the dicf contained in args. + + e.g: + >>> with open(path, 'w') as f: + json.dump(get_raw_dict(args), f, indent=2) + """ + if isinstance(args, argparse.Namespace): + return vars(args) + elif isinstance(args, dict): + return args + elif isinstance(args, SLConfig): + return args._cfg_dict + else: + raise NotImplementedError("Unknown type {}".format(type(args))) + + +def stat_tensors(tensor): + assert tensor.dim() == 1 + tensor_sm = tensor.softmax(0) + entropy = (tensor_sm * torch.log(tensor_sm + 1e-9)).sum() + + return { + "max": tensor.max(), + "min": tensor.min(), + "mean": tensor.mean(), + "var": tensor.var(), + "std": tensor.var() ** 0.5, + "entropy": entropy, + } + + +class NiceRepr: + """Inherit from this class and define ``__nice__`` to "nicely" print your + objects. + + Defines ``__str__`` and ``__repr__`` in terms of ``__nice__`` function + Classes that inherit from :class:`NiceRepr` should redefine ``__nice__``. + If the inheriting class has a ``__len__``, method then the default + ``__nice__`` method will return its length. + + Example: + >>> class Foo(NiceRepr): + ... def __nice__(self): + ... return 'info' + >>> foo = Foo() + >>> assert str(foo) == '' + >>> assert repr(foo).startswith('>> class Bar(NiceRepr): + ... pass + >>> bar = Bar() + >>> import pytest + >>> with pytest.warns(None) as record: + >>> assert 'object at' in str(bar) + >>> assert 'object at' in repr(bar) + + Example: + >>> class Baz(NiceRepr): + ... def __len__(self): + ... return 5 + >>> baz = Baz() + >>> assert str(baz) == '' + """ + + def __nice__(self): + """str: a "nice" summary string describing this module""" + if hasattr(self, "__len__"): + # It is a common pattern for objects to use __len__ in __nice__ + # As a convenience we define a default __nice__ for these objects + return str(len(self)) + else: + # In all other cases force the subclass to overload __nice__ + raise NotImplementedError(f"Define the __nice__ method for {self.__class__!r}") + + def __repr__(self): + """str: the string of the module""" + try: + nice = self.__nice__() + classname = self.__class__.__name__ + return f"<{classname}({nice}) at {hex(id(self))}>" + except NotImplementedError as ex: + warnings.warn(str(ex), category=RuntimeWarning) + return object.__repr__(self) + + def __str__(self): + """str: the string of the module""" + try: + classname = self.__class__.__name__ + nice = self.__nice__() + return f"<{classname}({nice})>" + except NotImplementedError as ex: + warnings.warn(str(ex), category=RuntimeWarning) + return object.__repr__(self) + + +def ensure_rng(rng=None): + """Coerces input into a random number generator. + + If the input is None, then a global random state is returned. + + If the input is a numeric value, then that is used as a seed to construct a + random state. Otherwise the input is returned as-is. + + Adapted from [1]_. + + Args: + rng (int | numpy.random.RandomState | None): + if None, then defaults to the global rng. Otherwise this can be an + integer or a RandomState class + Returns: + (numpy.random.RandomState) : rng - + a numpy random number generator + + References: + .. [1] https://gitlab.kitware.com/computer-vision/kwarray/blob/master/kwarray/util_random.py#L270 # noqa: E501 + """ + + if rng is None: + rng = np.random.mtrand._rand + elif isinstance(rng, int): + rng = np.random.RandomState(rng) + else: + rng = rng + return rng + + +def random_boxes(num=1, scale=1, rng=None): + """Simple version of ``kwimage.Boxes.random`` + + Returns: + Tensor: shape (n, 4) in x1, y1, x2, y2 format. + + References: + https://gitlab.kitware.com/computer-vision/kwimage/blob/master/kwimage/structs/boxes.py#L1390 + + Example: + >>> num = 3 + >>> scale = 512 + >>> rng = 0 + >>> boxes = random_boxes(num, scale, rng) + >>> print(boxes) + tensor([[280.9925, 278.9802, 308.6148, 366.1769], + [216.9113, 330.6978, 224.0446, 456.5878], + [405.3632, 196.3221, 493.3953, 270.7942]]) + """ + rng = ensure_rng(rng) + + tlbr = rng.rand(num, 4).astype(np.float32) + + tl_x = np.minimum(tlbr[:, 0], tlbr[:, 2]) + tl_y = np.minimum(tlbr[:, 1], tlbr[:, 3]) + br_x = np.maximum(tlbr[:, 0], tlbr[:, 2]) + br_y = np.maximum(tlbr[:, 1], tlbr[:, 3]) + + tlbr[:, 0] = tl_x * scale + tlbr[:, 1] = tl_y * scale + tlbr[:, 2] = br_x * scale + tlbr[:, 3] = br_y * scale + + boxes = torch.from_numpy(tlbr) + return boxes + + +class ModelEma(torch.nn.Module): + def __init__(self, model, decay=0.9997, device=None): + super(ModelEma, self).__init__() + # make a copy of the model for accumulating moving average of weights + self.module = deepcopy(model) + self.module.eval() + + # import ipdb; ipdb.set_trace() + + self.decay = decay + self.device = device # perform ema on different device from model if set + if self.device is not None: + self.module.to(device=device) + + def _update(self, model, update_fn): + with torch.no_grad(): + for ema_v, model_v in zip( + self.module.state_dict().values(), model.state_dict().values() + ): + if self.device is not None: + model_v = model_v.to(device=self.device) + ema_v.copy_(update_fn(ema_v, model_v)) + + def update(self, model): + self._update(model, update_fn=lambda e, m: self.decay * e + (1.0 - self.decay) * m) + + def set(self, model): + self._update(model, update_fn=lambda e, m: m) + + +class BestMetricSingle: + def __init__(self, init_res=0.0, better="large") -> None: + self.init_res = init_res + self.best_res = init_res + self.best_ep = -1 + + self.better = better + assert better in ["large", "small"] + + def isbetter(self, new_res, old_res): + if self.better == "large": + return new_res > old_res + if self.better == "small": + return new_res < old_res + + def update(self, new_res, ep): + if self.isbetter(new_res, self.best_res): + self.best_res = new_res + self.best_ep = ep + return True + return False + + def __str__(self) -> str: + return "best_res: {}\t best_ep: {}".format(self.best_res, self.best_ep) + + def __repr__(self) -> str: + return self.__str__() + + def summary(self) -> dict: + return { + "best_res": self.best_res, + "best_ep": self.best_ep, + } + + +class BestMetricHolder: + def __init__(self, init_res=0.0, better="large", use_ema=False) -> None: + self.best_all = BestMetricSingle(init_res, better) + self.use_ema = use_ema + if use_ema: + self.best_ema = BestMetricSingle(init_res, better) + self.best_regular = BestMetricSingle(init_res, better) + + def update(self, new_res, epoch, is_ema=False): + """ + return if the results is the best. + """ + if not self.use_ema: + return self.best_all.update(new_res, epoch) + else: + if is_ema: + self.best_ema.update(new_res, epoch) + return self.best_all.update(new_res, epoch) + else: + self.best_regular.update(new_res, epoch) + return self.best_all.update(new_res, epoch) + + def summary(self): + if not self.use_ema: + return self.best_all.summary() + + res = {} + res.update({f"all_{k}": v for k, v in self.best_all.summary().items()}) + res.update({f"regular_{k}": v for k, v in self.best_regular.summary().items()}) + res.update({f"ema_{k}": v for k, v in self.best_ema.summary().items()}) + return res + + def __repr__(self) -> str: + return json.dumps(self.summary(), indent=2) + + def __str__(self) -> str: + return self.__repr__() + + +def targets_to(targets: List[Dict[str, Any]], device): + """Moves the target dicts to the given device.""" + excluded_keys = [ + "questionId", + "tokens_positive", + "strings_positive", + "tokens", + "dataset_name", + "sentence_id", + "original_img_id", + "nb_eval", + "task_id", + "original_id", + "token_span", + "caption", + "dataset_type", + ] + return [ + {k: v.to(device) if k not in excluded_keys else v for k, v in t.items()} for t in targets + ] + + +def get_phrases_from_posmap( + posmap: torch.BoolTensor, tokenized: Dict, tokenizer: AutoTokenizer +): + assert isinstance(posmap, torch.Tensor), "posmap must be torch.Tensor" + if posmap.dim() == 1: + non_zero_idx = posmap.nonzero(as_tuple=True)[0].tolist() + token_ids = [tokenized["input_ids"][i] for i in non_zero_idx] + return tokenizer.decode(token_ids) + else: + raise NotImplementedError("posmap must be 1-dim") diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/visualizer.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7a1b7b101e9b73f75f9136bc67f2063c7c1cf1c1 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/visualizer.py @@ -0,0 +1,318 @@ +# -*- coding: utf-8 -*- +""" +@File : visualizer.py +@Time : 2022/04/05 11:39:33 +@Author : Shilong Liu +@Contact : slongliu86@gmail.com +""" + +import datetime +import os + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import torch +from matplotlib import transforms +from matplotlib.collections import PatchCollection +from matplotlib.patches import Polygon +from pycocotools import mask as maskUtils + + +def renorm( + img: torch.FloatTensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] +) -> torch.FloatTensor: + # img: tensor(3,H,W) or tensor(B,3,H,W) + # return: same as img + assert img.dim() == 3 or img.dim() == 4, "img.dim() should be 3 or 4 but %d" % img.dim() + if img.dim() == 3: + assert img.size(0) == 3, 'img.size(0) shoule be 3 but "%d". (%s)' % ( + img.size(0), + str(img.size()), + ) + img_perm = img.permute(1, 2, 0) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(2, 0, 1) + else: # img.dim() == 4 + assert img.size(1) == 3, 'img.size(1) shoule be 3 but "%d". (%s)' % ( + img.size(1), + str(img.size()), + ) + img_perm = img.permute(0, 2, 3, 1) + mean = torch.Tensor(mean) + std = torch.Tensor(std) + img_res = img_perm * std + mean + return img_res.permute(0, 3, 1, 2) + + +class ColorMap: + def __init__(self, basergb=[255, 255, 0]): + self.basergb = np.array(basergb) + + def __call__(self, attnmap): + # attnmap: h, w. np.uint8. + # return: h, w, 4. np.uint8. + assert attnmap.dtype == np.uint8 + h, w = attnmap.shape + res = self.basergb.copy() + res = res[None][None].repeat(h, 0).repeat(w, 1) # h, w, 3 + attn1 = attnmap.copy()[..., None] # h, w, 1 + res = np.concatenate((res, attn1), axis=-1).astype(np.uint8) + return res + + +def rainbow_text(x, y, ls, lc, **kw): + """ + Take a list of strings ``ls`` and colors ``lc`` and place them next to each + other, with text ls[i] being shown in color lc[i]. + + This example shows how to do both vertical and horizontal text, and will + pass all keyword arguments to plt.text, so you can set the font size, + family, etc. + """ + t = plt.gca().transData + fig = plt.gcf() + plt.show() + + # horizontal version + for s, c in zip(ls, lc): + text = plt.text(x, y, " " + s + " ", color=c, transform=t, **kw) + text.draw(fig.canvas.get_renderer()) + ex = text.get_window_extent() + t = transforms.offset_copy(text._transform, x=ex.width, units="dots") + + # #vertical version + # for s,c in zip(ls,lc): + # text = plt.text(x,y," "+s+" ",color=c, transform=t, + # rotation=90,va='bottom',ha='center',**kw) + # text.draw(fig.canvas.get_renderer()) + # ex = text.get_window_extent() + # t = transforms.offset_copy(text._transform, y=ex.height, units='dots') + + +class COCOVisualizer: + def __init__(self, coco=None, tokenlizer=None) -> None: + self.coco = coco + + def visualize(self, img, tgt, caption=None, dpi=180, savedir="vis"): + """ + img: tensor(3, H, W) + tgt: make sure they are all on cpu. + must have items: 'image_id', 'boxes', 'size' + """ + plt.figure(dpi=dpi) + plt.rcParams["font.size"] = "5" + ax = plt.gca() + img = renorm(img).permute(1, 2, 0) + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + ax.imshow(img) + + self.addtgt(tgt) + + if tgt is None: + image_id = 0 + elif "image_id" not in tgt: + image_id = 0 + else: + image_id = tgt["image_id"] + + if caption is None: + savename = "{}/{}-{}.png".format( + savedir, int(image_id), str(datetime.datetime.now()).replace(" ", "-") + ) + else: + savename = "{}/{}-{}-{}.png".format( + savedir, caption, int(image_id), str(datetime.datetime.now()).replace(" ", "-") + ) + print("savename: {}".format(savename)) + os.makedirs(os.path.dirname(savename), exist_ok=True) + plt.savefig(savename) + plt.close() + + def addtgt(self, tgt): + """ """ + if tgt is None or not "boxes" in tgt: + ax = plt.gca() + + if "caption" in tgt: + ax.set_title(tgt["caption"], wrap=True) + + ax.set_axis_off() + return + + ax = plt.gca() + H, W = tgt["size"] + numbox = tgt["boxes"].shape[0] + + color = [] + polygons = [] + boxes = [] + for box in tgt["boxes"].cpu(): + unnormbbox = box * torch.Tensor([W, H, W, H]) + unnormbbox[:2] -= unnormbbox[2:] / 2 + [bbox_x, bbox_y, bbox_w, bbox_h] = unnormbbox.tolist() + boxes.append([bbox_x, bbox_y, bbox_w, bbox_h]) + poly = [ + [bbox_x, bbox_y], + [bbox_x, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y], + ] + np_poly = np.array(poly).reshape((4, 2)) + polygons.append(Polygon(np_poly)) + c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] + color.append(c) + + p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.1) + ax.add_collection(p) + p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2) + ax.add_collection(p) + + if "strings_positive" in tgt and len(tgt["strings_positive"]) > 0: + assert ( + len(tgt["strings_positive"]) == numbox + ), f"{len(tgt['strings_positive'])} = {numbox}, " + for idx, strlist in enumerate(tgt["strings_positive"]): + cate_id = int(tgt["labels"][idx]) + _string = str(cate_id) + ":" + " ".join(strlist) + bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx] + # ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1}) + ax.text( + bbox_x, + bbox_y, + _string, + color="black", + bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1}, + ) + + if "box_label" in tgt: + assert len(tgt["box_label"]) == numbox, f"{len(tgt['box_label'])} = {numbox}, " + for idx, bl in enumerate(tgt["box_label"]): + _string = str(bl) + bbox_x, bbox_y, bbox_w, bbox_h = boxes[idx] + # ax.text(bbox_x, bbox_y, _string, color='black', bbox={'facecolor': 'yellow', 'alpha': 1.0, 'pad': 1}) + ax.text( + bbox_x, + bbox_y, + _string, + color="black", + bbox={"facecolor": color[idx], "alpha": 0.6, "pad": 1}, + ) + + if "caption" in tgt: + ax.set_title(tgt["caption"], wrap=True) + # plt.figure() + # rainbow_text(0.0,0.0,"all unicorns poop rainbows ! ! !".split(), + # ['red', 'orange', 'brown', 'green', 'blue', 'purple', 'black']) + + if "attn" in tgt: + # if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO': + # import ipdb; ipdb.set_trace() + if isinstance(tgt["attn"], tuple): + tgt["attn"] = [tgt["attn"]] + for item in tgt["attn"]: + attn_map, basergb = item + attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-3) + attn_map = (attn_map * 255).astype(np.uint8) + cm = ColorMap(basergb) + heatmap = cm(attn_map) + ax.imshow(heatmap) + ax.set_axis_off() + + def showAnns(self, anns, draw_bbox=False): + """ + Display the specified annotations. + :param anns (array of object): annotations to display + :return: None + """ + if len(anns) == 0: + return 0 + if "segmentation" in anns[0] or "keypoints" in anns[0]: + datasetType = "instances" + elif "caption" in anns[0]: + datasetType = "captions" + else: + raise Exception("datasetType not supported") + if datasetType == "instances": + ax = plt.gca() + ax.set_autoscale_on(False) + polygons = [] + color = [] + for ann in anns: + c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0] + if "segmentation" in ann: + if type(ann["segmentation"]) == list: + # polygon + for seg in ann["segmentation"]: + poly = np.array(seg).reshape((int(len(seg) / 2), 2)) + polygons.append(Polygon(poly)) + color.append(c) + else: + # mask + t = self.imgs[ann["image_id"]] + if type(ann["segmentation"]["counts"]) == list: + rle = maskUtils.frPyObjects( + [ann["segmentation"]], t["height"], t["width"] + ) + else: + rle = [ann["segmentation"]] + m = maskUtils.decode(rle) + img = np.ones((m.shape[0], m.shape[1], 3)) + if ann["iscrowd"] == 1: + color_mask = np.array([2.0, 166.0, 101.0]) / 255 + if ann["iscrowd"] == 0: + color_mask = np.random.random((1, 3)).tolist()[0] + for i in range(3): + img[:, :, i] = color_mask[i] + ax.imshow(np.dstack((img, m * 0.5))) + if "keypoints" in ann and type(ann["keypoints"]) == list: + # turn skeleton into zero-based index + sks = np.array(self.loadCats(ann["category_id"])[0]["skeleton"]) - 1 + kp = np.array(ann["keypoints"]) + x = kp[0::3] + y = kp[1::3] + v = kp[2::3] + for sk in sks: + if np.all(v[sk] > 0): + plt.plot(x[sk], y[sk], linewidth=3, color=c) + plt.plot( + x[v > 0], + y[v > 0], + "o", + markersize=8, + markerfacecolor=c, + markeredgecolor="k", + markeredgewidth=2, + ) + plt.plot( + x[v > 1], + y[v > 1], + "o", + markersize=8, + markerfacecolor=c, + markeredgecolor=c, + markeredgewidth=2, + ) + + if draw_bbox: + [bbox_x, bbox_y, bbox_w, bbox_h] = ann["bbox"] + poly = [ + [bbox_x, bbox_y], + [bbox_x, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y + bbox_h], + [bbox_x + bbox_w, bbox_y], + ] + np_poly = np.array(poly).reshape((4, 2)) + polygons.append(Polygon(np_poly)) + color.append(c) + + # p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4) + # ax.add_collection(p) + p = PatchCollection(polygons, facecolor="none", edgecolors=color, linewidths=2) + ax.add_collection(p) + elif datasetType == "captions": + for ann in anns: + print(ann["caption"]) diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/vl_utils.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/vl_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c91bb02f584398f08a28e6b7719e2b99f6e28616 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/util/vl_utils.py @@ -0,0 +1,100 @@ +import os +import random +from typing import List + +import torch + + +def create_positive_map_from_span(tokenized, token_span, max_text_len=256): + """construct a map such that positive_map[i,j] = True iff box i is associated to token j + Input: + - tokenized: + - input_ids: Tensor[1, ntokens] + - attention_mask: Tensor[1, ntokens] + - token_span: list with length num_boxes. + - each item: [start_idx, end_idx] + """ + positive_map = torch.zeros((len(token_span), max_text_len), dtype=torch.float) + for j, tok_list in enumerate(token_span): + for (beg, end) in tok_list: + beg_pos = tokenized.char_to_token(beg) + end_pos = tokenized.char_to_token(end - 1) + if beg_pos is None: + try: + beg_pos = tokenized.char_to_token(beg + 1) + if beg_pos is None: + beg_pos = tokenized.char_to_token(beg + 2) + except: + beg_pos = None + if end_pos is None: + try: + end_pos = tokenized.char_to_token(end - 2) + if end_pos is None: + end_pos = tokenized.char_to_token(end - 3) + except: + end_pos = None + if beg_pos is None or end_pos is None: + continue + + assert beg_pos is not None and end_pos is not None + if os.environ.get("SHILONG_DEBUG_ONLY_ONE_POS", None) == "TRUE": + positive_map[j, beg_pos] = 1 + break + else: + positive_map[j, beg_pos : end_pos + 1].fill_(1) + + return positive_map / (positive_map.sum(-1)[:, None] + 1e-6) + + +def build_captions_and_token_span(cat_list, force_lowercase): + """ + Return: + captions: str + cat2tokenspan: dict + { + 'dog': [[0, 2]], + ... + } + """ + + cat2tokenspan = {} + captions = "" + for catname in cat_list: + class_name = catname + if force_lowercase: + class_name = class_name.lower() + if "/" in class_name: + class_name_list: List = class_name.strip().split("/") + class_name_list.append(class_name) + class_name: str = random.choice(class_name_list) + + tokens_positive_i = [] + subnamelist = [i.strip() for i in class_name.strip().split(" ")] + for subname in subnamelist: + if len(subname) == 0: + continue + if len(captions) > 0: + captions = captions + " " + strat_idx = len(captions) + end_idx = strat_idx + len(subname) + tokens_positive_i.append([strat_idx, end_idx]) + captions = captions + subname + + if len(tokens_positive_i) > 0: + captions = captions + " ." + cat2tokenspan[class_name] = tokens_positive_i + + return captions, cat2tokenspan + + +def build_id2posspan_and_caption(category_dict: dict): + """Build id2pos_span and caption from category_dict + + Args: + category_dict (dict): category_dict + """ + cat_list = [item["name"].lower() for item in category_dict] + id2catname = {item["id"]: item["name"].lower() for item in category_dict} + caption, cat2posspan = build_captions_and_token_span(cat_list, force_lowercase=True) + id2posspan = {catid: cat2posspan[catname] for catid, catname in id2catname.items()} + return id2posspan, caption diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/version.py b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/version.py new file mode 100644 index 0000000000000000000000000000000000000000..b794fd409a5e3b3b65ad76a43d6a01a318877640 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/groundingdino/version.py @@ -0,0 +1 @@ +__version__ = '0.1.0' diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/pyproject.toml b/ArtiAgent - DefectFill/src/GroundingDINO/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..24dcc68d94ea5aaee6bb7a903a0e1638cf14e6b1 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/pyproject.toml @@ -0,0 +1,8 @@ +[build-system] +requires = [ + "setuptools", + "torch", + "wheel", + "torch" +] +build-backend = "setuptools.build_meta" diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/requirements.txt b/ArtiAgent - DefectFill/src/GroundingDINO/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..32e2538a959167f9ce248a5c99cf275bc53cc51b --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/requirements.txt @@ -0,0 +1,12 @@ +torch +torchvision +transformers +addict +yapf +timm +numpy +opencv-python +supervision==0.21.0 +pycocotools +lpips +openai \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/GroundingDINO/setup.py b/ArtiAgent - DefectFill/src/GroundingDINO/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..a58340d44eca86b09cb69630465dfbdfe8acb742 --- /dev/null +++ b/ArtiAgent - DefectFill/src/GroundingDINO/setup.py @@ -0,0 +1,216 @@ +# coding=utf-8 +# Copyright 2022 The IDEA Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ------------------------------------------------------------------------------------------------ +# Modified from +# https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/setup.py +# https://github.com/facebookresearch/detectron2/blob/main/setup.py +# https://github.com/open-mmlab/mmdetection/blob/master/setup.py +# https://github.com/Oneflow-Inc/libai/blob/main/setup.py +# ------------------------------------------------------------------------------------------------ + +import glob +import os +import subprocess + +import torch +from setuptools import find_packages, setup +from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension + +# groundingdino version info +version = "0.1.0" +package_name = "groundingdino" +cwd = os.path.dirname(os.path.abspath(__file__)) + + +sha = "Unknown" +try: + sha = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=cwd).decode("ascii").strip() +except Exception: + pass + + +def write_version_file(): + version_path = os.path.join(cwd, "groundingdino", "version.py") + with open(version_path, "w") as f: + f.write(f"__version__ = '{version}'\n") + # f.write(f"git_version = {repr(sha)}\n") + + +requirements = ["torch", "torchvision"] + +torch_ver = [int(x) for x in torch.__version__.split(".")[:2]] + + +def get_extensions(): + this_dir = os.path.dirname(os.path.abspath(__file__)) + extensions_dir = os.path.join(this_dir, "groundingdino", "models", "GroundingDINO", "csrc") + + main_source = os.path.join(extensions_dir, "vision.cpp") + sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp")) + source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu")) + glob.glob( + os.path.join(extensions_dir, "*.cu") + ) + + sources = [main_source] + sources + + # We need these variables to build with CUDA when we create the Docker image + # It solves https://github.com/IDEA-Research/Grounded-Segment-Anything/issues/53 + # and https://github.com/IDEA-Research/Grounded-Segment-Anything/issues/84 when running + # inside a Docker container. + am_i_docker = os.environ.get('AM_I_DOCKER', '').casefold() in ['true', '1', 't'] + use_cuda = os.environ.get('BUILD_WITH_CUDA', '').casefold() in ['true', '1', 't'] + + extension = CppExtension + + extra_compile_args = {"cxx": []} + define_macros = [] + + if (torch.cuda.is_available() and CUDA_HOME is not None) or \ + (am_i_docker and use_cuda): + print("Compiling with CUDA") + extension = CUDAExtension + sources += source_cuda + define_macros += [("WITH_CUDA", None)] + extra_compile_args["nvcc"] = [ + "-DCUDA_HAS_FP16=1", + "-D__CUDA_NO_HALF_OPERATORS__", + "-D__CUDA_NO_HALF_CONVERSIONS__", + "-D__CUDA_NO_HALF2_OPERATORS__", + ] + else: + print("Compiling without CUDA") + define_macros += [("WITH_HIP", None)] + extra_compile_args["nvcc"] = [] + return None + + sources = [os.path.join(extensions_dir, s) for s in sources] + include_dirs = [extensions_dir] + + ext_modules = [ + extension( + "groundingdino._C", + sources, + include_dirs=include_dirs, + define_macros=define_macros, + extra_compile_args=extra_compile_args, + ) + ] + + return ext_modules + + +def parse_requirements(fname="requirements.txt", with_version=True): + """Parse the package dependencies listed in a requirements file but strips + specific versioning information. + + Args: + fname (str): path to requirements file + with_version (bool, default=False): if True include version specs + + Returns: + List[str]: list of requirements items + + CommandLine: + python -c "import setup; print(setup.parse_requirements())" + """ + import re + import sys + from os.path import exists + + require_fpath = fname + + def parse_line(line): + """Parse information from a line in a requirements text file.""" + if line.startswith("-r "): + # Allow specifying requirements in other files + target = line.split(" ")[1] + for info in parse_require_file(target): + yield info + else: + info = {"line": line} + if line.startswith("-e "): + info["package"] = line.split("#egg=")[1] + elif "@git+" in line: + info["package"] = line + else: + # Remove versioning from the package + pat = "(" + "|".join([">=", "==", ">"]) + ")" + parts = re.split(pat, line, maxsplit=1) + parts = [p.strip() for p in parts] + + info["package"] = parts[0] + if len(parts) > 1: + op, rest = parts[1:] + if ";" in rest: + # Handle platform specific dependencies + # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies + version, platform_deps = map(str.strip, rest.split(";")) + info["platform_deps"] = platform_deps + else: + version = rest # NOQA + info["version"] = (op, version) + yield info + + def parse_require_file(fpath): + with open(fpath, "r") as f: + for line in f.readlines(): + line = line.strip() + if line and not line.startswith("#"): + for info in parse_line(line): + yield info + + def gen_packages_items(): + if exists(require_fpath): + for info in parse_require_file(require_fpath): + parts = [info["package"]] + if with_version and "version" in info: + parts.extend(info["version"]) + if not sys.version.startswith("3.4"): + # apparently package_deps are broken in 3.4 + platform_deps = info.get("platform_deps") + if platform_deps is not None: + parts.append(";" + platform_deps) + item = "".join(parts) + yield item + + packages = list(gen_packages_items()) + return packages + + +if __name__ == "__main__": + print(f"Building wheel {package_name}-{version}") + + with open("LICENSE", "r", encoding="utf-8") as f: + license = f.read() + + write_version_file() + + setup( + name="groundingdino", + version="0.1.0", + author="International Digital Economy Academy, Shilong Liu", + url="https://github.com/IDEA-Research/GroundingDINO", + description="open-set object detector", + license=license, + install_requires=parse_requirements("requirements.txt"), + packages=find_packages( + exclude=( + "configs", + "tests", + ) + ), + ext_modules=get_extensions(), + cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, + ) diff --git a/ArtiAgent - DefectFill/src/__pycache__/align_image_to_reference.cpython-310.pyc b/ArtiAgent - DefectFill/src/__pycache__/align_image_to_reference.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e1b7288222712f329a5c679b39945c411563d2f Binary files /dev/null and b/ArtiAgent - DefectFill/src/__pycache__/align_image_to_reference.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/__pycache__/artiagent_orchestrator.cpython-310.pyc b/ArtiAgent - DefectFill/src/__pycache__/artiagent_orchestrator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ae414099555ce963afd5d7b7c8d03a22965d1f7 Binary files /dev/null and b/ArtiAgent - DefectFill/src/__pycache__/artiagent_orchestrator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/__pycache__/detect_similar_product.cpython-310.pyc b/ArtiAgent - DefectFill/src/__pycache__/detect_similar_product.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebc0ce4d40cd8495b84d770d01272461bb562fae Binary files /dev/null and b/ArtiAgent - DefectFill/src/__pycache__/detect_similar_product.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/__pycache__/model.cpython-310.pyc b/ArtiAgent - DefectFill/src/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e803b8e159cf3a2c71c1324db73c562c37dcde7 Binary files /dev/null and b/ArtiAgent - DefectFill/src/__pycache__/model.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/__pycache__/utils.cpython-310.pyc b/ArtiAgent - DefectFill/src/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b022605e8fe0209d9d65e0110c5e4b5ef63556e Binary files /dev/null and b/ArtiAgent - DefectFill/src/__pycache__/utils.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/align_image_to_reference - Copy.py b/ArtiAgent - DefectFill/src/align_image_to_reference - Copy.py new file mode 100644 index 0000000000000000000000000000000000000000..7ae90daf52189b0618bae30b4faabc1af6e12728 --- /dev/null +++ b/ArtiAgent - DefectFill/src/align_image_to_reference - Copy.py @@ -0,0 +1,137 @@ +import os +import cv2 +import numpy as np +from flask import Flask, request, jsonify, send_file +from werkzeug.utils import secure_filename + +app = Flask(__name__) + +# ============================================================================== +# [CONFIG: STORAGE DIRECTORIES] +# ------------------------------------------------------------------------------ +# Modify these paths to change where incoming raw uploads and processed output +# images are stored on your server's disk. +# ============================================================================== +UPLOAD_FOLDER = os.path.join(os.getcwd(), "storage", "inputs") # <--- INPUT DIRECTORY +OUTPUT_FOLDER = os.path.join(os.getcwd(), "storage", "outputs") # <--- OUTPUT DIRECTORY + +app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER +app.config["OUTPUT_FOLDER"] = OUTPUT_FOLDER + +# Ensure local directories exist on startup +os.makedirs(UPLOAD_FOLDER, exist_ok=True) +os.makedirs(OUTPUT_FOLDER, exist_ok=True) + + +def align_image_to_reference(ref_path: str, target_path: str, output_path: str): + """Aligns target_path image to match ref_path image geometry and saves to output_path.""" + ref_img = cv2.imread(ref_path) + targ_img = cv2.imread(target_path) + + if ref_img is None or targ_img is None: + raise ValueError("Could not read input images from storage.") + + # Convert to grayscale for SIFT feature extraction + gray_ref = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY) + gray_targ = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY) + + # 1. Detect SIFT features + sift = cv2.SIFT_create() + kp_ref, des_ref = sift.detectAndCompute(gray_ref, None) + kp_targ, des_targ = sift.detectAndCompute(gray_targ, None) + + if des_ref is None or des_targ is None: + raise ValueError("Failed to extract keypoints from one or both images.") + + # 2. Match keypoints using FLANN + INDEX_KDTREE = 1 + flann = cv2.FlannBasedMatcher( + dict(algorithm=INDEX_KDTREE, trees=5), dict(checks=50) + ) + matches = flann.knnMatch(des_ref, des_targ, k=2) + + # 3. Apply Lowe's ratio test to filter matches + good_matches = [m for m, n in matches if m.distance < 0.7 * n.distance] + + if len(good_matches) < 10: + raise ValueError("Insufficient matching features found between images.") + + # 4. Extract keypoint coordinates + src_pts = np.float32([kp_ref[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2) + dst_pts = np.float32([kp_targ[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2) + + # 5. Compute Affine Transformation Matrix (rigid: rotation, scale, translation) + matrix, _ = cv2.estimateAffinePartial2D(dst_pts, src_pts, method=cv2.RANSAC) + + if matrix is None: + raise ValueError("Failed to compute valid alignment transformation matrix.") + + # 6. Warp target image to match reference frame dimensions + h, w = ref_img.shape[:2] + aligned_img = cv2.warpAffine( + targ_img, + matrix, + (w, h), + flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_CONSTANT, + borderValue=(0, 0, 0) + ) + + # ============================================================================== + # [OUTPUT IMAGE STORAGE LOCATION - WRITE TO DISK] + # ------------------------------------------------------------------------------ + # The rotated/aligned image is saved here to output_path + # ============================================================================== + cv2.imwrite(output_path, aligned_img) + + +@app.route("/align", methods=["POST"]) +def align_endpoint(): + # Validate request payload + if "reference" not in request.files or "target" not in request.files: + return jsonify({"error": "Missing 'reference' or 'target' file in request form-data."}), 400 + + ref_file = request.files["reference"] + target_file = request.files["target"] + + if ref_file.filename == "" or target_file.filename == "": + return jsonify({"error": "No file selected."}), 400 + + # Sanitize filenames + ref_name = secure_filename(ref_file.filename) + target_name = secure_filename(target_file.filename) + + # ============================================================================== + # [INPUT IMAGE STORAGE LOCATION - SAVE RECEIVED FILES] + # ------------------------------------------------------------------------------ + # Input files are saved into 'app.config["UPLOAD_FOLDER"]' + # ============================================================================== + input_ref_path = os.path.join(app.config["UPLOAD_FOLDER"], f"ref_{ref_name}") + input_target_path = os.path.join(app.config["UPLOAD_FOLDER"], f"target_{target_name}") + + ref_file.save(input_ref_path) # <-- Input Reference saved here + target_file.save(input_target_path) # <-- Input Target saved here + # ============================================================================== + + # ============================================================================== + # [OUTPUT IMAGE STORAGE LOCATION - DEFINE TARGET PATH] + # ------------------------------------------------------------------------------ + # Rotated image path in 'app.config["OUTPUT_FOLDER"]' + # ============================================================================== + output_aligned_path = os.path.join(app.config["OUTPUT_FOLDER"], f"aligned_{target_name}") + # ============================================================================== + + try: + # Run alignment pipeline + align_image_to_reference(input_ref_path, input_target_path, output_aligned_path) + + # Return the processed image directly in the response + return send_file(output_aligned_path, mimetype="image/png") + + except Exception as e: + return jsonify({"status": "error", "message": str(e)}), 500 + + +if __name__ == "__main__": + # Run API server on http://0.0.0.0:5000 + app.run(host="0.0.0.0", port=5000, debug=True) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/align_image_to_reference.py b/ArtiAgent - DefectFill/src/align_image_to_reference.py new file mode 100644 index 0000000000000000000000000000000000000000..5950ca55904c49336fb0746473d55c77fe44bbe0 --- /dev/null +++ b/ArtiAgent - DefectFill/src/align_image_to_reference.py @@ -0,0 +1,372 @@ +import os +import cv2 +import numpy as np +from flask import Flask, request, jsonify, send_file +from werkzeug.utils import secure_filename + +app = Flask(__name__) + +UPLOAD_FOLDER = os.path.join(os.getcwd(), "storage", "inputs") +OUTPUT_FOLDER = os.path.join(os.getcwd(), "storage", "outputs") + +app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER +app.config["OUTPUT_FOLDER"] = OUTPUT_FOLDER + +os.makedirs(UPLOAD_FOLDER, exist_ok=True) +os.makedirs(OUTPUT_FOLDER, exist_ok=True) + + +def align_image_to_reference(ref_path: str, target_path: str, output_path: str): + """Aligns target_path image to match ref_path image geometry and saves to output_path. + + Uses multiple fallback strategies for robustness: + 1. SIFT with scale normalization and contrast enhancement + 2. ORB (good for low-texture / small images) + 3. AKAZE + 4. Phase correlation (frequency domain, good when features fail) + 5. Image moments (center-of-mass + principal axis) + 6. Ultimate fallback: center-crop + resize + """ + ref_img = cv2.imread(ref_path) + targ_img = cv2.imread(target_path) + + if ref_img is None or targ_img is None: + raise ValueError("Could not read input images from storage.") + + h, w = ref_img.shape[:2] + + # Try strategies in order of preference + aligned = None + method_used = "unknown" + + # --- Strategy 1: SIFT with scale normalization --- + try: + aligned = _align_sift(ref_img, targ_img, w, h) + if aligned is not None: + method_used = "sift" + except Exception: + pass + + # --- Strategy 2: ORB (better for small/blurry images) --- + if aligned is None: + try: + aligned = _align_orb(ref_img, targ_img, w, h) + if aligned is not None: + method_used = "orb" + except Exception: + pass + + # --- Strategy 3: AKAZE --- + if aligned is None: + try: + aligned = _align_akaze(ref_img, targ_img, w, h) + if aligned is not None: + method_used = "akaze" + except Exception: + pass + + # --- Strategy 4: Phase correlation (rotation + translation in frequency domain) --- + if aligned is None: + try: + aligned = _align_phase_correlation(ref_img, targ_img, w, h) + if aligned is not None: + method_used = "phase" + except Exception: + pass + + # --- Strategy 5: Image moments (centroid + principal axis) --- + if aligned is None: + try: + aligned = _align_moments(ref_img, targ_img, w, h) + if aligned is not None: + method_used = "moments" + except Exception: + pass + + # --- Strategy 6: Ultimate fallback --- + if aligned is None: + aligned = _fallback_resize_center(targ_img, w, h) + method_used = "fallback" + + cv2.imwrite(output_path, aligned) + return method_used + + +def _preprocess(gray): + """Enhance contrast to improve feature detection on blurry/low-contrast images.""" + clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) + return clahe.apply(gray) + + +def _normalize_scales(ref_img, targ_img, max_dim=1024): + """Resize images to similar scales before feature matching.""" + # Scale target to reference scale if they differ too much + ref_max = max(ref_img.shape[:2]) + targ_max = max(targ_img.shape[:2]) + + if ref_max / targ_max > 2.0 or targ_max / ref_max > 2.0: + scale = ref_max / targ_max + new_h = int(targ_img.shape[0] * scale) + new_w = int(targ_img.shape[1] * scale) + targ_img = cv2.resize(targ_img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4) + + # Cap at max_dim for performance + if ref_max > max_dim: + s = max_dim / ref_max + ref_img = cv2.resize(ref_img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA) + if max(targ_img.shape[:2]) > max_dim: + s = max_dim / max(targ_img.shape[:2]) + targ_img = cv2.resize(targ_img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA) + + return ref_img, targ_img + + +def _get_affine_matrix(src_pts, dst_pts): + """Estimate affine matrix with relaxed RANSAC for difficult cases.""" + if len(src_pts) < 3 or len(dst_pts) < 3: + return None + matrix, inliers = cv2.estimateAffinePartial2D( + src_pts, dst_pts, + method=cv2.RANSAC, + ransacReprojThreshold=5.0, + maxIters=5000, + confidence=0.99 + ) + if matrix is None: + return None + if inliers is not None and np.sum(inliers) < 3: + # Very few inliers - try without RANSAC as last resort + matrix, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts, method=cv2.LMEDS) + return matrix + + +def _align_sift(ref_img, targ_img, w, h): + ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY) + targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY) + + # Normalize scales + ref_norm, targ_norm = _normalize_scales(ref_img, targ_img) + ref_g = _preprocess(cv2.cvtColor(ref_norm, cv2.COLOR_BGR2GRAY)) + targ_g = _preprocess(cv2.cvtColor(targ_norm, cv2.COLOR_BGR2GRAY)) + + sift = cv2.SIFT_create(nfeatures=5000) + kp1, des1 = sift.detectAndCompute(ref_g, None) + kp2, des2 = sift.detectAndCompute(targ_g, None) + + if des1 is None or des2 is None or len(kp1) < 6 or len(kp2) < 6: + return None + + # Use BFMatcher instead of FLANN - more stable across scale differences + bf = cv2.BFMatcher(cv2.NORM_L2) + matches = bf.knnMatch(des1, des2, k=2) + + good = [m for m, n in matches if m.distance < 0.75 * n.distance] + if len(good) < 6: + return None + + # Scale keypoints back to original image coordinates + scale_ref = ref_img.shape[1] / ref_norm.shape[1] + scale_targ = targ_img.shape[1] / targ_norm.shape[1] + + src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) * scale_ref + dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) * scale_targ + + matrix = _get_affine_matrix(dst_pts, src_pts) + if matrix is None: + return None + + return cv2.warpAffine(targ_img, matrix, (w, h), + flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_CONSTANT, + borderValue=(0, 0, 0)) + + +def _align_orb(ref_img, targ_img, w, h): + ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY) + targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY) + ref_gray = _preprocess(ref_gray) + targ_gray = _preprocess(targ_gray) + + orb = cv2.ORB_create(nfeatures=5000, scaleFactor=1.2, nlevels=8) + kp1, des1 = orb.detectAndCompute(ref_gray, None) + kp2, des2 = orb.detectAndCompute(targ_gray, None) + + if des1 is None or des2 is None or len(kp1) < 6 or len(kp2) < 6: + return None + + bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False) + matches = bf.knnMatch(des1, des2, k=2) + + good = [] + for pair in matches: + if len(pair) == 2: + m, n = pair + if m.distance < 0.8 * n.distance: + good.append(m) + + if len(good) < 6: + return None + + src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) + dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) + + matrix = _get_affine_matrix(dst_pts, src_pts) + if matrix is None: + return None + + return cv2.warpAffine(targ_img, matrix, (w, h), + flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_CONSTANT, + borderValue=(0, 0, 0)) + + +def _align_akaze(ref_img, targ_img, w, h): + ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY) + targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY) + + akaze = cv2.AKAZE_create() + kp1, des1 = akaze.detectAndCompute(ref_gray, None) + kp2, des2 = akaze.detectAndCompute(targ_gray, None) + + if des1 is None or des2 is None or len(kp1) < 6 or len(kp2) < 6: + return None + + bf = cv2.BFMatcher(cv2.NORM_HAMMING) + matches = bf.knnMatch(des1, des2, k=2) + + good = [m for m, n in matches if m.distance < 0.8 * n.distance] + if len(good) < 6: + return None + + src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) + dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) + + matrix = _get_affine_matrix(dst_pts, src_pts) + if matrix is None: + return None + + return cv2.warpAffine(targ_img, matrix, (w, h), + flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_CONSTANT, + borderValue=(0, 0, 0)) + + +def _align_phase_correlation(ref_img, targ_img, w, h): + """Frequency-domain alignment for translation/rotation.""" + ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY).astype(np.float32) + targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY) + + # Resize target to reference size + targ_resized = cv2.resize(targ_gray, (w, h)).astype(np.float32) + + # Hanning window to reduce edge artifacts + window = cv2.createHanningWindow((w, h), cv2.CV_32F) + + shift, response = cv2.phaseCorrelate(ref_gray * window, targ_resized * window) + + matrix = np.array([[1, 0, shift[0]], [0, 1, shift[1]]], dtype=np.float32) + return cv2.warpAffine(targ_img, matrix, (w, h), + flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_CONSTANT, + borderValue=(0, 0, 0)) + + +def _align_moments(ref_img, targ_img, w, h): + """Align using centroid and principal axis - works even with almost no texture.""" + ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY) + targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY) + + # Otsu threshold to isolate component from cyan background + _, ref_thresh = cv2.threshold(ref_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) + _, targ_thresh = cv2.threshold(targ_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) + + # Clean up noise + kernel = np.ones((5, 5), np.uint8) + ref_thresh = cv2.morphologyEx(ref_thresh, cv2.MORPH_CLOSE, kernel) + targ_thresh = cv2.morphologyEx(targ_thresh, cv2.MORPH_CLOSE, kernel) + + ref_m = cv2.moments(ref_thresh) + targ_m = cv2.moments(targ_thresh) + + if ref_m["m00"] == 0 or targ_m["m00"] == 0: + return None + + # Centroids + rcx, rcy = ref_m["m10"] / ref_m["m00"], ref_m["m01"] / ref_m["m00"] + tcx, tcy = targ_m["m10"] / targ_m["m00"], targ_m["m01"] / targ_m["m00"] + + # Principal axis angles + def principal_angle(m): + return 0.5 * np.arctan2(2 * m["mu11"], m["mu20"] - m["mu02"]) + + r_angle = principal_angle(ref_m) + t_angle = principal_angle(targ_m) + rotation = r_angle - t_angle + + # Scale from area ratio + scale = np.sqrt(ref_m["m00"] / targ_m["m00"]) if targ_m["m00"] > 0 else 1.0 + + cos_r = np.cos(rotation) * scale + sin_r = np.sin(rotation) * scale + + tx = rcx - (cos_r * tcx - sin_r * tcy) + ty = rcy - (sin_r * tcx + cos_r * tcy) + + matrix = np.array([[cos_r, -sin_r, tx], + [sin_r, cos_r, ty]], dtype=np.float32) + + return cv2.warpAffine(targ_img, matrix, (w, h), + flags=cv2.INTER_LANCZOS4, + borderMode=cv2.BORDER_CONSTANT, + borderValue=(0, 0, 0)) + + +def _fallback_resize_center(targ_img, w, h): + """Last resort: center the target in a canvas of reference size.""" + th, tw = targ_img.shape[:2] + + # Scale to fit within reference while preserving aspect ratio + scale = min(w / tw, h / th) * 0.9 # 90% fill + new_w, new_h = int(tw * scale), int(th * scale) + resized = cv2.resize(targ_img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4) + + # Create black canvas and center the image + canvas = np.zeros((h, w, 3), dtype=np.uint8) + y_off = (h - new_h) // 2 + x_off = (w - new_w) // 2 + canvas[y_off:y_off+new_h, x_off:x_off+new_w] = resized + return canvas + + +@app.route("/align", methods=["POST"]) +def align_endpoint(): + if "reference" not in request.files or "target" not in request.files: + return jsonify({"error": "Missing 'reference' or 'target' file in request form-data."}), 400 + + ref_file = request.files["reference"] + target_file = request.files["target"] + + if ref_file.filename == "" or target_file.filename == "": + return jsonify({"error": "No file selected."}), 400 + + ref_name = secure_filename(ref_file.filename) + target_name = secure_filename(target_file.filename) + + input_ref_path = os.path.join(app.config["UPLOAD_FOLDER"], f"ref_{ref_name}") + input_target_path = os.path.join(app.config["UPLOAD_FOLDER"], f"target_{target_name}") + + ref_file.save(input_ref_path) + target_file.save(input_target_path) + + output_aligned_path = os.path.join(app.config["OUTPUT_FOLDER"], f"aligned_{target_name}") + + try: + method_used = align_image_to_reference(input_ref_path, input_target_path, output_aligned_path) + return send_file(output_aligned_path, mimetype="image/png") + + except Exception as e: + return jsonify({"status": "error", "message": str(e)}), 500 + + +if __name__ == "__main__": + app.run(host="0.0.0.0", port=5000, debug=True) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy (2).py b/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy (2).py new file mode 100644 index 0000000000000000000000000000000000000000..574ec490ed229e815ca63a61b1c8ed185bdeb767 --- /dev/null +++ b/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy (2).py @@ -0,0 +1,800 @@ +""" +ArtiAgent Orchestrator โ€” DefectFill Edition (with VLM list selection) + +The VLM selects object_class from a given list and defect_type from the +corresponding per-object-class list. Product description drives the selection. + +Usage: + python artiagent_orchestrator.py \\ + --product-desc "VCSEL laser diode with glass lens cap" \\ + --image ./clean_chip.png \\ + --output-dir ./defect_output \\ + --checkpoint-dir "C:/.../checkpoints" \\ + --valid-object-classes '["xray_PCB","vcsel"]' \\ + --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \\ + --device cuda +""" + +import os +import sys +import json +import argparse +import uuid +import traceback +from pathlib import Path +from typing import Dict, List, Optional, Tuple +from datetime import datetime + +import numpy as np +import torch +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from pipeline.local_vlm_client import LocalVLMClient +from pipeline.prompts import ( + plan_defects_for_product, + artifact_description, + MoneyManager +) +from pipeline.gsam_detector import GSAMDetector +from pipeline.defectfill_generator import DefectFillGenerator, DefectFillConfig +from pipeline.defect_rag import get_rag +from pipeline.domain_router import get_router + +import cv2 + + +def blend_defect_onto_real_image( + real_image: np.ndarray, + defect_image: Image.Image, + defect_mask: np.ndarray, + target_bbox: List[int], + max_defect_ratio: Optional[float] = None, + mask_shape: str = "free" +) -> Tuple[np.ndarray, np.ndarray]: + """Injects a DefectFill defect patch onto a real clean factory image.""" + defect_np = np.array(defect_image) + mask_uint8 = (defect_mask.astype(np.uint8) * 255) if defect_mask.dtype == bool else defect_mask.astype(np.uint8) + + ys, xs = np.where(mask_uint8 > 0) + if len(ys) == 0 or len(xs) == 0: + return real_image.copy(), np.zeros(real_image.shape[:2], dtype=np.uint8) + + y1_d, y2_d = ys.min(), ys.max() + x1_d, x2_d = xs.min(), xs.max() + + defect_patch = defect_np[y1_d:y2_d + 1, x1_d:x2_d + 1] + mask_patch = mask_uint8[y1_d:y2_d + 1, x1_d:x2_d + 1] + + x1_t, y1_t, x2_t, y2_t = target_bbox + target_w = max(1, x2_t - x1_t) + target_h = max(1, y2_t - y1_t) + + if max_defect_ratio is None or max_defect_ratio >= 1.0: + final_w = target_w + final_h = target_h + else: + if max_defect_ratio < 0.2: + max_defect_ratio = 0.2 + scale_factor = np.sqrt(max_defect_ratio) + scaled_w = int(target_w * scale_factor) + scaled_h = int(target_h * scale_factor) + + patch_h, patch_w = defect_patch.shape[:2] + aspect_ratio = patch_w / max(1, patch_h) + + if aspect_ratio > 1: + final_w = max(15, scaled_w) + final_h = max(15, int(final_w / aspect_ratio)) + else: + final_h = max(15, scaled_h) + final_w = max(15, int(final_h * aspect_ratio)) + + defect_patch_resized = cv2.resize(defect_patch, (final_w, final_h), interpolation=cv2.INTER_AREA) + mask_patch_resized = cv2.resize(mask_patch, (final_w, final_h), interpolation=cv2.INTER_NEAREST) + + shape_type = mask_shape.lower().strip() if mask_shape else "free" + + if shape_type == "circle": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + center = (final_w // 2, final_h // 2) + radius = max(1, min(final_w, final_h) // 2 - 1) + cv2.circle(geom_mask, center, radius, 255, thickness=-1) + mask_patch_resized = geom_mask + + elif shape_type == "square": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + side = max(1, min(final_w, final_h) - 2) + top_left_x = (final_w - side) // 2 + top_left_y = (final_h - side) // 2 + cv2.rectangle( + geom_mask, + (top_left_x, top_left_y), + (top_left_x + side, top_left_y + side), + 255, + thickness=-1 + ) + mask_patch_resized = geom_mask + + elif shape_type == "rectangle": + mask_patch_resized = np.full((final_h, final_w), 255, dtype=np.uint8) + + center_x = x1_t + target_w // 2 + center_y = y1_t + target_h // 2 + center = (center_x, center_y) + + real_bgr = cv2.cvtColor(real_image, cv2.COLOR_RGB2BGR) + patch_bgr = cv2.cvtColor(defect_patch_resized, cv2.COLOR_RGB2BGR) + + patch_mean = np.mean(defect_patch_resized) + clone_mode = cv2.NORMAL_CLONE if patch_mean < 30 else cv2.MIXED_CLONE + + blended_bgr = cv2.seamlessClone( + patch_bgr, + real_bgr, + mask_patch_resized, + center, + clone_mode + ) + blended_rgb = cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB) + + full_mask = np.zeros(real_image.shape[:2], dtype=np.uint8) + top_left_x = max(0, center_x - final_w // 2) + top_left_y = max(0, center_y - final_h // 2) + + h_end = min(real_image.shape[0], top_left_y + final_h) + w_end = min(real_image.shape[1], top_left_x + final_w) + + mask_crop_h = h_end - top_left_y + mask_crop_w = w_end - top_left_x + + if mask_crop_h > 0 and mask_crop_w > 0: + full_mask[top_left_y:h_end, top_left_x:w_end] = ( + mask_patch_resized[:mask_crop_h, :mask_crop_w] > 128 + ).astype(np.uint8) + + return blended_rgb, full_mask + + +def create_visual_prompt_image(full_image: np.ndarray, bbox: list) -> np.ndarray: + """Draws a bright neon bounding box on the full image around the target ROI.""" + viz_img = full_image.copy() + x1, y1, x2, y2 = bbox + cv2.rectangle(viz_img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2) + return viz_img + + +def resolve_checkpoint_path(checkpoint_dir: str, object_class: str, defect_type: str) -> str: + """Resolve DefectFill checkpoint path from object_class + defect_type.""" + path = Path(checkpoint_dir) / object_class / defect_type / "checkpoints" / "checkpoint_final.pt" + if not path.exists(): + alt = Path(checkpoint_dir) / object_class / defect_type / "checkpoint_final.pt" + if alt.exists(): + return str(alt) + raise FileNotFoundError( + f"Checkpoint not found for object_class='{object_class}', defect_type='{defect_type}'.\n" + f"Tried: {path}\nAlso tried: {alt}" + ) + return str(path) + + +class ArtiAgentOrchestrator: + """Agentic orchestrator for directed defect generation with DefectFill.""" + + def __init__( + self, + device='cuda', + output_dir='./defect_output', + vlm_model='gemma3:12b', + checkpoint_dir: str = "", + object_class: str = "", + defect_type: str = "", + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + image_size: int = 512, + num_steps: int = 50, + guidance_scale: float = 7.5 + ): + self.device = device + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.vlm_client = LocalVLMClient(model=vlm_model) + self.money_manager = MoneyManager(model="gpt-4o") + + self.gsam_detector = None + self.defectfill_generator = None + + self.checkpoint_dir = checkpoint_dir + self.object_class = object_class + self.defect_type = defect_type + self.valid_object_classes = valid_object_classes or [] + self.valid_defect_types = valid_defect_types or {} + self.image_size = image_size + self.num_steps = num_steps + self.guidance_scale = guidance_scale + + self.rag = get_rag() + self.router = get_router() + + # ------------------------------------------------------------------ + # Lazy initializers + # ------------------------------------------------------------------ + + def _init_gsam(self): + if self.gsam_detector is None or getattr(self.gsam_detector, 'sam_predictor', None) is None: + print("[Agent] Initializing GSAM detector...") + self.gsam_detector = GSAMDetector( + device=self.device, + openai_client=self.vlm_client + ) + + def _init_defectfill(self, object_class: Optional[str] = None, defect_type: Optional[str] = None): + obj_cls = object_class or self.object_class + dfc_type = defect_type or self.defect_type + + if not obj_cls or not dfc_type: + raise ValueError("Both object_class and defect_type must be provided.") + + ckpt_path = resolve_checkpoint_path(self.checkpoint_dir, obj_cls, dfc_type) + + if self.defectfill_generator is None: + print(f"[Agent] Initializing DefectFill: {ckpt_path}") + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + else: + current_ckpt = getattr(self.defectfill_generator, 'ckpt_path', None) + if current_ckpt != ckpt_path: + print(f"[Agent] Switching checkpoint: {ckpt_path}") + self.defectfill_generator.unload_models() + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + else: + print("[Agent] Reusing DefectFill generator.") + + # ------------------------------------------------------------------ + # Step 1: Planning (with list-constrained selection) + # ------------------------------------------------------------------ + + def plan(self, product_description: str, image: np.ndarray, + defect_type: Optional[str] = None, object_class: Optional[str] = None, + max_defects: int = 3): + """Agent plans defects. VLM selects object_class/defect_type from valid lists if not provided.""" + print(f"\n{'='*60}") + print("[Agent] Step 1: Planning defects from product description...") + print(f"Product: {product_description}") + if object_class: + print(f"[Agent] User-provided object_class: {object_class}") + else: + print(f"[Agent] No object_class provided โ€” VLM will select from: {self.valid_object_classes}") + if defect_type: + print(f"[Agent] User-provided defect_type: {defect_type}") + else: + print(f"[Agent] No defect_type provided โ€” VLM will select from valid list.") + + plan = plan_defects_for_product( + self.vlm_client, + product_description, + image, + money_manager=self.money_manager, + target_defect_type=defect_type, + object_class=object_class, + valid_object_classes=self.valid_object_classes if self.valid_object_classes else None, + valid_defect_types=self.valid_defect_types if self.valid_defect_types else None, + max_defects=max_defects + ) + + if plan is None: + raise RuntimeError("Defect planning failed") + + print(f"[Agent] Product type: {plan.product_type}") + print(f"[Agent] Analysis: {plan.analysis}") + print(f"[Agent] Proposed {len(plan.possible_defects)} defects:") + for i, d in enumerate(plan.possible_defects, 1): + print(f" {i}. [{d.defect_type.upper()}] {d.description}") + print(f" Object class: {d.object_class}") + print(f" Target: {d.target_entity} / {d.target_subentity or '(whole)'}") + print(f" Location: {d.location_hint}") + print(f" Coverage: {d.defect_coverage_ratio} | Shape: {d.mask_shape}") + + return plan + + # ------------------------------------------------------------------ + # Step 2: Perception + # ------------------------------------------------------------------ + + def perceive(self, image: np.ndarray, defect_plan): + print(f"\n{'='*60}") + print("[Agent] Step 2: Directed perception (verification bbox)...") + + self._init_gsam() + + entity = defect_plan.target_entity + synonym_map = { + "metal can package": ["TO-can", "metal can", "can body", "package body", "metal ring"], + "lens cap": ["glass lens cap", "lens", "glass dome", "optical window"], + "electrode bars": ["vertical bars", "electrodes", "metal lines"], + "die surface": ["die", "chip die", "semiconductor die", "IC die", "black rectangle", "dark square", "central chip"], + "solder joint": ["solder ball", "BGA ball", "solder bump", "joint"], + } + search_terms = [entity] + synonym_map.get(entity.lower(), []) + + predictions = [] + for term in search_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[defect_plan.target_subentity] if defect_plan.target_subentity else [], + entity_subentity_mapping={}, + min_area_ratio=0.005, + max_area_ratio=0.5, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + if term != entity: + print(f"[Agent] Fallback: detected '{term}' instead of '{entity}'") + break + + # Fallback 1: If entity is tiny (solder balls), ask VLM to detect the GROUP instead of one-by-one + if not predictions and any(k in entity.lower() for k in ['solder', 'ball', 'bump', 'joint']): + group_terms = [ + "array of solder balls", + "BGA ball grid", + "group of solder joints", + "solder ball array", + "BGA array", + "ball grid array" + ] + for term in group_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[], + entity_subentity_mapping={}, + min_area_ratio=0.05, # Group is large enough + max_area_ratio=0.8, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + print(f"[Agent] Group detection: found '{term}' with {len(preds)} prediction(s)") + break + + if not predictions: + print(f"[Agent] Warning: No detections for {entity}; using full image") + h, w = image.shape[:2] + best_pred = { + 'bbox': [0, 0, w, h], + 'pred_mask': torch.ones((h, w), dtype=torch.bool) + } + else: + best_pred = max(predictions, key=lambda p: p.get('area_ratio', 0)) + bbox = best_pred['bbox'] + h, w = image.shape[:2] + + cx = (bbox[0] + bbox[2]) / 2 + cy = (bbox[1] + bbox[3]) / 2 + if cx < 0.1 * w or cx > 0.9 * w or cy < 0.1 * h or cy > 0.9 * h: + print(f"[Agent] Warning: Bbox {bbox} on edge. Falling back to image center.") + best_pred['bbox'] = [int(0.25 * w), int(0.25 * h), int(0.75 * w), int(0.75 * h)] + + return best_pred + + # ------------------------------------------------------------------ + # Step 3: Prepare DefectFill generation conditions + # ------------------------------------------------------------------ + + def prepare_generation_conditions(self, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 3: Preparing DefectFill generation conditions...") + + obj_cls = defect_plan.object_class + if not obj_cls: + raise ValueError("object_class is required for DefectFill prompt standardization.") + + prompt = f"A {obj_cls} with " + defect_class = defect_plan.defect_type + + print(f"[Agent] Standardized prompt: '{prompt}'") + print(f"[Agent] Defect type (checkpoint): {defect_class}") + print(f"[Agent] Object class (checkpoint): {obj_cls}") + + return { + 'prompt': prompt, + 'defect_type': defect_class, + 'object_class': obj_cls, + 'defect_plan': defect_plan + } + + # ------------------------------------------------------------------ + # Step 4: Synthesize with DefectFill + # ------------------------------------------------------------------ + + def synthesize(self, image_patch: Image.Image, mask_patch: np.ndarray, + gen_conditions: Dict) -> Tuple[Image.Image, np.ndarray]: + print(f"\n{'='*60}") + print("[Agent] Step 4: Synthesizing defect with DefectFill...") + + obj_cls = gen_conditions['object_class'] + dfc_type = gen_conditions['defect_type'] + + self._init_defectfill(object_class=obj_cls, defect_type=dfc_type) + + # DefectFill expects PIL image + numpy mask (0-255) + inpainted = self.defectfill_generator.inpaint( + image=image_patch, + mask=mask_patch, + prompt=gen_conditions['prompt'], + defect_type=dfc_type, + seed=42 + ) + return inpainted, mask_patch + + # ------------------------------------------------------------------ + # Step 5: Verification + # ------------------------------------------------------------------ + + def verify(self, original_image: np.ndarray, generated_image: Image.Image, + defect_mask: np.ndarray, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 5: Verifying generated defect...") + + mask_bool = defect_mask.astype(bool) + if mask_bool.sum() == 0: + print("[Agent] Warning: Empty mask; verifying full image") + y1, x1 = 0, 0 + y2, x2 = original_image.shape[0], original_image.shape[1] + else: + ys, xs = np.where(mask_bool) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + margin = 32 + h, w = original_image.shape[:2] + y1 = max(0, y1 - margin) + x1 = max(0, x1 - margin) + y2 = min(h, y2 + margin) + x2 = min(w, x2 + margin) + + gen_crop = np.array(generated_image)[y1:y2, x1:x2] + orig_crop = original_image[y1:y2, x1:x2] + + obj_name = f"a {defect_plan.target_subentity or defect_plan.target_entity}" + artifact_type = "addition" + + result = artifact_description( + self.vlm_client, + original_image, + orig_crop, + gen_crop, + obj_name, + artifact_type, + self.money_manager + ) + + print(f"[Agent] Verification result: has_artifact={result.has_artifact}") + print(f"[Agent] Explanation: {result.explanation}") + print(f"[Agent] Label: {result.label}") + + return { + 'passed': result.has_artifact, + 'explanation': result.explanation, + 'label': result.label + } + + # ------------------------------------------------------------------ + # Main pipeline + # ------------------------------------------------------------------ + + def run(self, product_description: str, image_path: str, + caption: Optional[str] = None, max_defects: int = 3, + defect_type: Optional[str] = None, + object_class: Optional[str] = None) -> Dict: + """Run the full agentic pipeline with DefectFill.""" + start_time = datetime.now() + exp_id = str(uuid.uuid4())[:8] + + # ALWAYS resize to 512x512 โ€” DefectFill model requirement + orig_image = np.array(Image.open(image_path).convert('RGB')) + image = np.array(Image.fromarray(orig_image).resize( + (self.image_size, self.image_size), Image.LANCZOS + )) + print(f"[Agent] Loaded image: {orig_image.shape} -> working at {self.image_size}x{self.image_size}") + + plan = self.plan( + product_description, + image, + defect_type=defect_type, + object_class=object_class, + max_defects=max_defects + ) + + results = [] + defects_to_process = plan.possible_defects[:max_defects] + print(f"[Agent] Processing top {len(defects_to_process)} of {len(plan.possible_defects)} planned defects") + + for i, defect_plan in enumerate(defects_to_process): + print(f"\n{'='*80}") + print(f"[Agent] Defect {i+1}/{len(defects_to_process)}: [{defect_plan.defect_type.upper()}] {defect_plan.description}") + print(f"{'='*80}") + + try: + prediction = self.perceive(image, defect_plan) + target_bbox = prediction.get('bbox', [0, 0, image.shape[1], image.shape[0]]) + + gen_conditions = self.prepare_generation_conditions(defect_plan) + + # ========================================================================= + # Step 4: Smart Crop + DefectFill (inference.py style, always 512x512) + # ========================================================================= + x1, y1, x2, y2 = target_bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) + + # 1. Build full-image mask (0-255) โ€” DEFECT-FIRST: use full bbox + h, w = image.shape[:2] + full_mask = np.zeros((h, w), dtype=np.uint8) + + # Use the ENTIRE detected bbox as the mask region. + # This gives the model enough pixels to generate a realistic defect. + cv2.rectangle(full_mask, (x1, y1), (x2, y2), 255, -1) + + # Optional: slightly dilate the mask so the defect can "bleed" + # naturally into neighboring texture (better for hip/void defects) + dilation_kernel = np.ones((21, 21), np.uint8) + full_mask = cv2.dilate(full_mask, dilation_kernel, iterations=1) + + # 2. Smart crop around defect (exactly like inference.py) + y_idx, x_idx = np.where(full_mask > 0) + if len(y_idx) > 0: + min_y, max_y = np.min(y_idx), np.max(y_idx) + min_x, max_x = np.min(x_idx), np.max(x_idx) + cy_crop = (min_y + max_y) // 2 + cx_crop = (min_x + max_x) // 2 + max_dim = max(max_y - min_y, max_x - min_x) + else: + cy_crop, cx_crop = h // 2, w // 2 + max_dim = 0 + + padding = 50 + crop_size = max(self.image_size, max_dim + padding) + half = crop_size // 2 + + x1c = cx_crop - half + y1c = cy_crop - half + x2c = x1c + crop_size + y2c = y1c + crop_size + + if x1c < 0: x2c -= x1c; x1c = 0 + if y1c < 0: y2c -= y1c; y1c = 0 + if x2c > w: x1c -= (x2c - w); x2c = w + if y2c > h: y1c -= (y2c - h); y2c = h + x1c = max(0, x1c); y1c = max(0, y1c) + x2c = min(w, x2c); y2c = min(h, y2c) + + crop_img = image[y1c:y2c, x1c:x2c] + crop_mask = full_mask[y1c:y2c, x1c:x2c] + + if crop_img.shape[0] != self.image_size or crop_img.shape[1] != self.image_size: + crop_img = cv2.resize(crop_img, (self.image_size, self.image_size), + interpolation=cv2.INTER_AREA) + crop_mask = cv2.resize(crop_mask, (self.image_size, self.image_size), + interpolation=cv2.INTER_NEAREST) + + # 3. DefectFill on the crop + generated_image, _ = self.synthesize( + image_patch=Image.fromarray(crop_img), + mask_patch=crop_mask, + gen_conditions=gen_conditions + ) + gen_np = np.array(generated_image) + + # 4. Paste crop back into full 512x512 canvas + # 4. Paste crop back โ€” MASK-AWARE blending + crop_h, crop_w = y2c - y1c, x2c - x1c + if gen_np.shape[:2] != (crop_h, crop_w): + gen_np = cv2.resize(gen_np, (crop_w, crop_h), interpolation=cv2.INTER_AREA) + mask_back = cv2.resize(crop_mask, (crop_w, crop_h), interpolation=cv2.INTER_NEAREST) + else: + mask_back = crop_mask + + # Prepare original crop for blending + crop_orig = image[y1c:y2c, x1c:x2c].copy() + + # Build soft 3-channel mask (0.0~1.0) for smooth edges + mask_float = (mask_back > 0).astype(np.float32) + mask_float = cv2.GaussianBlur(mask_float, (7, 7), sigmaX=2.0) + mask_3ch = np.stack([mask_float] * 3, axis=-1) + + # Composite: generated inside mask, original background outside + blended_crop = ( + gen_np.astype(np.float32) * mask_3ch + + crop_orig.astype(np.float32) * (1.0 - mask_3ch) + ).astype(np.uint8) + + blended_image = image.copy() + blended_image[y1c:y2c, x1c:x2c] = blended_crop + + full_defect_mask = np.zeros((h, w), dtype=np.uint8) + full_defect_mask[y1c:y2c, x1c:x2c] = mask_back + + ys, xs = np.where(full_defect_mask > 0) + if len(ys) > 0 and len(xs) > 0: + mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] + else: + mask_bbox = target_bbox + + blended_with_green_box = create_visual_prompt_image(blended_image, mask_bbox) + + verification = self.verify(image, Image.fromarray(blended_image), full_defect_mask, defect_plan) + + if not verification['passed']: + defect_dir = self.output_dir / "failed" / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + else: + defect_dir = self.output_dir / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + + defect_dir.mkdir(parents=True, exist_ok=True) + + Image.fromarray(image).save(defect_dir / "real_clean_image.png") + Image.fromarray(blended_image).save(defect_dir / "blended_factory_defect.png") + generated_image.save(defect_dir / "raw_defectfill_patch.png") + + mask_img = Image.fromarray(full_defect_mask) + mask_img.save(defect_dir / "defect_mask.png") + + metadata = { + 'experiment_id': exp_id, + 'product_description': product_description, + 'product_type': plan.product_type, + 'object_class': defect_plan.object_class, + 'defect_type': defect_plan.defect_type, + 'defect_plan': defect_plan.dict() if hasattr(defect_plan, 'dict') else vars(defect_plan), + 'generation_conditions': {k: v for k, v in gen_conditions.items() if k != 'defect_plan'}, + 'verification': verification, + 'timestamp': datetime.now().isoformat() + } + with open(defect_dir / "metadata.json", 'w') as f: + json.dump(metadata, f, indent=2, default=str) + + if not verification['passed']: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': False, + 'verification_passed': False, + 'error': f"VLM verification failed: {verification.get('explanation', 'no explanation')}" + }) + print(f"[Agent] Verification FAILED for {defect_plan.defect_type}.") + else: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': True, + 'verification_passed': verification['passed'], + 'output_dir': str(defect_dir) + }) + print(f"[Agent] Defect {i+1} complete. Saved to {defect_dir}") + + except Exception as e: + print(f"[Agent] ERROR processing defect {i+1}: {str(e)}") + traceback.print_exc() + results.append({ + 'defect_type': getattr(defect_plan, 'defect_type', 'unknown'), + 'object_class': getattr(defect_plan, 'object_class', 'unknown'), + 'success': False, + 'error': str(e) + }) + + if hasattr(self, '_gsam_cache'): + keys_to_remove = [k for k in self._gsam_cache if k.startswith(str(image_path) + "::")] + for k in keys_to_remove: + self._gsam_cache.pop(k, None) + + elapsed = (datetime.now() - start_time).total_seconds() + print(f"\n{'='*60}") + print(f"[Agent] Pipeline complete in {elapsed:.1f}s") + print(f"[Agent] Results: {sum(1 for r in results if r['success'])}/{len(results)} succeeded") + + return { + 'experiment_id': exp_id, + 'product_type': plan.product_type, + 'results': results, + 'output_dir': str(self.output_dir), + 'elapsed_time': elapsed + } + + def cleanup(self): + if self.gsam_detector: + self.gsam_detector.cleanup() + self.gsam_detector = None + if self.defectfill_generator: + self.defectfill_generator.unload_models() + self.defectfill_generator = None + print("[Agent] All models cleaned up.") + + +# ============================================================================= +# CLI +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description='ArtiAgent โ€” DefectFill Edition (with VLM list selection)') + parser.add_argument('--product-desc', required=True, help='Product description (drives VLM selection)') + parser.add_argument('--image', required=True, help='Path to clean product image') + + # DefectFill checkpoint routing + parser.add_argument('--checkpoint-dir', required=True, + help='Root directory containing object_class/defect_type checkpoint subfolders') + parser.add_argument('--object-class', default=None, + help='Object class (optional; VLM selects from valid list if omitted)') + parser.add_argument('--defect-type', default=None, + help='Defect type (optional; VLM selects from valid list if omitted)') + + # Valid lists (auto-discovered from checkpoint-dir if not provided) + parser.add_argument('--valid-object-classes', default=None, + help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" ) + parser.add_argument('--valid-defect-types', default=None, + help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" ) + # Generation control + parser.add_argument('--output-dir', default='./defect_output', help='Output directory') + parser.add_argument('--caption', default=None, help='Optional image caption') + parser.add_argument('--max-defects', type=int, default=3, help='Max defects to generate') + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, help='Generation resolution') + parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps') + parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale') + + args = parser.parse_args() + + # Parse valid lists + valid_object_classes = None + valid_defect_types = None + if args.valid_object_classes: + valid_object_classes = json.loads(args.valid_object_classes) + if args.valid_defect_types: + valid_defect_types = json.loads(args.valid_defect_types) + + orchestrator = ArtiAgentOrchestrator( + device=args.device, + output_dir=args.output_dir, + vlm_model=args.vlm_model, + checkpoint_dir=args.checkpoint_dir, + object_class=args.object_class or "", + defect_type=args.defect_type or "", + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + image_size=args.image_size, + num_steps=args.num_steps, + guidance_scale=args.guidance_scale + ) + + result = orchestrator.run( + product_description=args.product_desc, + image_path=args.image, + caption=args.caption, + max_defects=args.max_defects, + defect_type=args.defect_type, + object_class=args.object_class + ) + + print(f"\nFinal output saved to: {result['output_dir']}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy (3).py b/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy (3).py new file mode 100644 index 0000000000000000000000000000000000000000..fdbfaa8c218a89cf83ea2430cd041ae063034e6f --- /dev/null +++ b/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy (3).py @@ -0,0 +1,833 @@ +""" +ArtiAgent Orchestrator โ€” DefectFill Edition (with VLM list selection) + +The VLM selects object_class from a given list and defect_type from the +corresponding per-object-class list. Product description drives the selection. + +Usage: + python artiagent_orchestrator.py \\ + --product-desc "VCSEL laser diode with glass lens cap" \\ + --image ./clean_chip.png \\ + --output-dir ./defect_output \\ + --checkpoint-dir "C:/.../checkpoints" \\ + --valid-object-classes '["xray_PCB","vcsel"]' \\ + --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \\ + --device cuda +""" + +import os +import sys +import json +import argparse +import uuid +import traceback +from pathlib import Path +from typing import Dict, List, Optional, Tuple +from datetime import datetime + +import numpy as np +import torch +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from pipeline.local_vlm_client import LocalVLMClient +from pipeline.prompts import ( + plan_defects_for_product, + artifact_description, + MoneyManager +) +from pipeline.gsam_detector import GSAMDetector +from pipeline.defectfill_generator import DefectFillGenerator, DefectFillConfig +from pipeline.defect_rag import get_rag +from pipeline.domain_router import get_router + +import cv2 + + +def blend_defect_onto_real_image( + real_image: np.ndarray, + defect_image: Image.Image, + defect_mask: np.ndarray, + target_bbox: List[int], + max_defect_ratio: Optional[float] = None, + mask_shape: str = "free" +) -> Tuple[np.ndarray, np.ndarray]: + """Injects a DefectFill defect patch onto a real clean factory image.""" + defect_np = np.array(defect_image) + mask_uint8 = (defect_mask.astype(np.uint8) * 255) if defect_mask.dtype == bool else defect_mask.astype(np.uint8) + + ys, xs = np.where(mask_uint8 > 0) + if len(ys) == 0 or len(xs) == 0: + return real_image.copy(), np.zeros(real_image.shape[:2], dtype=np.uint8) + + y1_d, y2_d = ys.min(), ys.max() + x1_d, x2_d = xs.min(), xs.max() + + defect_patch = defect_np[y1_d:y2_d + 1, x1_d:x2_d + 1] + mask_patch = mask_uint8[y1_d:y2_d + 1, x1_d:x2_d + 1] + + x1_t, y1_t, x2_t, y2_t = target_bbox + target_w = max(1, x2_t - x1_t) + target_h = max(1, y2_t - y1_t) + + if max_defect_ratio is None or max_defect_ratio >= 1.0: + final_w = target_w + final_h = target_h + else: + if max_defect_ratio < 0.2: + max_defect_ratio = 0.2 + scale_factor = np.sqrt(max_defect_ratio) + scaled_w = int(target_w * scale_factor) + scaled_h = int(target_h * scale_factor) + + patch_h, patch_w = defect_patch.shape[:2] + aspect_ratio = patch_w / max(1, patch_h) + + if aspect_ratio > 1: + final_w = max(15, scaled_w) + final_h = max(15, int(final_w / aspect_ratio)) + else: + final_h = max(15, scaled_h) + final_w = max(15, int(final_h * aspect_ratio)) + + defect_patch_resized = cv2.resize(defect_patch, (final_w, final_h), interpolation=cv2.INTER_AREA) + mask_patch_resized = cv2.resize(mask_patch, (final_w, final_h), interpolation=cv2.INTER_NEAREST) + + shape_type = mask_shape.lower().strip() if mask_shape else "free" + + if shape_type == "circle": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + center = (final_w // 2, final_h // 2) + radius = max(1, min(final_w, final_h) // 2 - 1) + cv2.circle(geom_mask, center, radius, 255, thickness=-1) + mask_patch_resized = geom_mask + + elif shape_type == "square": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + side = max(1, min(final_w, final_h) - 2) + top_left_x = (final_w - side) // 2 + top_left_y = (final_h - side) // 2 + cv2.rectangle( + geom_mask, + (top_left_x, top_left_y), + (top_left_x + side, top_left_y + side), + 255, + thickness=-1 + ) + mask_patch_resized = geom_mask + + elif shape_type == "rectangle": + mask_patch_resized = np.full((final_h, final_w), 255, dtype=np.uint8) + + center_x = x1_t + target_w // 2 + center_y = y1_t + target_h // 2 + center = (center_x, center_y) + + real_bgr = cv2.cvtColor(real_image, cv2.COLOR_RGB2BGR) + patch_bgr = cv2.cvtColor(defect_patch_resized, cv2.COLOR_RGB2BGR) + + patch_mean = np.mean(defect_patch_resized) + clone_mode = cv2.NORMAL_CLONE if patch_mean < 30 else cv2.MIXED_CLONE + + blended_bgr = cv2.seamlessClone( + patch_bgr, + real_bgr, + mask_patch_resized, + center, + clone_mode + ) + blended_rgb = cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB) + + full_mask = np.zeros(real_image.shape[:2], dtype=np.uint8) + top_left_x = max(0, center_x - final_w // 2) + top_left_y = max(0, center_y - final_h // 2) + + h_end = min(real_image.shape[0], top_left_y + final_h) + w_end = min(real_image.shape[1], top_left_x + final_w) + + mask_crop_h = h_end - top_left_y + mask_crop_w = w_end - top_left_x + + if mask_crop_h > 0 and mask_crop_w > 0: + full_mask[top_left_y:h_end, top_left_x:w_end] = ( + mask_patch_resized[:mask_crop_h, :mask_crop_w] > 128 + ).astype(np.uint8) + + return blended_rgb, full_mask + + +def create_visual_prompt_image(full_image: np.ndarray, bbox: list) -> np.ndarray: + """Draws a bright neon bounding box on the full image around the target ROI.""" + viz_img = full_image.copy() + x1, y1, x2, y2 = bbox + cv2.rectangle(viz_img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2) + return viz_img + + +def resolve_checkpoint_path(checkpoint_dir: str, object_class: str, defect_type: str) -> str: + """Resolve DefectFill checkpoint path from object_class + defect_type.""" + path = Path(checkpoint_dir) / object_class / defect_type / "checkpoints" / "checkpoint_final.pt" + if not path.exists(): + alt = Path(checkpoint_dir) / object_class / defect_type / "checkpoint_final.pt" + if alt.exists(): + return str(alt) + raise FileNotFoundError( + f"Checkpoint not found for object_class='{object_class}', defect_type='{defect_type}'.\n" + f"Tried: {path}\nAlso tried: {alt}" + ) + return str(path) + + +class ArtiAgentOrchestrator: + """Agentic orchestrator for directed defect generation with DefectFill.""" + + def __init__( + self, + device='cuda', + output_dir='./defect_output', + vlm_model='gemma3:12b', + checkpoint_dir: str = "", + object_class: str = "", + defect_type: str = "", + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + image_size: int = 512, + num_steps: int = 50, + guidance_scale: float = 7.5 + ): + self.device = device + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.vlm_client = LocalVLMClient(model=vlm_model) + self.money_manager = MoneyManager(model="gpt-4o") + + self.gsam_detector = None + self.defectfill_generator = None + + self.checkpoint_dir = checkpoint_dir + self.object_class = object_class + self.defect_type = defect_type + self.valid_object_classes = valid_object_classes or [] + self.valid_defect_types = valid_defect_types or {} + self.image_size = image_size + self.num_steps = num_steps + self.guidance_scale = guidance_scale + + self.rag = get_rag() + self.router = get_router() + + # ------------------------------------------------------------------ + # Lazy initializers + # ------------------------------------------------------------------ + + def _init_gsam(self): + if self.gsam_detector is None or getattr(self.gsam_detector, 'sam_predictor', None) is None: + print("[Agent] Initializing GSAM detector...") + self.gsam_detector = GSAMDetector( + device=self.device, + openai_client=self.vlm_client + ) + + def _init_defectfill(self, object_class: Optional[str] = None, defect_type: Optional[str] = None): + obj_cls = object_class or self.object_class + dfc_type = defect_type or self.defect_type + + if not obj_cls or not dfc_type: + raise ValueError("Both object_class and defect_type must be provided.") + + ckpt_path = resolve_checkpoint_path(self.checkpoint_dir, obj_cls, dfc_type) + + if self.defectfill_generator is None: + print(f"[Agent] Initializing DefectFill: {ckpt_path}") + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + else: + current_ckpt = getattr(self.defectfill_generator, 'ckpt_path', None) + if current_ckpt != ckpt_path: + print(f"[Agent] Switching checkpoint: {ckpt_path}") + self.defectfill_generator.unload_models() + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + else: + print("[Agent] Reusing DefectFill generator.") + + # ------------------------------------------------------------------ + # Step 1: Planning (with list-constrained selection) + # ------------------------------------------------------------------ + + def plan(self, product_description: str, image: np.ndarray, + defect_type: Optional[str] = None, object_class: Optional[str] = None, + max_defects: int = 3): + """Agent plans defects. VLM selects object_class/defect_type from valid lists if not provided.""" + print(f"\n{'='*60}") + print("[Agent] Step 1: Planning defects from product description...") + print(f"Product: {product_description}") + if object_class: + print(f"[Agent] User-provided object_class: {object_class}") + else: + print(f"[Agent] No object_class provided โ€” VLM will select from: {self.valid_object_classes}") + if defect_type: + print(f"[Agent] User-provided defect_type: {defect_type}") + else: + print(f"[Agent] No defect_type provided โ€” VLM will select from valid list.") + + plan = plan_defects_for_product( + self.vlm_client, + product_description, + image, + money_manager=self.money_manager, + target_defect_type=defect_type, + object_class=object_class, + valid_object_classes=self.valid_object_classes if self.valid_object_classes else None, + valid_defect_types=self.valid_defect_types if self.valid_defect_types else None, + max_defects=max_defects + ) + + if plan is None: + raise RuntimeError("Defect planning failed") + + print(f"[Agent] Product type: {plan.product_type}") + print(f"[Agent] Analysis: {plan.analysis}") + print(f"[Agent] Proposed {len(plan.possible_defects)} defects:") + for i, d in enumerate(plan.possible_defects, 1): + print(f" {i}. [{d.defect_type.upper()}] {d.description}") + print(f" Object class: {d.object_class}") + print(f" Target: {d.target_entity} / {d.target_subentity or '(whole)'}") + print(f" Location: {d.location_hint}") + print(f" Coverage: {d.defect_coverage_ratio} | Shape: {d.mask_shape}") + + return plan + + # ------------------------------------------------------------------ + # Step 2: Perception + # ------------------------------------------------------------------ + + def perceive(self, image: np.ndarray, defect_plan): + print(f"\n{'='*60}") + print("[Agent] Step 2: Directed perception (verification bbox)...") + + self._init_gsam() + + entity = defect_plan.target_entity + + is_solder_target = any(k in entity.lower() for k in ['solder', 'ball', 'bump', 'joint', 'bga']) + + # ------------------------------------------------------------------ + # FAST PATH: For solder-ball arrays, use OpenCV blob detection directly + # ------------------------------------------------------------------ + if is_solder_target: + print(f"[Agent] Solder-ball target detected. Trying blob-detection fast path...") + preds, _, viz = self.gsam_detector.detect_solder_ball_array(image) + if len(preds) > 0: + pred = preds[0] + print(f"[Agent] Blob detection succeeded: bbox={preds[0]['bbox']}, " + f"area_ratio={pred.get('area_ratio', 0):.4f}") + return pred + + + # ------------------------------------------------------------------ + # FALLBACK: VLM + SAM (with relaxed area threshold for small parts) + # ------------------------------------------------------------------ + synonym_map = { + "metal can package": ["TO-can", "metal can", "can body", "package body", "metal ring"], + "lens cap": ["glass lens cap", "lens", "glass dome", "optical window"], + "electrode bars": ["vertical bars", "electrodes", "metal lines"], + "die surface": ["die", "chip die", "semiconductor die", "IC die", "black rectangle", "dark square", "central chip"], + "solder joint": ["solder ball", "BGA ball", "solder bump", "joint"], + } + search_terms = [entity] + synonym_map.get(entity.lower(), []) + + predictions = [] + for term in search_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[defect_plan.target_subentity] if defect_plan.target_subentity else [], + entity_subentity_mapping={}, + min_area_ratio=0.005, + max_area_ratio=0.5, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + if term != entity: + print(f"[Agent] Fallback: detected '{term}' instead of '{entity}'") + break + + # Fallback: group detection for solder balls + if not predictions and any(k in entity.lower() for k in ['solder', 'ball', 'bump', 'joint']): + group_terms = [ + "array of solder balls", + "BGA ball grid", + "group of solder joints", + "solder ball array", + "BGA array", + "ball grid array" + ] + for term in group_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[], + entity_subentity_mapping={}, + min_area_ratio=0.05, # Group is large enough + max_area_ratio=0.8, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + print(f"[Agent] Group detection: found '{term}' with {len(preds)} prediction(s)") + break + + if not predictions: + print(f"[Agent] Warning: No detections for {entity}; using full image") + h, w = image.shape[:2] + best_pred = { + 'bbox': [0, 0, w, h], + 'pred_mask': torch.ones((h, w), dtype=torch.bool) + } + else: + best_pred = max(predictions, key=lambda p: p.get('area_ratio', 0)) + bbox = best_pred['bbox'] + h, w = image.shape[:2] + + # Sanity check: solder-ball array should be in upper half of X-ray + if is_solder_target: + cy = (bbox[1] + bbox[3]) / 2 + if cy > h * 0.6: + print(f"[Agent] Warning: Detected bbox {bbox} is too low for solder array. " + f"Forcing upper-half fallback.") + best_pred['bbox'] = [int(0.15 * w), int(0.10 * h), + int(0.85 * w), int(0.50 * h)] + + # General edge check + cx = (bbox[0] + bbox[2]) / 2 + cy = (bbox[1] + bbox[3]) / 2 + if cx < 0.05 * w or cx > 0.95 * w or cy < 0.05 * h or cy > 0.95 * h: + print(f"[Agent] Warning: Bbox {bbox} on extreme edge. Falling back to center.") + best_pred['bbox'] = [int(0.20 * w), int(0.20 * h), + int(0.80 * w), int(0.80 * h)] + + return best_pred + + # ------------------------------------------------------------------ + # Step 3: Prepare DefectFill generation conditions + # ------------------------------------------------------------------ + + def prepare_generation_conditions(self, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 3: Preparing DefectFill generation conditions...") + + obj_cls = defect_plan.object_class + if not obj_cls: + raise ValueError("object_class is required for DefectFill prompt standardization.") + + prompt = f"A {obj_cls} with " + defect_class = defect_plan.defect_type + + print(f"[Agent] Standardized prompt: '{prompt}'") + print(f"[Agent] Defect type (checkpoint): {defect_class}") + print(f"[Agent] Object class (checkpoint): {obj_cls}") + + return { + 'prompt': prompt, + 'defect_type': defect_class, + 'object_class': obj_cls, + 'defect_plan': defect_plan + } + + # ------------------------------------------------------------------ + # Step 4: Synthesize with DefectFill + # ------------------------------------------------------------------ + + def synthesize(self, image_patch: Image.Image, mask_patch: np.ndarray, + gen_conditions: Dict) -> Tuple[Image.Image, np.ndarray]: + print(f"\n{'='*60}") + print("[Agent] Step 4: Synthesizing defect with DefectFill...") + + obj_cls = gen_conditions['object_class'] + dfc_type = gen_conditions['defect_type'] + + self._init_defectfill(object_class=obj_cls, defect_type=dfc_type) + + # DefectFill expects PIL image + numpy mask (0-255) + inpainted = self.defectfill_generator.inpaint( + image=image_patch, + mask=mask_patch, + prompt=gen_conditions['prompt'], + defect_type=dfc_type, + seed=42 + ) + return inpainted, mask_patch + + # ------------------------------------------------------------------ + # Step 5: Verification + # ------------------------------------------------------------------ + + def verify(self, original_image: np.ndarray, generated_image: Image.Image, + defect_mask: np.ndarray, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 5: Verifying generated defect...") + + mask_bool = defect_mask.astype(bool) + if mask_bool.sum() == 0: + print("[Agent] Warning: Empty mask; verifying full image") + y1, x1 = 0, 0 + y2, x2 = original_image.shape[0], original_image.shape[1] + else: + ys, xs = np.where(mask_bool) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + margin = 32 + h, w = original_image.shape[:2] + y1 = max(0, y1 - margin) + x1 = max(0, x1 - margin) + y2 = min(h, y2 + margin) + x2 = min(w, x2 + margin) + + gen_crop = np.array(generated_image)[y1:y2, x1:x2] + orig_crop = original_image[y1:y2, x1:x2] + + obj_name = f"a {defect_plan.target_subentity or defect_plan.target_entity}" + artifact_type = "addition" + + result = artifact_description( + self.vlm_client, + original_image, + orig_crop, + gen_crop, + obj_name, + artifact_type, + self.money_manager + ) + + print(f"[Agent] Verification result: has_artifact={result.has_artifact}") + print(f"[Agent] Explanation: {result.explanation}") + print(f"[Agent] Label: {result.label}") + + return { + 'passed': result.has_artifact, + 'explanation': result.explanation, + 'label': result.label + } + + # ------------------------------------------------------------------ + # Main pipeline + # ------------------------------------------------------------------ + + def run(self, product_description: str, image_path: str, + caption: Optional[str] = None, max_defects: int = 3, + defect_type: Optional[str] = None, + object_class: Optional[str] = None) -> Dict: + """Run the full agentic pipeline with DefectFill.""" + start_time = datetime.now() + exp_id = str(uuid.uuid4())[:8] + + # ALWAYS resize to 512x512 โ€” DefectFill model requirement + orig_image = np.array(Image.open(image_path).convert('RGB')) + image = np.array(Image.fromarray(orig_image).resize( + (self.image_size, self.image_size), Image.LANCZOS + )) + print(f"[Agent] Loaded image: {orig_image.shape} -> working at {self.image_size}x{self.image_size}") + + plan = self.plan( + product_description, + image, + defect_type=defect_type, + object_class=object_class, + max_defects=max_defects + ) + + results = [] + defects_to_process = plan.possible_defects[:max_defects] + print(f"[Agent] Processing top {len(defects_to_process)} of {len(plan.possible_defects)} planned defects") + + for i, defect_plan in enumerate(defects_to_process): + print(f"\n{'='*80}") + print(f"[Agent] Defect {i+1}/{len(defects_to_process)}: [{defect_plan.defect_type.upper()}] {defect_plan.description}") + print(f"{'='*80}") + + try: + prediction = self.perceive(image, defect_plan) + target_bbox = prediction.get('bbox', [0, 0, image.shape[1], image.shape[0]]) + + gen_conditions = self.prepare_generation_conditions(defect_plan) + + # ========================================================================= + # Step 4: Smart Crop + DefectFill (inference.py style, always 512x512) + # ========================================================================= + x1, y1, x2, y2 = target_bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) + + # 1. Build full-image mask (0-255) โ€” DEFECT-FIRST: use full bbox + h, w = image.shape[:2] + full_mask = np.zeros((h, w), dtype=np.uint8) + + # Use the ENTIRE detected bbox as the mask region. + # This gives the model enough pixels to generate a realistic defect. + cv2.rectangle(full_mask, (x1, y1), (x2, y2), 255, -1) + + # Optional: slightly dilate the mask so the defect can "bleed" + # naturally into neighboring texture (better for hip/void defects) + dilation_kernel = np.ones((21, 21), np.uint8) + full_mask = cv2.dilate(full_mask, dilation_kernel, iterations=1) + + # 2. Smart crop around defect (exactly like inference.py) + y_idx, x_idx = np.where(full_mask > 0) + if len(y_idx) > 0: + min_y, max_y = np.min(y_idx), np.max(y_idx) + min_x, max_x = np.min(x_idx), np.max(x_idx) + cy_crop = (min_y + max_y) // 2 + cx_crop = (min_x + max_x) // 2 + max_dim = max(max_y - min_y, max_x - min_x) + else: + cy_crop, cx_crop = h // 2, w // 2 + max_dim = 0 + + padding = 50 + crop_size = max(self.image_size, max_dim + padding) + half = crop_size // 2 + + x1c = cx_crop - half + y1c = cy_crop - half + x2c = x1c + crop_size + y2c = y1c + crop_size + + if x1c < 0: x2c -= x1c; x1c = 0 + if y1c < 0: y2c -= y1c; y1c = 0 + if x2c > w: x1c -= (x2c - w); x2c = w + if y2c > h: y1c -= (y2c - h); y2c = h + x1c = max(0, x1c); y1c = max(0, y1c) + x2c = min(w, x2c); y2c = min(h, y2c) + + crop_img = image[y1c:y2c, x1c:x2c] + crop_mask = full_mask[y1c:y2c, x1c:x2c] + + if crop_img.shape[0] != self.image_size or crop_img.shape[1] != self.image_size: + crop_img = cv2.resize(crop_img, (self.image_size, self.image_size), + interpolation=cv2.INTER_AREA) + crop_mask = cv2.resize(crop_mask, (self.image_size, self.image_size), + interpolation=cv2.INTER_NEAREST) + + # 3. DefectFill on the crop + generated_image, _ = self.synthesize( + image_patch=Image.fromarray(crop_img), + mask_patch=crop_mask, + gen_conditions=gen_conditions + ) + gen_np = np.array(generated_image) + + # 4. Paste crop back into full 512x512 canvas + # 4. Paste crop back โ€” MASK-AWARE blending + crop_h, crop_w = y2c - y1c, x2c - x1c + if gen_np.shape[:2] != (crop_h, crop_w): + gen_np = cv2.resize(gen_np, (crop_w, crop_h), interpolation=cv2.INTER_AREA) + mask_back = cv2.resize(crop_mask, (crop_w, crop_h), interpolation=cv2.INTER_NEAREST) + else: + mask_back = crop_mask + + # Get original crop for blending + orig_crop = image[y1c:y2c, x1c:x2c].copy() + + # Build soft alpha mask (0.0 ~ 1.0) + # Gaussian blur creates feathered edges for smooth transition + mask_float = (mask_back > 0).astype(np.float32) + mask_float = cv2.GaussianBlur(mask_float, (15, 15), sigmaX=5.0) + + # Expand to 3 channels + alpha = np.stack([mask_float] * 3, axis=-1) + + # Alpha blend: generated inside mask, original outside + blended_crop = ( + gen_np.astype(np.float32) * alpha + + orig_crop.astype(np.float32) * (1.0 - alpha) + ).astype(np.uint8) + + blended_image = image.copy() + blended_image[y1c:y2c, x1c:x2c] = blended_crop + + full_defect_mask = np.zeros((h, w), dtype=np.uint8) + full_defect_mask[y1c:y2c, x1c:x2c] = mask_back + + ys, xs = np.where(full_defect_mask > 0) + if len(ys) > 0 and len(xs) > 0: + mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] + else: + mask_bbox = target_bbox + + blended_with_green_box = create_visual_prompt_image(blended_image, mask_bbox) + + verification = self.verify(image, Image.fromarray(blended_image), full_defect_mask, defect_plan) + + if not verification['passed']: + defect_dir = self.output_dir / "failed" / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + else: + defect_dir = self.output_dir / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + + defect_dir.mkdir(parents=True, exist_ok=True) + + Image.fromarray(image).save(defect_dir / "real_clean_image.png") + Image.fromarray(blended_image).save(defect_dir / "blended_factory_defect.png") + generated_image.save(defect_dir / "raw_defectfill_patch.png") + + mask_img = Image.fromarray(full_defect_mask) + mask_img.save(defect_dir / "defect_mask.png") + + metadata = { + 'experiment_id': exp_id, + 'product_description': product_description, + 'product_type': plan.product_type, + 'object_class': defect_plan.object_class, + 'defect_type': defect_plan.defect_type, + 'defect_plan': defect_plan.dict() if hasattr(defect_plan, 'dict') else vars(defect_plan), + 'generation_conditions': {k: v for k, v in gen_conditions.items() if k != 'defect_plan'}, + 'verification': verification, + 'timestamp': datetime.now().isoformat() + } + with open(defect_dir / "metadata.json", 'w') as f: + json.dump(metadata, f, indent=2, default=str) + + if not verification['passed']: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': False, + 'verification_passed': False, + 'error': f"VLM verification failed: {verification.get('explanation', 'no explanation')}" + }) + print(f"[Agent] Verification FAILED for {defect_plan.defect_type}.") + else: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': True, + 'verification_passed': verification['passed'], + 'output_dir': str(defect_dir) + }) + print(f"[Agent] Defect {i+1} complete. Saved to {defect_dir}") + + except Exception as e: + print(f"[Agent] ERROR processing defect {i+1}: {str(e)}") + traceback.print_exc() + results.append({ + 'defect_type': getattr(defect_plan, 'defect_type', 'unknown'), + 'object_class': getattr(defect_plan, 'object_class', 'unknown'), + 'success': False, + 'error': str(e) + }) + + if hasattr(self, '_gsam_cache'): + keys_to_remove = [k for k in self._gsam_cache if k.startswith(str(image_path) + "::")] + for k in keys_to_remove: + self._gsam_cache.pop(k, None) + + elapsed = (datetime.now() - start_time).total_seconds() + print(f"\n{'='*60}") + print(f"[Agent] Pipeline complete in {elapsed:.1f}s") + print(f"[Agent] Results: {sum(1 for r in results if r['success'])}/{len(results)} succeeded") + + return { + 'experiment_id': exp_id, + 'product_type': plan.product_type, + 'results': results, + 'output_dir': str(self.output_dir), + 'elapsed_time': elapsed + } + + def cleanup(self): + if self.gsam_detector: + self.gsam_detector.cleanup() + self.gsam_detector = None + if self.defectfill_generator: + self.defectfill_generator.unload_models() + self.defectfill_generator = None + print("[Agent] All models cleaned up.") + + +# ============================================================================= +# CLI +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description='ArtiAgent โ€” DefectFill Edition (with VLM list selection)') + parser.add_argument('--product-desc', required=True, help='Product description (drives VLM selection)') + parser.add_argument('--image', required=True, help='Path to clean product image') + + # DefectFill checkpoint routing + parser.add_argument('--checkpoint-dir', required=True, + help='Root directory containing object_class/defect_type checkpoint subfolders') + parser.add_argument('--object-class', default=None, + help='Object class (optional; VLM selects from valid list if omitted)') + parser.add_argument('--defect-type', default=None, + help='Defect type (optional; VLM selects from valid list if omitted)') + + # Valid lists (auto-discovered from checkpoint-dir if not provided) + parser.add_argument('--valid-object-classes', default=None, + help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" ) + parser.add_argument('--valid-defect-types', default=None, + help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" ) + # Generation control + parser.add_argument('--output-dir', default='./defect_output', help='Output directory') + parser.add_argument('--caption', default=None, help='Optional image caption') + parser.add_argument('--max-defects', type=int, default=3, help='Max defects to generate') + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, help='Generation resolution') + parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps') + parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale') + + args = parser.parse_args() + + # Parse valid lists + valid_object_classes = None + valid_defect_types = None + if args.valid_object_classes: + valid_object_classes = json.loads(args.valid_object_classes) + if args.valid_defect_types: + valid_defect_types = json.loads(args.valid_defect_types) + + orchestrator = ArtiAgentOrchestrator( + device=args.device, + output_dir=args.output_dir, + vlm_model=args.vlm_model, + checkpoint_dir=args.checkpoint_dir, + object_class=args.object_class or "", + defect_type=args.defect_type or "", + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + image_size=args.image_size, + num_steps=args.num_steps, + guidance_scale=args.guidance_scale + ) + + result = orchestrator.run( + product_description=args.product_desc, + image_path=args.image, + caption=args.caption, + max_defects=args.max_defects, + defect_type=args.defect_type, + object_class=args.object_class + ) + + print(f"\nFinal output saved to: {result['output_dir']}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy.py b/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy.py new file mode 100644 index 0000000000000000000000000000000000000000..c9aa17ec85c43e5dc94fb5dffa1e572c50471199 --- /dev/null +++ b/ArtiAgent - DefectFill/src/artiagent_orchestrator - Copy.py @@ -0,0 +1,806 @@ +""" +ArtiAgent Orchestrator โ€” DefectFill Edition (with VLM list selection) + +The VLM selects object_class from a given list and defect_type from the +corresponding per-object-class list. Product description drives the selection. + +Usage: + python artiagent_orchestrator.py \\ + --product-desc "VCSEL laser diode with glass lens cap" \\ + --image ./clean_chip.png \\ + --output-dir ./defect_output \\ + --checkpoint-dir "C:/.../checkpoints" \\ + --valid-object-classes '["xray_PCB","vcsel"]' \\ + --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \\ + --device cuda +""" + +import os +import sys +import json +import argparse +import uuid +import traceback +from pathlib import Path +from typing import Dict, List, Optional, Tuple +from datetime import datetime + +import numpy as np +import torch +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from pipeline.local_vlm_client import LocalVLMClient +from pipeline.prompts import ( + plan_defects_for_product, + artifact_description, + MoneyManager +) +from pipeline.gsam_detector import GSAMDetector +from pipeline.defectfill_generator import DefectFillGenerator, DefectFillConfig +from pipeline.defect_rag import get_rag +from pipeline.domain_router import get_router + +import cv2 + + +def blend_defect_onto_real_image( + real_image: np.ndarray, + defect_image: Image.Image, + defect_mask: np.ndarray, + target_bbox: List[int], + max_defect_ratio: Optional[float] = None, + mask_shape: str = "free" +) -> Tuple[np.ndarray, np.ndarray]: + """Injects a DefectFill defect patch onto a real clean factory image.""" + defect_np = np.array(defect_image) + mask_uint8 = (defect_mask.astype(np.uint8) * 255) if defect_mask.dtype == bool else defect_mask.astype(np.uint8) + + ys, xs = np.where(mask_uint8 > 0) + if len(ys) == 0 or len(xs) == 0: + return real_image.copy(), np.zeros(real_image.shape[:2], dtype=np.uint8) + + y1_d, y2_d = ys.min(), ys.max() + x1_d, x2_d = xs.min(), xs.max() + + defect_patch = defect_np[y1_d:y2_d + 1, x1_d:x2_d + 1] + mask_patch = mask_uint8[y1_d:y2_d + 1, x1_d:x2_d + 1] + + x1_t, y1_t, x2_t, y2_t = target_bbox + target_w = max(1, x2_t - x1_t) + target_h = max(1, y2_t - y1_t) + + if max_defect_ratio is None or max_defect_ratio >= 1.0: + final_w = target_w + final_h = target_h + else: + if max_defect_ratio < 0.2: + max_defect_ratio = 0.2 + scale_factor = np.sqrt(max_defect_ratio) + scaled_w = int(target_w * scale_factor) + scaled_h = int(target_h * scale_factor) + + patch_h, patch_w = defect_patch.shape[:2] + aspect_ratio = patch_w / max(1, patch_h) + + if aspect_ratio > 1: + final_w = max(15, scaled_w) + final_h = max(15, int(final_w / aspect_ratio)) + else: + final_h = max(15, scaled_h) + final_w = max(15, int(final_h * aspect_ratio)) + + defect_patch_resized = cv2.resize(defect_patch, (final_w, final_h), interpolation=cv2.INTER_AREA) + mask_patch_resized = cv2.resize(mask_patch, (final_w, final_h), interpolation=cv2.INTER_NEAREST) + + shape_type = mask_shape.lower().strip() if mask_shape else "free" + + if shape_type == "circle": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + center = (final_w // 2, final_h // 2) + radius = max(1, min(final_w, final_h) // 2 - 1) + cv2.circle(geom_mask, center, radius, 255, thickness=-1) + mask_patch_resized = geom_mask + + elif shape_type == "square": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + side = max(1, min(final_w, final_h) - 2) + top_left_x = (final_w - side) // 2 + top_left_y = (final_h - side) // 2 + cv2.rectangle( + geom_mask, + (top_left_x, top_left_y), + (top_left_x + side, top_left_y + side), + 255, + thickness=-1 + ) + mask_patch_resized = geom_mask + + elif shape_type == "rectangle": + mask_patch_resized = np.full((final_h, final_w), 255, dtype=np.uint8) + + center_x = x1_t + target_w // 2 + center_y = y1_t + target_h // 2 + center = (center_x, center_y) + + real_bgr = cv2.cvtColor(real_image, cv2.COLOR_RGB2BGR) + patch_bgr = cv2.cvtColor(defect_patch_resized, cv2.COLOR_RGB2BGR) + + patch_mean = np.mean(defect_patch_resized) + clone_mode = cv2.NORMAL_CLONE if patch_mean < 30 else cv2.MIXED_CLONE + + blended_bgr = cv2.seamlessClone( + patch_bgr, + real_bgr, + mask_patch_resized, + center, + clone_mode + ) + blended_rgb = cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB) + + full_mask = np.zeros(real_image.shape[:2], dtype=np.uint8) + top_left_x = max(0, center_x - final_w // 2) + top_left_y = max(0, center_y - final_h // 2) + + h_end = min(real_image.shape[0], top_left_y + final_h) + w_end = min(real_image.shape[1], top_left_x + final_w) + + mask_crop_h = h_end - top_left_y + mask_crop_w = w_end - top_left_x + + if mask_crop_h > 0 and mask_crop_w > 0: + full_mask[top_left_y:h_end, top_left_x:w_end] = ( + mask_patch_resized[:mask_crop_h, :mask_crop_w] > 128 + ).astype(np.uint8) + + return blended_rgb, full_mask + + +def create_visual_prompt_image(full_image: np.ndarray, bbox: list) -> np.ndarray: + """Draws a bright neon bounding box on the full image around the target ROI.""" + viz_img = full_image.copy() + x1, y1, x2, y2 = bbox + cv2.rectangle(viz_img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2) + return viz_img + + +def resolve_checkpoint_path(checkpoint_dir: str, object_class: str, defect_type: str) -> str: + """Resolve DefectFill checkpoint path from object_class + defect_type.""" + path = Path(checkpoint_dir) / object_class / defect_type / "checkpoints" / "checkpoint_final.pt" + if not path.exists(): + alt = Path(checkpoint_dir) / object_class / defect_type / "checkpoint_final.pt" + if alt.exists(): + return str(alt) + raise FileNotFoundError( + f"Checkpoint not found for object_class='{object_class}', defect_type='{defect_type}'.\n" + f"Tried: {path}\nAlso tried: {alt}" + ) + return str(path) + + +class ArtiAgentOrchestrator: + """Agentic orchestrator for directed defect generation with DefectFill.""" + + def __init__( + self, + device='cuda', + output_dir='./defect_output', + vlm_model='gemma3:12b', + checkpoint_dir: str = "", + object_class: str = "", + defect_type: str = "", + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + image_size: int = 512, + num_steps: int = 50, + guidance_scale: float = 7.5 + ): + self.device = device + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.vlm_client = LocalVLMClient(model=vlm_model) + self.money_manager = MoneyManager(model="gpt-4o") + + self.gsam_detector = None + self.defectfill_generator = None + + self.checkpoint_dir = checkpoint_dir + self.object_class = object_class + self.defect_type = defect_type + self.valid_object_classes = valid_object_classes or [] + self.valid_defect_types = valid_defect_types or {} + self.image_size = image_size + self.num_steps = num_steps + self.guidance_scale = guidance_scale + + self.rag = get_rag() + self.router = get_router() + + # ------------------------------------------------------------------ + # Lazy initializers + # ------------------------------------------------------------------ + + def _init_gsam(self): + if self.gsam_detector is None or getattr(self.gsam_detector, 'sam_predictor', None) is None: + print("[Agent] Initializing GSAM detector...") + self.gsam_detector = GSAMDetector( + device=self.device, + openai_client=self.vlm_client + ) + + def _init_defectfill(self, object_class: Optional[str] = None, defect_type: Optional[str] = None): + """Initialize DefectFill generator with the correct checkpoint.""" + obj_cls = object_class or self.object_class + dfc_type = defect_type or self.defect_type + + if not obj_cls or not dfc_type: + raise ValueError("Both object_class and defect_type must be provided to initialize DefectFill.") + + ckpt_path = resolve_checkpoint_path(self.checkpoint_dir, obj_cls, dfc_type) + + if self.defectfill_generator is None: + print(f"[Agent] Initializing DefectFill generator...") + print(f"[Agent] Checkpoint: {ckpt_path}") + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + # After DefectFillGenerator is created, force-load the checkpoint properly + from utils import load_checkpoint + load_checkpoint(self.defectfill_generator.model, None, ckpt_path) + else: + current_ckpt = getattr(self.defectfill_generator, 'ckpt_path', None) + if current_ckpt != ckpt_path: + print(f"[Agent] Switching DefectFill checkpoint to: {ckpt_path}") + self.defectfill_generator.unload_models() + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + # After DefectFillGenerator is created, force-load the checkpoint properly + from utils import load_checkpoint + load_checkpoint(self.defectfill_generator.model, None, ckpt_path) + + # ------------------------------------------------------------------ + # Step 1: Planning (with list-constrained selection) + # ------------------------------------------------------------------ + + def plan(self, product_description: str, image: np.ndarray, + defect_type: Optional[str] = None, object_class: Optional[str] = None, + max_defects: int = 3): + """Agent plans defects. VLM selects object_class/defect_type from valid lists if not provided.""" + print(f"\n{'='*60}") + print("[Agent] Step 1: Planning defects from product description...") + print(f"Product: {product_description}") + if object_class: + print(f"[Agent] User-provided object_class: {object_class}") + else: + print(f"[Agent] No object_class provided โ€” VLM will select from: {self.valid_object_classes}") + if defect_type: + print(f"[Agent] User-provided defect_type: {defect_type}") + else: + print(f"[Agent] No defect_type provided โ€” VLM will select from valid list.") + + plan = plan_defects_for_product( + self.vlm_client, + product_description, + image, + money_manager=self.money_manager, + target_defect_type=defect_type, + object_class=object_class, + valid_object_classes=self.valid_object_classes if self.valid_object_classes else None, + valid_defect_types=self.valid_defect_types if self.valid_defect_types else None, + max_defects=max_defects + ) + + if plan is None: + raise RuntimeError("Defect planning failed") + + print(f"[Agent] Product type: {plan.product_type}") + print(f"[Agent] Analysis: {plan.analysis}") + print(f"[Agent] Proposed {len(plan.possible_defects)} defects:") + for i, d in enumerate(plan.possible_defects, 1): + print(f" {i}. [{d.defect_type.upper()}] {d.description}") + print(f" Object class: {d.object_class}") + print(f" Target: {d.target_entity} / {d.target_subentity or '(whole)'}") + print(f" Location: {d.location_hint}") + print(f" Coverage: {d.defect_coverage_ratio} | Shape: {d.mask_shape}") + + return plan + + # ------------------------------------------------------------------ + # Step 2: Perception + # ------------------------------------------------------------------ + + def perceive(self, image: np.ndarray, defect_plan): + print(f"\n{'='*60}") + print("[Agent] Step 2: Directed perception (verification bbox)...") + + self._init_gsam() + + entity = defect_plan.target_entity + synonym_map = { + "metal can package": ["TO-can", "metal can", "can body", "package body", "metal ring"], + "lens cap": ["glass lens cap", "lens", "glass dome", "optical window"], + "electrode bars": ["vertical bars", "electrodes", "metal lines"], + "die surface": ["die", "chip die", "semiconductor die", "IC die", "black rectangle", "dark square", "central chip"], + "solder joint": ["solder ball", "BGA ball", "solder bump", "joint"], + } + search_terms = [entity] + synonym_map.get(entity.lower(), []) + + predictions = [] + for term in search_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[defect_plan.target_subentity] if defect_plan.target_subentity else [], + entity_subentity_mapping={}, + min_area_ratio=0.005, + max_area_ratio=0.5, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + if term != entity: + print(f"[Agent] Fallback: detected '{term}' instead of '{entity}'") + break + + # Fallback 1: If entity is tiny (solder balls), ask VLM to detect the GROUP instead of one-by-one + if not predictions and any(k in entity.lower() for k in ['solder', 'ball', 'bump', 'joint']): + group_terms = [ + "array of solder balls", + "BGA ball grid", + "group of solder joints", + "solder ball array", + "BGA array", + "ball grid array" + ] + for term in group_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[], + entity_subentity_mapping={}, + min_area_ratio=0.05, # Group is large enough + max_area_ratio=0.8, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + print(f"[Agent] Group detection: found '{term}' with {len(preds)} prediction(s)") + break + + if not predictions: + print(f"[Agent] Warning: No detections for {entity}; using full image") + h, w = image.shape[:2] + best_pred = { + 'bbox': [0, 0, w, h], + 'pred_mask': torch.ones((h, w), dtype=torch.bool) + } + else: + best_pred = max(predictions, key=lambda p: p.get('area_ratio', 0)) + bbox = best_pred['bbox'] + h, w = image.shape[:2] + + cx = (bbox[0] + bbox[2]) / 2 + cy = (bbox[1] + bbox[3]) / 2 + if cx < 0.1 * w or cx > 0.9 * w or cy < 0.1 * h or cy > 0.9 * h: + print(f"[Agent] Warning: Bbox {bbox} on edge. Falling back to image center.") + best_pred['bbox'] = [int(0.25 * w), int(0.25 * h), int(0.75 * w), int(0.75 * h)] + + return best_pred + + # ------------------------------------------------------------------ + # Step 3: Prepare DefectFill generation conditions + # ------------------------------------------------------------------ + + def prepare_generation_conditions(self, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 3: Preparing DefectFill generation conditions...") + + obj_cls = defect_plan.object_class + if not obj_cls: + raise ValueError("object_class is required for DefectFill prompt standardization.") + + prompt = f"A {obj_cls} with " + defect_class = defect_plan.defect_type + + print(f"[Agent] Standardized prompt: '{prompt}'") + print(f"[Agent] Defect type (checkpoint): {defect_class}") + print(f"[Agent] Object class (checkpoint): {obj_cls}") + + return { + 'prompt': prompt, + 'defect_type': defect_class, + 'object_class': obj_cls, + 'defect_plan': defect_plan + } + + # ------------------------------------------------------------------ + # Step 4: Synthesize with DefectFill + # ------------------------------------------------------------------ + + def synthesize(self, image_patch: Image.Image, mask_patch: np.ndarray, + gen_conditions: Dict) -> Tuple[Image.Image, np.ndarray]: + print(f"\n{'='*60}") + print("[Agent] Step 4: Synthesizing defect with DefectFill...") + + obj_cls = gen_conditions['object_class'] + dfc_type = gen_conditions['defect_type'] + + self._init_defectfill(object_class=obj_cls, defect_type=dfc_type) + + # DefectFill expects PIL image + numpy mask (0-255) + inpainted = self.defectfill_generator.inpaint( + image=image_patch, + mask=mask_patch, + prompt=gen_conditions['prompt'], + defect_type=dfc_type, + seed=42 + ) + return inpainted, mask_patch + + # ------------------------------------------------------------------ + # Step 5: Verification + # ------------------------------------------------------------------ + + def verify(self, original_image: np.ndarray, generated_image: Image.Image, + defect_mask: np.ndarray, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 5: Verifying generated defect...") + + mask_bool = defect_mask.astype(bool) + if mask_bool.sum() == 0: + print("[Agent] Warning: Empty mask; verifying full image") + y1, x1 = 0, 0 + y2, x2 = original_image.shape[0], original_image.shape[1] + else: + ys, xs = np.where(mask_bool) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + margin = 32 + h, w = original_image.shape[:2] + y1 = max(0, y1 - margin) + x1 = max(0, x1 - margin) + y2 = min(h, y2 + margin) + x2 = min(w, x2 + margin) + + gen_crop = np.array(generated_image)[y1:y2, x1:x2] + orig_crop = original_image[y1:y2, x1:x2] + + obj_name = f"a {defect_plan.target_subentity or defect_plan.target_entity}" + artifact_type = "addition" + + result = artifact_description( + self.vlm_client, + original_image, + orig_crop, + gen_crop, + obj_name, + artifact_type, + self.money_manager + ) + + print(f"[Agent] Verification result: has_artifact={result.has_artifact}") + print(f"[Agent] Explanation: {result.explanation}") + print(f"[Agent] Label: {result.label}") + + return { + 'passed': result.has_artifact, + 'explanation': result.explanation, + 'label': result.label + } + + # ------------------------------------------------------------------ + # Main pipeline + # ------------------------------------------------------------------ + + def run(self, product_description: str, image_path: str, + caption: Optional[str] = None, max_defects: int = 3, + defect_type: Optional[str] = None, + object_class: Optional[str] = None) -> Dict: + """Run the full agentic pipeline with DefectFill.""" + start_time = datetime.now() + exp_id = str(uuid.uuid4())[:8] + + # ALWAYS resize to 512x512 โ€” DefectFill model requirement + orig_image = np.array(Image.open(image_path).convert('RGB')) + image = np.array(Image.fromarray(orig_image).resize( + (self.image_size, self.image_size), Image.LANCZOS + )) + print(f"[Agent] Loaded image: {orig_image.shape} -> working at {self.image_size}x{self.image_size}") + + plan = self.plan( + product_description, + image, + defect_type=defect_type, + object_class=object_class, + max_defects=max_defects + ) + + results = [] + defects_to_process = plan.possible_defects[:max_defects] + print(f"[Agent] Processing top {len(defects_to_process)} of {len(plan.possible_defects)} planned defects") + + for i, defect_plan in enumerate(defects_to_process): + print(f"\n{'='*80}") + print(f"[Agent] Defect {i+1}/{len(defects_to_process)}: [{defect_plan.defect_type.upper()}] {defect_plan.description}") + print(f"{'='*80}") + + try: + prediction = self.perceive(image, defect_plan) + target_bbox = prediction.get('bbox', [0, 0, image.shape[1], image.shape[0]]) + + gen_conditions = self.prepare_generation_conditions(defect_plan) + + # ========================================================================= + # Step 4: Smart Crop + DefectFill (inference.py style, always 512x512) + # ========================================================================= + x1, y1, x2, y2 = target_bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) + + # 1. Build full-image mask (0-255) with small defect region + h, w = image.shape[:2] + full_mask = np.zeros((h, w), dtype=np.uint8) + cx = (x1 + x2) // 2 + cy = (y1 + y2) // 2 + target_w = max(1, x2 - x1) + target_h = max(1, y2 - y1) + + vlm_ratio = getattr(defect_plan, 'defect_coverage_ratio', 0.25) + vlm_shape = getattr(defect_plan, 'mask_shape', 'free').lower().strip() + scale_factor = np.sqrt(vlm_ratio) + defect_w = max(8, int(target_w * scale_factor)) + defect_h = max(8, int(target_h * scale_factor)) + + if vlm_shape == 'circle': + r = max(4, min(defect_w, defect_h) // 2) + cv2.circle(full_mask, (cx, cy), r, 255, -1) + elif vlm_shape == 'square': + half = max(4, min(defect_w, defect_h) // 2) + cv2.rectangle(full_mask, (cx-half, cy-half), (cx+half, cy+half), 255, -1) + elif vlm_shape == 'rectangle': + cv2.rectangle(full_mask, (cx-defect_w//2, cy-defect_h//2), + (cx+defect_w//2, cy+defect_h//2), 255, -1) + else: # free/irregular + cv2.ellipse(full_mask, (cx, cy), + (max(4, defect_w//2), max(4, defect_h//2)), + 0, 0, 360, 255, -1) + + # 2. Smart crop around defect (exactly like inference.py) + y_idx, x_idx = np.where(full_mask > 0) + if len(y_idx) > 0: + min_y, max_y = np.min(y_idx), np.max(y_idx) + min_x, max_x = np.min(x_idx), np.max(x_idx) + cy_crop = (min_y + max_y) // 2 + cx_crop = (min_x + max_x) // 2 + max_dim = max(max_y - min_y, max_x - min_x) + else: + cy_crop, cx_crop = h // 2, w // 2 + max_dim = 0 + + padding = 50 + crop_size = max(self.image_size, max_dim + padding) + half = crop_size // 2 + + x1c = cx_crop - half + y1c = cy_crop - half + x2c = x1c + crop_size + y2c = y1c + crop_size + + if x1c < 0: x2c -= x1c; x1c = 0 + if y1c < 0: y2c -= y1c; y1c = 0 + if x2c > w: x1c -= (x2c - w); x2c = w + if y2c > h: y1c -= (y2c - h); y2c = h + x1c = max(0, x1c); y1c = max(0, y1c) + x2c = min(w, x2c); y2c = min(h, y2c) + + crop_img = image[y1c:y2c, x1c:x2c] + crop_mask = full_mask[y1c:y2c, x1c:x2c] + + if crop_img.shape[0] != self.image_size or crop_img.shape[1] != self.image_size: + crop_img = cv2.resize(crop_img, (self.image_size, self.image_size), + interpolation=cv2.INTER_AREA) + crop_mask = cv2.resize(crop_mask, (self.image_size, self.image_size), + interpolation=cv2.INTER_NEAREST) + + # 3. DefectFill on the crop + generated_image, _ = self.synthesize( + image_patch=Image.fromarray(crop_img), + mask_patch=crop_mask, + gen_conditions=gen_conditions + ) + gen_np = np.array(generated_image) + + # 4. Paste crop back into full 512x512 canvas + crop_h, crop_w = y2c - y1c, x2c - x1c + if gen_np.shape[:2] != (crop_h, crop_w): + gen_np = cv2.resize(gen_np, (crop_w, crop_h), interpolation=cv2.INTER_AREA) + mask_back = cv2.resize(crop_mask, (crop_w, crop_h), interpolation=cv2.INTER_NEAREST) + else: + mask_back = crop_mask + + blended_image = image.copy() + blended_image[y1c:y2c, x1c:x2c] = gen_np + + full_defect_mask = np.zeros((h, w), dtype=np.uint8) + full_defect_mask[y1c:y2c, x1c:x2c] = mask_back + + ys, xs = np.where(full_defect_mask > 0) + if len(ys) > 0 and len(xs) > 0: + mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] + else: + mask_bbox = target_bbox + + blended_with_green_box = create_visual_prompt_image(blended_image, mask_bbox) + + verification = self.verify(image, Image.fromarray(blended_image), full_defect_mask, defect_plan) + + if not verification['passed']: + defect_dir = self.output_dir / "failed" / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + else: + defect_dir = self.output_dir / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + + defect_dir.mkdir(parents=True, exist_ok=True) + + Image.fromarray(image).save(defect_dir / "real_clean_image.png") + Image.fromarray(blended_image).save(defect_dir / "blended_factory_defect.png") + generated_image.save(defect_dir / "raw_defectfill_patch.png") + + mask_img = Image.fromarray(full_defect_mask) + mask_img.save(defect_dir / "defect_mask.png") + + metadata = { + 'experiment_id': exp_id, + 'product_description': product_description, + 'product_type': plan.product_type, + 'object_class': defect_plan.object_class, + 'defect_type': defect_plan.defect_type, + 'defect_plan': defect_plan.dict() if hasattr(defect_plan, 'dict') else vars(defect_plan), + 'generation_conditions': {k: v for k, v in gen_conditions.items() if k != 'defect_plan'}, + 'verification': verification, + 'timestamp': datetime.now().isoformat() + } + with open(defect_dir / "metadata.json", 'w') as f: + json.dump(metadata, f, indent=2, default=str) + + if not verification['passed']: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': False, + 'verification_passed': False, + 'error': f"VLM verification failed: {verification.get('explanation', 'no explanation')}" + }) + print(f"[Agent] Verification FAILED for {defect_plan.defect_type}.") + else: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': True, + 'verification_passed': verification['passed'], + 'output_dir': str(defect_dir) + }) + print(f"[Agent] Defect {i+1} complete. Saved to {defect_dir}") + + except Exception as e: + print(f"[Agent] ERROR processing defect {i+1}: {str(e)}") + traceback.print_exc() + results.append({ + 'defect_type': getattr(defect_plan, 'defect_type', 'unknown'), + 'object_class': getattr(defect_plan, 'object_class', 'unknown'), + 'success': False, + 'error': str(e) + }) + + if hasattr(self, '_gsam_cache'): + keys_to_remove = [k for k in self._gsam_cache if k.startswith(str(image_path) + "::")] + for k in keys_to_remove: + self._gsam_cache.pop(k, None) + + elapsed = (datetime.now() - start_time).total_seconds() + print(f"\n{'='*60}") + print(f"[Agent] Pipeline complete in {elapsed:.1f}s") + print(f"[Agent] Results: {sum(1 for r in results if r['success'])}/{len(results)} succeeded") + + return { + 'experiment_id': exp_id, + 'product_type': plan.product_type, + 'results': results, + 'output_dir': str(self.output_dir), + 'elapsed_time': elapsed + } + + def cleanup(self): + if self.gsam_detector: + self.gsam_detector.cleanup() + self.gsam_detector = None + if self.defectfill_generator: + self.defectfill_generator.unload_models() + self.defectfill_generator = None + print("[Agent] All models cleaned up.") + + +# ============================================================================= +# CLI +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description='ArtiAgent โ€” DefectFill Edition (with VLM list selection)') + parser.add_argument('--product-desc', required=True, help='Product description (drives VLM selection)') + parser.add_argument('--image', required=True, help='Path to clean product image') + + # DefectFill checkpoint routing + parser.add_argument('--checkpoint-dir', required=True, + help='Root directory containing object_class/defect_type checkpoint subfolders') + parser.add_argument('--object-class', default=None, + help='Object class (optional; VLM selects from valid list if omitted)') + parser.add_argument('--defect-type', default=None, + help='Defect type (optional; VLM selects from valid list if omitted)') + + # Valid lists (auto-discovered from checkpoint-dir if not provided) + parser.add_argument('--valid-object-classes', default=None, + help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" ) + parser.add_argument('--valid-defect-types', default=None, + help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" ) + # Generation control + parser.add_argument('--output-dir', default='./defect_output', help='Output directory') + parser.add_argument('--caption', default=None, help='Optional image caption') + parser.add_argument('--max-defects', type=int, default=3, help='Max defects to generate') + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, help='Generation resolution') + parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps') + parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale') + + args = parser.parse_args() + + # Parse valid lists + valid_object_classes = None + valid_defect_types = None + if args.valid_object_classes: + valid_object_classes = json.loads(args.valid_object_classes) + if args.valid_defect_types: + valid_defect_types = json.loads(args.valid_defect_types) + + orchestrator = ArtiAgentOrchestrator( + device=args.device, + output_dir=args.output_dir, + vlm_model=args.vlm_model, + checkpoint_dir=args.checkpoint_dir, + object_class=args.object_class or "", + defect_type=args.defect_type or "", + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + image_size=args.image_size, + num_steps=args.num_steps, + guidance_scale=args.guidance_scale + ) + + result = orchestrator.run( + product_description=args.product_desc, + image_path=args.image, + caption=args.caption, + max_defects=args.max_defects, + defect_type=args.defect_type, + object_class=args.object_class + ) + + print(f"\nFinal output saved to: {result['output_dir']}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/artiagent_orchestrator.py b/ArtiAgent - DefectFill/src/artiagent_orchestrator.py new file mode 100644 index 0000000000000000000000000000000000000000..ac9cb7af6ceb5f89fa2cfd63a536a5e92522de3b --- /dev/null +++ b/ArtiAgent - DefectFill/src/artiagent_orchestrator.py @@ -0,0 +1,1297 @@ +""" +ArtiAgent Orchestrator โ€” DefectFill Edition (with VLM list selection) + +The VLM selects object_class from a given list and defect_type from the +corresponding per-object-class list. Product description drives the selection. + +Usage: + python artiagent_orchestrator.py \\ + --product-desc "VCSEL laser diode with glass lens cap" \\ + --image ./clean_chip.png \\ + --output-dir ./defect_output \\ + --checkpoint-dir "C:/.../checkpoints" \\ + --valid-object-classes '["xray_PCB","vcsel"]' \\ + --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \\ + --device cuda +""" + +import os +import sys +import json +import argparse +import uuid +import traceback +from pathlib import Path +from typing import Dict, List, Optional, Tuple +from datetime import datetime + +import numpy as np +import torch +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from pipeline.local_vlm_client import LocalVLMClient +from pipeline.prompts import ( + plan_defects_for_product, + artifact_description, + MoneyManager +) +from pipeline.gsam_detector import GSAMDetector +from pipeline.defectfill_generator import DefectFillGenerator, DefectFillConfig +from pipeline.defect_rag import get_rag +from pipeline.domain_router import get_router + +import cv2 + + +def blend_defect_onto_real_image( + real_image: np.ndarray, + defect_image: Image.Image, + defect_mask: np.ndarray, + target_bbox: List[int], + max_defect_ratio: Optional[float] = None, + mask_shape: str = "rectangle" +) -> Tuple[np.ndarray, np.ndarray]: + """Injects a DefectFill defect patch onto a real clean factory image.""" + defect_np = np.array(defect_image) + mask_uint8 = (defect_mask.astype(np.uint8) * 255) if defect_mask.dtype == bool else defect_mask.astype(np.uint8) + + ys, xs = np.where(mask_uint8 > 0) + if len(ys) == 0 or len(xs) == 0: + return real_image.copy(), np.zeros(real_image.shape[:2], dtype=np.uint8) + + y1_d, y2_d = ys.min(), ys.max() + x1_d, x2_d = xs.min(), xs.max() + + defect_patch = defect_np[y1_d:y2_d + 1, x1_d:x2_d + 1] + mask_patch = mask_uint8[y1_d:y2_d + 1, x1_d:x2_d + 1] + + x1_t, y1_t, x2_t, y2_t = target_bbox + target_w = max(1, x2_t - x1_t) + target_h = max(1, y2_t - y1_t) + + if max_defect_ratio is None or max_defect_ratio >= 1.0: + final_w = target_w + final_h = target_h + else: + if max_defect_ratio < 0.2: + max_defect_ratio = 0.2 + scale_factor = np.sqrt(max_defect_ratio) + scaled_w = int(target_w * scale_factor) + scaled_h = int(target_h * scale_factor) + + patch_h, patch_w = defect_patch.shape[:2] + aspect_ratio = patch_w / max(1, patch_h) + + if aspect_ratio > 1: + final_w = max(15, scaled_w) + final_h = max(15, int(final_w / aspect_ratio)) + else: + final_h = max(15, scaled_h) + final_w = max(15, int(final_h * aspect_ratio)) + + defect_patch_resized = cv2.resize(defect_patch, (final_w, final_h), interpolation=cv2.INTER_AREA) + mask_patch_resized = cv2.resize(mask_patch, (final_w, final_h), interpolation=cv2.INTER_NEAREST) + + shape_type = mask_shape.lower().strip() if mask_shape else "rectangle" + + if shape_type == "circle": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + center = (final_w // 2, final_h // 2) + radius = max(1, min(final_w, final_h) // 2 - 1) + cv2.circle(geom_mask, center, radius, 255, thickness=-1) + mask_patch_resized = geom_mask + + elif shape_type == "square": + geom_mask = np.zeros((final_h, final_w), dtype=np.uint8) + side = max(1, min(final_w, final_h) - 2) + top_left_x = (final_w - side) // 2 + top_left_y = (final_h - side) // 2 + cv2.rectangle( + geom_mask, + (top_left_x, top_left_y), + (top_left_x + side, top_left_y + side), + 255, + thickness=-1 + ) + mask_patch_resized = geom_mask + + elif shape_type == "rectangle": + mask_patch_resized = np.full((final_h, final_w), 255, dtype=np.uint8) + + center_x = x1_t + target_w // 2 + center_y = y1_t + target_h // 2 + center = (center_x, center_y) + + real_bgr = cv2.cvtColor(real_image, cv2.COLOR_RGB2BGR) + patch_bgr = cv2.cvtColor(defect_patch_resized, cv2.COLOR_RGB2BGR) + + patch_mean = np.mean(defect_patch_resized) + clone_mode = cv2.NORMAL_CLONE if patch_mean < 30 else cv2.MIXED_CLONE + + blended_bgr = cv2.seamlessClone( + patch_bgr, + real_bgr, + mask_patch_resized, + center, + clone_mode + ) + blended_rgb = cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB) + + full_mask = np.zeros(real_image.shape[:2], dtype=np.uint8) + top_left_x = max(0, center_x - final_w // 2) + top_left_y = max(0, center_y - final_h // 2) + + h_end = min(real_image.shape[0], top_left_y + final_h) + w_end = min(real_image.shape[1], top_left_x + final_w) + + mask_crop_h = h_end - top_left_y + mask_crop_w = w_end - top_left_x + + if mask_crop_h > 0 and mask_crop_w > 0: + full_mask[top_left_y:h_end, top_left_x:w_end] = ( + mask_patch_resized[:mask_crop_h, :mask_crop_w] > 128 + ).astype(np.uint8) + + return blended_rgb, full_mask + + +def create_visual_prompt_image(full_image: np.ndarray, bbox: list) -> np.ndarray: + """Draws a bright neon bounding box on the full image around the target ROI.""" + viz_img = full_image.copy() + x1, y1, x2, y2 = bbox + cv2.rectangle(viz_img, (x1, y1), (x2, y2), (0, 255, 0), thickness=2) + return viz_img + + +def resolve_checkpoint_path(checkpoint_dir: str, object_class: str, defect_type: str) -> str: + """Resolve DefectFill checkpoint path from object_class + defect_type.""" + path = Path(checkpoint_dir) / object_class / defect_type / "checkpoints" / "checkpoint_final.pt" + if not path.exists(): + alt = Path(checkpoint_dir) / object_class / defect_type / "checkpoint_final.pt" + if alt.exists(): + return str(alt) + raise FileNotFoundError( + f"Checkpoint not found for object_class='{object_class}', defect_type='{defect_type}'.\n" + f"Tried: {path}\nAlso tried: {alt}" + ) + return str(path) + + +class ArtiAgentOrchestrator: + """Agentic orchestrator for directed defect generation with DefectFill.""" + + def __init__( + self, + device='cuda', + output_dir='./defect_output', + vlm_model='gemma3:12b', + checkpoint_dir: str = "", + object_class: str = "", + defect_type: str = "", + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + image_size: int = 512, + num_steps: int = 50, + guidance_scale: float = 7.0, + domain_hint: str = "" + ): + self.device = device + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.vlm_client = LocalVLMClient(model=vlm_model) + self.money_manager = MoneyManager(model="gpt-4o") + + self.gsam_detector = None + self.defectfill_generator = None + + self.checkpoint_dir = checkpoint_dir + self.object_class = object_class + self.defect_type = defect_type + self.valid_object_classes = valid_object_classes or [] + self.valid_defect_types = valid_defect_types or {} + self.image_size = image_size + self.num_steps = num_steps + self.guidance_scale = guidance_scale + self.domain_hint = domain_hint + + self.rag = get_rag() + self.router = get_router() + + # ------------------------------------------------------------------ + # Lazy initializers + # ------------------------------------------------------------------ + + def _init_gsam(self): + if self.gsam_detector is None or getattr(self.gsam_detector, 'sam_predictor', None) is None: + print("[Agent] Initializing GSAM detector...") + self.gsam_detector = GSAMDetector( + device=self.device, + openai_client=self.vlm_client + ) + + def _init_defectfill(self, object_class: Optional[str] = None, defect_type: Optional[str] = None): + obj_cls = object_class or self.object_class + dfc_type = defect_type or self.defect_type + + if not obj_cls or not dfc_type: + raise ValueError("Both object_class and defect_type must be provided.") + + ckpt_path = resolve_checkpoint_path(self.checkpoint_dir, obj_cls, dfc_type) + + if self.defectfill_generator is None: + print(f"[Agent] Initializing DefectFill: {ckpt_path}") + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + else: + current_ckpt = getattr(self.defectfill_generator, 'ckpt_path', None) + if current_ckpt != ckpt_path: + print(f"[Agent] Switching checkpoint: {ckpt_path}") + self.defectfill_generator.unload_models() + config = DefectFillConfig( + ckpt_path=ckpt_path, + image_size=self.image_size, + num_steps=self.num_steps, + device=self.device, + guidance_scale=self.guidance_scale + ) + self.defectfill_generator = DefectFillGenerator(config) + else: + print("[Agent] Reusing DefectFill generator.") + + # ------------------------------------------------------------------ + # Step 1: Planning (with list-constrained selection) + # ------------------------------------------------------------------ + + def plan(self, product_description: str, image: np.ndarray, + defect_type: Optional[str] = None, object_class: Optional[str] = None, + num_defects: int = 3, domain_hint: Optional[str] = None): + """Agent plans defects. VLM selects object_class/defect_type from valid lists if not provided.""" + print(f"\n{'='*60}") + print("[Agent] Step 1: Planning defects from product description...") + print(f"Product: {product_description}") + if object_class: + print(f"[Agent] User-provided object_class: {object_class}") + else: + print(f"[Agent] No object_class provided โ€” VLM will select from: {self.valid_object_classes}") + if defect_type: + print(f"[Agent] User-provided defect_type: {defect_type}") + else: + print(f"[Agent] No defect_type provided โ€” VLM will select from valid list.") + + plan = plan_defects_for_product( + self.vlm_client, + product_description, + image, + money_manager=self.money_manager, + target_defect_type=defect_type, + object_class=object_class, + valid_object_classes=self.valid_object_classes if self.valid_object_classes else None, + valid_defect_types=self.valid_defect_types if self.valid_defect_types else None, + num_defects=num_defects, + domain_hint=domain_hint or self.domain_hint + ) + + if plan is None: + raise RuntimeError("Defect planning failed") + + print(f"[Agent] Product type: {plan.product_type}") + print(f"[Agent] Analysis: {plan.analysis}") + print(f"[Agent] Proposed {len(plan.possible_defects)} defects:") + for i, d in enumerate(plan.possible_defects, 1): + print(f" {i}. [{d.defect_type.upper()}] {d.description}") + print(f" Object class: {d.object_class}") + print(f" Target: {d.target_entity} / {d.target_subentity or '(whole)'}") + print(f" Location: {d.location_hint}") + print(f" Coverage: {d.defect_coverage_ratio} | Shape: {d.mask_shape}") + + return plan + + # ------------------------------------------------------------------ + # Step 2: Perception + # ------------------------------------------------------------------ + + def perceive(self, image: np.ndarray, defect_plan): + print(f"\n{'='*60}") + print("[Agent] Step 2: Directed perception (verification bbox)...") + + self._init_gsam() + + entity = defect_plan.target_entity + + is_solder_target = any(k in entity.lower() for k in ['solder', 'ball', 'bump', 'joint', 'bga', 'die']) + is_leg_target = any(k in entity.lower() for k in ['leg', 'pin', 'lead']) + + # ------------------------------------------------------------------ + # FAST PATH: For solder-ball arrays, use OpenCV blob detection directly + # ------------------------------------------------------------------ + if is_solder_target: + print(f"[Agent] Solder-ball target detected. Trying blob-detection fast path...") + preds, _, viz = self.gsam_detector.detect_feature_array( + image, + feature_type="dots", + blob_color=0, + min_area=5, + max_area=300, + min_circularity=0.5, + min_inertia_ratio=0.1, + pad_x=60, + pad_y=40, + max_y_span=90, + min_cluster_size=4, + entity_name="solder_ball_array" + ) + if len(preds) > 0: + pred = preds[0] + print(f"[Agent] Blob detection succeeded: bbox={preds[0]['bbox']}, " + f"area_ratio={pred.get('area_ratio', 0):.4f}") + return pred + + # ------------------------------------------------------------------ + # FAST PATH: For leg / line , use OpenCV blob detection directly + # ------------------------------------------------------------------ + if is_leg_target: + print(f"[Agent] Leg/Pin target detected. Trying blob-detection fast path...") + preds, _, viz = self.gsam_detector.detect_feature_array( + image, + feature_type="lines", + blob_color=0, + min_area_ratio=0.0005, # ~130 px at 512x512 + max_area_ratio=0.2, # ~39,000 px at 512x512 + min_circularity=0.01, # lines are NOT circular + min_inertia_ratio=0.4, # lines ARE elongated + pad_x=40, # tight horizontal padding + pad_y=20, + max_y_span=200, # allow full height span + min_cluster_size=2, # need at least 2 legs + aspect_ratio_range=(1.5, 100.0), # tall and thin + vertical_align_threshold=0.15, # x-gap tolerance + entity_name="triac_legs", + use_adaptive_threshold=True, # NEW: Use adaptive threshold + morph_kernel_size=3 # NEW: Connect broken pixel lines + ) + if len(preds) > 0: + pred = preds[0] + + # POST-PROCESS SAFETY: Force bbox to cover the bottom edge (pin tips) + # bbox = pred['bbox'] + # h, w = image.shape[:2] + # if bbox[3] < 0.85 * h: # If detected lines didn't reach the bottom + # print(f"[Agent] Correcting leg bbox: extending to bottom and narrowing width.") + # new_y2 = h + # new_width = bbox[2] - bbox[0] + # # If the detected width is too large (plastic body), narrow it to pin-width + # if new_width > w * 0.3: + # new_width = int(w * 0.15) + # center_x = (bbox[0] + bbox[2]) // 2 + # new_x1 = max(0, center_x - new_width // 2) + # new_x2 = min(w, center_x + new_width // 2) + # pred['bbox'] = [new_x1, bbox[1], new_x2, new_y2] + # pred['area_ratio'] = (new_x2 - new_x1) * (new_y2 - bbox[1]) / (h * w) + + print(f"[Agent] Blob detection succeeded: bbox={preds[0]['bbox']}, " + f"area_ratio={pred.get('area_ratio', 0):.4f}") + return pred + print(f"[Agent] Line detection failed, falling back to VLM...") + + # ------------------------------------------------------------------ + # FAST PATH: For screws / holes, use OpenCV blob detection directly + # ------------------------------------------------------------------ + # Check target_entity, target_subentity, description, AND defect_type + is_screw_target = ( + any(k in entity.lower() for k in ['screw', 'hole', 'fastener', 'mounting']) or + any(k in (defect_plan.target_subentity or '').lower() for k in ['screw', 'hole', 'fastener', 'mounting']) or + any(k in defect_plan.description.lower() for k in ['screw', 'hole', 'fastener', 'mounting']) or + 'screw' in defect_plan.defect_type.lower() # Catches "extra_screw", "missing_screw" + ) + if is_screw_target: + print(f"[Agent] Screw/hole target detected. Trying blob-detection fast path...") + # blob_color=255 finds bright screw heads, blob_color=0 finds dark empty holes + # - extra_screw: Need an EMPTY HOLE (dark) to place the new screw into. + # - missing_screw: Need an EXISTING SCREW (bright) to remove it from. + if defect_plan.defect_type == "extra_screw": + blob_color = 0 # Find dark empty holes + elif defect_plan.defect_type == "missing_screw": + blob_color = 255 # Find bright screw heads + else: + blob_color = 0 if "hole" in entity.lower() else 255 + preds, _, viz = self.gsam_detector.detect_feature_array( + image, + feature_type="single_dot", + blob_color=blob_color, # 0 for holes, 255 for heads + min_area=10, + max_area=200, + min_circularity=0.5, + min_inertia_ratio=0.1, + pad_x=15, + pad_y=15, + max_y_span=30, + min_cluster_size=1, + entity_name="screw_target", + + location_hint=defect_plan.location_hint # Pass the hint so it knows which screw to pick! + ) + if len(preds) > 0: + best_pred = preds[0] + print(f"[Agent] Blob detection succeeded: bbox={best_pred['bbox']}") + return best_pred + + # ------------------------------------------------------------------ + # FALLBACK: VLM + SAM (with relaxed area threshold for small parts) + # ------------------------------------------------------------------ + synonym_map = { + "metal can package": ["TO-can", "metal can", "can body", "package body", "metal ring"], + "lens cap": ["glass lens cap", "lens", "glass dome", "optical window"], + "electrode bars": ["vertical bars", "electrodes", "metal lines"], + "die surface": ["die", "chip die", "semiconductor die", "IC die", "black rectangle", "dark square", "central chip"], + "solder joint": ["solder ball", "BGA ball", "solder bump", "joint"], + "3 pin": ["three metal legs", "vertical metal pins", "metal leads", "power pins"], + "3-pin": ["three metal legs", "vertical metal pins", "metal leads", "power pins"], + "pin": ["metal leg", "lead", "component terminal"], + } + search_terms = [entity] + synonym_map.get(entity.lower(), []) + + predictions = [] + for term in search_terms: + # Lower threshold for small features like holes + is_small_feature = any(k in term.lower() for k in ['hole', 'screw', 'dot', 'pin', 'via']) + min_ratio = 0.0001 if is_small_feature else 0.005 + + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[defect_plan.target_subentity] if defect_plan.target_subentity else [], + entity_subentity_mapping={}, + location_hint=defect_plan.location_hint, # <--- ADD THIS LINE + min_area_ratio=min_ratio, + max_area_ratio=0.5, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + if term != entity: + print(f"[Agent] Fallback: detected '{term}' instead of '{entity}'") + break + + # Fallback: group detection for solder balls + if not predictions and any(k in entity.lower() for k in ['solder', 'ball', 'bump', 'joint']): + group_terms = [ + "array of solder balls", + "BGA ball grid", + "group of solder joints", + "solder ball array", + "BGA array", + "ball grid array" + ] + for term in group_terms: + preds, _, viz = self.gsam_detector.detect_parts( + image=image, + entities=[term], + subentities=[], + entity_subentity_mapping={}, + location_hint=defect_plan.location_hint, # <--- ADD THIS LINE + min_area_ratio=0.05, # Group is large enough + max_area_ratio=0.8, + openai_client=self.vlm_client + ) + if len(preds) > 0: + predictions = preds + print(f"[Agent] Group detection: found '{term}' with {len(preds)} prediction(s)") + break + + h, w = image.shape[:2] + + if not predictions: + print(f"[Agent] Warning: No detections for {entity}; using full image") + + best_pred = { + 'bbox': [0, 0, w, h], + 'pred_mask': torch.ones((h, w), dtype=torch.bool) + } + else: + # ------------------------------------------------------------------ + # If multiple leg/pin detections, MERGE them into one group bbox + # ------------------------------------------------------------------ + if is_leg_target and len(predictions) >= 2: + print(f"[Agent] Merging {len(predictions)} leg detections into group bbox...") + all_x1 = [p['bbox'][0] for p in predictions] + all_y1 = [p['bbox'][1] for p in predictions] + all_x2 = [p['bbox'][2] for p in predictions] + all_y2 = [p['bbox'][3] for p in predictions] + + pad_x = 15 + pad_y = 10 + + x1 = max(0, min(all_x1) - pad_x) + x2 = min(w, max(all_x2) + pad_x) + + # Position-agnostic bounds: depend ONLY on detected pin coordinates, not image height 'h' + y1 = max(0, min(all_y1) - pad_y) + y2 = min(h, max(all_y2) + pad_y) + + merged_bbox = [x1, y1, x2, y2] + + best_pred = { + 'bbox': merged_bbox, + 'pred_box': torch.tensor(merged_bbox).float(), + 'pred_mask': torch.ones((h, w), dtype=torch.bool), + 'area_ratio': (x2 - x1) * (y2 - y1) / (h * w), + 'entity': entity + } + print(f"[Agent] Merged leg group bbox: {merged_bbox}") + + else: + best_pred = max(predictions, key=lambda p: p.get('area_ratio', 0)) + bbox = best_pred['bbox'] + + # Edge / size sanity checks + cx = (bbox[0] + bbox[2]) / 2 + cy = (bbox[1] + bbox[3]) / 2 + bbox_area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) + total_area = h * w + + on_extreme_edge = (cx < 0.02 * w or cx > 0.98 * w or + cy < 0.02 * h or cy > 0.98 * h) + too_small = bbox_area < total_area * 0.001 + too_large = bbox_area > total_area * 0.35 + + if (on_extreme_edge and too_small) or too_large: + print(f"[Agent] Warning: Bad bbox {bbox}. Using center fallback.") + best_pred['bbox'] = [int(0.30 * w), int(0.30 * h), + int(0.70 * w), int(0.70 * h)] + + # General edge check + bbox = best_pred['bbox'] + cx = (bbox[0] + bbox[2]) / 2 + cy = (bbox[1] + bbox[3]) / 2 + if cx < 0.05 * w or cx > 0.95 * w or cy < 0.05 * h or cy > 0.95 * h: + print(f"[Agent] Warning: Bbox {bbox} on extreme edge. Falling back to center.") + best_pred['bbox'] = [int(0.20 * w), int(0.20 * h), + int(0.80 * w), int(0.80 * h)] + + return best_pred + + # ------------------------------------------------------------------ + # Step 3: Prepare DefectFill generation conditions + # ------------------------------------------------------------------ + + def prepare_generation_conditions(self, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 3: Preparing DefectFill generation conditions...") + + obj_cls = defect_plan.object_class + if not obj_cls: + raise ValueError("object_class is required for DefectFill prompt standardization.") + + defect_class = defect_plan.defect_type + + # Build a material-aware prompt + material_hint = "" + if defect_class == "missing_screw": + material_hint = "by only dark empty threaded hole, black shadowed interior, no screw" + elif defect_class == "extra_screw": + material_hint = "shiny silver Phillips head screw in empty hole, metallic cross slot, bright reflection" + + prompt = f"A {obj_cls} with {defect_class} {material_hint}" + # prompt = f"A {obj_cls} with " + + print(f"[Agent] Standardized prompt: '{prompt}'") + print(f"[Agent] Defect type (checkpoint): {defect_class}") + print(f"[Agent] Object class (checkpoint): {obj_cls}") + + return { + 'prompt': prompt, + 'defect_type': defect_class, + 'object_class': obj_cls, + 'defect_plan': defect_plan + } + + # ------------------------------------------------------------------ + # Step 4: Synthesize with DefectFill + # ------------------------------------------------------------------ + + def synthesize(self, image_patch: Image.Image, mask_patch: np.ndarray, + gen_conditions: Dict) -> Tuple[Image.Image, np.ndarray]: + print(f"\n{'='*60}") + print("[Agent] Step 4: Synthesizing defect with DefectFill...") + + obj_cls = gen_conditions['object_class'] + dfc_type = gen_conditions['defect_type'] + + self._init_defectfill(object_class=obj_cls, defect_type=dfc_type) + + # DefectFill expects PIL image + numpy mask (0-255) + inpainted = self.defectfill_generator.inpaint( + image=image_patch, + mask=mask_patch, + prompt=gen_conditions['prompt'], + defect_type=dfc_type, + seed=42 + ) + + return inpainted, mask_patch + + # ------------------------------------------------------------------ + # Step 5: Verification + # ------------------------------------------------------------------ + + def verify(self, original_image: np.ndarray, generated_image: Image.Image, + defect_mask: np.ndarray, defect_plan) -> Dict: + print(f"\n{'='*60}") + print("[Agent] Step 5: Verifying generated defect...") + + mask_bool = defect_mask.astype(bool) + if mask_bool.sum() == 0: + print("[Agent] Warning: Empty mask; verifying full image") + y1, x1 = 0, 0 + y2, x2 = original_image.shape[0], original_image.shape[1] + else: + ys, xs = np.where(mask_bool) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + margin = 64 + h, w = original_image.shape[:2] + y1 = max(0, y1 - margin) + x1 = max(0, x1 - margin) + y2 = min(h, y2 + margin) + x2 = min(w, x2 + margin) + + gen_crop = np.array(generated_image)[y1:y2, x1:x2] + orig_crop = original_image[y1:y2, x1:x2] + + obj_name = f"a {defect_plan.target_subentity or defect_plan.target_entity}" + artifact_type = "addition" + + result = artifact_description( + self.vlm_client, + original_image, + orig_crop, + gen_crop, + obj_name, + artifact_type, + self.money_manager + ) + + print(f"[Agent] Verification result: has_artifact={result.has_artifact}") + print(f"[Agent] Explanation: {result.explanation}") + print(f"[Agent] Label: {result.label}") + + return { + 'passed': result.has_artifact, + 'explanation': result.explanation, + 'label': result.label + } + + # ------------------------------------------------------------------ + # Main pipeline + # ------------------------------------------------------------------ + + def run(self, product_description: str, image_path: str, + caption: Optional[str] = None, num_defects: int = 3, + defect_type: Optional[str] = None, + object_class: Optional[str] = None) -> Dict: + """Run the full agentic pipeline with DefectFill.""" + start_time = datetime.now() + exp_id = str(uuid.uuid4())[:8] + + # ALWAYS resize to 512x512 โ€” DefectFill model requirement + orig_image = np.array(Image.open(image_path).convert('RGB')) + image = np.array(Image.fromarray(orig_image).resize( + (self.image_size, self.image_size), Image.LANCZOS + )) + print(f"[Agent] Loaded image: {orig_image.shape} -> working at {self.image_size}x{self.image_size}") + + plan = self.plan( + product_description, + image, + defect_type=defect_type, + object_class=object_class, + num_defects=num_defects + ) + + results = [] + defects_to_process = plan.possible_defects[:num_defects] + print(f"[Agent] Processing top {len(defects_to_process)} of {len(plan.possible_defects)} planned defects") + + for i, defect_plan in enumerate(defects_to_process): + print(f"\n{'='*80}") + print(f"[Agent] Defect {i+1}/{len(defects_to_process)}: [{defect_plan.defect_type.upper()}] {defect_plan.description}") + print(f"{'='*80}") + + try: + prediction = self.perceive(image, defect_plan) + + # ========================================================================= + # PATCH: Force location to the 16-dot area for xray_PCB type defects + # ========================================================================= + if defect_plan.object_class == "xray_PCB": + self._init_gsam() # Ensure the detector is loaded + solder_preds, _, _ = self.gsam_detector.detect_feature_array(image) + if solder_preds and len(solder_preds) > 0: + print(f"[Agent] Forcing xray_PCB type defect to target the 16-dot solder ball array.") + prediction = solder_preds[0] # Override the prediction + else: + print(f"[Agent] Warning: Could not detect 16-dot array. Using original detection.") + # ========================================================================= + + target_bbox = prediction.get('bbox', [0, 0, image.shape[1], image.shape[0]]) + + # ============================================================================= + # COMPOSITING FAST-PATH: extra_screw โ€” copy a real screw instead of generating + # ============================================================================= + skip_synthesis = False + if defect_plan.defect_type == "extra_screw": + self._init_gsam() + + # Find all bright screw heads in the image (donor candidates) + donor_preds, _, _ = self.gsam_detector.detect_feature_array( + image, + feature_type="single_dot", + blob_color=255, # Bright metal screws + min_area=20, + max_area=250, + min_circularity=0.5, + min_inertia_ratio=0.1, + pad_x=20, + pad_y=20, + min_cluster_size=1, + entity_name="donor_screw" + ) + + tx1, ty1, tx2, ty2 = target_bbox + tw = max(1, tx2 - tx1) + th = max(1, ty2 - ty1) + target_area = tw * th + + # Pick the first donor that does NOT overlap with the target hole + donor_bbox = None + best_donor_area = 0 + for d in donor_preds: + dx1, dy1, dx2, dy2 = d['bbox'] + # Overlap check: if donor is far from target, use it + if not (dx2 < tx1 or dx1 > tx2 or dy2 < ty1 or dy1 > ty2): + continue + area = (dx2 - dx1) * (dy2 - dy1) + # Donor must be at least 70% of target size so it doesn't stretch into mush + + if area > best_donor_area and area >= target_area * 0.7: + best_donor_area = area + donor_bbox = [dx1, dy1, dx2, dy2] + + if donor_bbox is None: + print("[Agent] No suitable donor screw found. Falling back to diffusion.") + skip_synthesis = False + else: + dx1, dy1, dx2, dy2 = donor_bbox + + # Extract donor screw patch + donor_patch = Image.fromarray(image[dy1:dy2, dx1:dx2]) + donor_mask = np.ones((dy2 - dy1, dx2 - dx1), dtype=np.uint8) * 255 + + # A real screw head is ~1.4ร— larger than the hole opening. + # Expand the target bbox so the head naturally overhangs. + head_scale = 1.4 + cx = (tx1 + tx2) // 2 + cy = (ty1 + ty2) // 2 + half_w = int((tx2 - tx1) * head_scale / 2) + half_h = int((ty2 - ty1) * head_scale / 2) + expanded_bbox = [ + max(0, cx - half_w), + max(0, cy - half_h), + min(image.shape[1], cx + half_w), + min(image.shape[0], cy + half_h) + ] + + # Blend using your existing Poisson blending function + blended_rgb, full_mask = blend_defect_onto_real_image( + real_image=image, + defect_image=donor_patch, + defect_mask=donor_mask, + target_bbox=expanded_bbox, + max_defect_ratio=1.0, + mask_shape="circle" + ) + + blended_image = blended_rgb + blended_with_green_box = create_visual_prompt_image(blended_image, target_bbox) + full_defect_mask = full_mask + generated_image = Image.fromarray(blended_image) + + skip_synthesis = True + print(f"[Agent] extra_screw: Used copy-paste compositing from donor screw at {donor_bbox}") + + verification = self.verify(image, Image.fromarray(blended_image), full_defect_mask, defect_plan) + + # ============================================================================= + # END COMPOSITING FAST-PATH + # ============================================================================= + + if not skip_synthesis: + # --- PUT THE EXISTING SYNTHESIS BLOCK HERE --- + # (Everything from gen_conditions = self.prepare_generation_conditions(...) + # down to the seamlessClone blending goes inside this if-block) + + gen_conditions = self.prepare_generation_conditions(defect_plan) + + # ========================================================================= + # Step 4: Smart Crop + DefectFill (inference.py style, always 512x512) + # ========================================================================= + x1, y1, x2, y2 = target_bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2) + + # 1. Build full-image mask respecting defect plan shape and coverage ratio + h, w = image.shape[:2] + full_mask = np.zeros((h, w), dtype=np.uint8) + + # mask_shape = defect_plan.mask_shape + mask_shape = "rectangle" + coverage_ratio = defect_plan.defect_coverage_ratio + + # Calculate scaled mask dimensions based on coverage ratio + bbox_w = x2 - x1 + bbox_h = y2 - y1 + + # SAFETY: If detection failed and bbox is the full image, shrink it + if bbox_w > w * 0.9 and bbox_h > h * 0.9: + print("[Agent] Warning: Full-image bbox detected. Using centered quarter region.") + bbox_w = int(w * 0.5) + bbox_h = int(h * 0.5) + x1 = (w - bbox_w) // 2 + y1 = (h - bbox_h) // 2 + x2 = x1 + bbox_w + y2 = y1 + bbox_h + + scale_factor = np.sqrt(max(0.05, coverage_ratio)) # Clamp to avoid zero + scaled_w = int(bbox_w * scale_factor) + scaled_h = int(bbox_h * scale_factor) + + center_x = (x1 + x2) // 2 + center_y = (y1 + y2) // 2 + + # ============================================================================= + # NEW: MASK SIZE OVERRIDE FOR SCREW/HOLE TARGETS + # ============================================================================= + # is_screw_target = any(k in defect_plan.target_entity.lower() for k in ['screw', 'hole']) + # if is_screw_target: + # # FORCE the mask to cover the ENTIRE detected hole/screw-head bbox + # scaled_w = bbox_w + # scaled_h = bbox_h + # # Enforce a minimum size so the generator has enough pixels to work with + # min_mask_size = 32 + # if scaled_w < min_mask_size: scaled_w = min_mask_size + # if scaled_h < min_mask_size: scaled_h = min_mask_size + # print(f"[Agent] Screw target detected. Mask resized to {scaled_w}x{scaled_h} (was {int(bbox_w*scale_factor)}x{int(bbox_h*scale_factor)})") + + # ========================================================================= + # NEW PATCH: Prevent mask width from covering multiple pins + # (Fixes reddish artifacts in fallback scenarios) + # ========================================================================= + # is_leg_target = any(k in defect_plan.target_entity.lower() for k in ['leg', 'pin', 'lead']) + + # if is_leg_target and bbox_w > 35: # bbox_w > 35 means it covers more than 1 pin + # # Check if the defect targets a specific side (left, middle, right) + # pin_keyword = "" + # if defect_plan.location_hint: + # pin_keyword = defect_plan.location_hint.lower() + # if not pin_keyword and defect_plan.target_subentity: + # pin_keyword = defect_plan.target_subentity.lower() + + # # Only override width if they specified a specific pin + # if any(k in pin_keyword for k in ['left', 'right', 'middle', 'center']): + # # Force the mask width to match a single metallic pin (~18px) + # single_pin_w = 18 + # scaled_w = single_pin_w + + # # Shift the horizontal center to the correct pin + # pin_offsets = { + # 'left': -int(bbox_w * 0.20), + # 'middle': 0, + # 'center': 0, + # 'right': int(bbox_w * 0.20) + # } + # for k, v in pin_offsets.items(): + # if k in pin_keyword: + # center_x = center_x + v + # break + # ========================================================================= + + if mask_shape == "circle": + radius = max(1, min(scaled_w, scaled_h) // 2 - 1) + cv2.circle(full_mask, (center_x, center_y), radius, 255, -1) + elif mask_shape == "square": + side = max(1, min(scaled_w, scaled_h) - 2) + top_left_x = center_x - side // 2 + top_left_y = center_y - side // 2 + cv2.rectangle(full_mask, (top_left_x, top_left_y), (top_left_x + side, top_left_y + side), 255, -1) + elif mask_shape == "rectangle": + top_left_x = center_x - scaled_w // 2 + top_left_y = center_y - scaled_h // 2 + # --- NEW PATCH --- + entity = defect_plan.target_entity + is_leg_target = any(k in entity.lower() for k in ['leg', 'pin', 'lead']) + if is_leg_target: # Use the variable defined earlier in perceive or calculate again + # Instead of centering on the bbox center (which is the plastic body), + # anchor the mask to the bottom of the bbox (where the legs are). + # Place the mask so its bottom touches the bbox's bottom. + top_left_y = y2 - scaled_h + top_left_y = max(y1, top_left_y) + else: + top_left_y = center_y - scaled_h // 2 + # ------------------- + cv2.rectangle(full_mask, (top_left_x, top_left_y), (top_left_x + scaled_w, top_left_y + scaled_h), 255, -1) + else: # "free" - use an ellipse centered in the die + axes = (max(1, scaled_w // 2), max(1, scaled_h // 2)) + cv2.ellipse(full_mask, (center_x, center_y), axes, 0, 0, 360, 255, -1) + + # 2. Smart crop around defect (exactly like inference.py) + y_idx, x_idx = np.where(full_mask > 0) + if len(y_idx) > 0: + min_y, max_y = np.min(y_idx), np.max(y_idx) + min_x, max_x = np.min(x_idx), np.max(x_idx) + cy_crop = (min_y + max_y) // 2 + cx_crop = (min_x + max_x) // 2 + max_dim = max(max_y - min_y, max_x - min_x) + else: + cy_crop, cx_crop = h // 2, w // 2 + max_dim = 0 + + padding = 50 + crop_size = max(self.image_size, max_dim + padding) + # CAP CROP SIZE to focus only on the local defect area (cures hallucinated duplicates) + # max_crop_size = 256 + # crop_size = max_dim + padding + # if crop_size < 128: + # crop_size = 128 + # if crop_size > max_crop_size: + # crop_size = max_crop_size + + half = crop_size // 2 + + x1c = cx_crop - half + y1c = cy_crop - half + x2c = x1c + crop_size + y2c = y1c + crop_size + + if x1c < 0: x2c -= x1c; x1c = 0 + if y1c < 0: y2c -= y1c; y1c = 0 + if x2c > w: x1c -= (x2c - w); x2c = w + if y2c > h: y1c -= (y2c - h); y2c = h + x1c = max(0, x1c); y1c = max(0, y1c) + x2c = min(w, x2c); y2c = min(h, y2c) + + crop_img = image[y1c:y2c, x1c:x2c] + crop_mask = full_mask[y1c:y2c, x1c:x2c] + + if crop_img.shape[0] != self.image_size or crop_img.shape[1] != self.image_size: + crop_img = cv2.resize(crop_img, (self.image_size, self.image_size), + interpolation=cv2.INTER_AREA) + crop_mask = cv2.resize(crop_mask, (self.image_size, self.image_size), + interpolation=cv2.INTER_NEAREST) + + # 3. DefectFill on the crop + generated_image, _ = self.synthesize( + image_patch=Image.fromarray(crop_img), + mask_patch=crop_mask, + gen_conditions=gen_conditions + ) + gen_np = np.array(generated_image) + + if defect_plan.defect_type == "missing_screw": + gen_np = (gen_np * 0.35).astype(np.uint8) # Crush to ~35% brightness + # # Darken the interior of the hole for more depth + # mask_bool = crop_mask.astype(bool) + # if mask_bool.sum() > 0: + # # Slightly darken the generated region + # gen_np[mask_bool] = (gen_np[mask_bool] * 0.6).astype(np.uint8) + + if defect_plan.defect_type == "extra_screw": + # Screws are bright โ€” force the patch to be light and contrasty + gen_np = np.clip(gen_np * 1.2 + 30, 0, 255).astype(np.uint8) + + # 4. Paste crop back โ€” SEAMLESS CLONE BLENDING (Poisson) + crop_h, crop_w = y2c - y1c, x2c - x1c + if gen_np.shape[:2] != (crop_h, crop_w): + gen_np = cv2.resize(gen_np, (crop_w, crop_h), interpolation=cv2.INTER_AREA) + mask_back = cv2.resize(crop_mask, (crop_w, crop_h), interpolation=cv2.INTER_NEAREST) + else: + mask_back = crop_mask + + # Get original crop for blending + orig_crop = image[y1c:y2c, x1c:x2c].copy() + + # Use Poisson blending (seamlessClone) to perfectly match gradients and remove halos + mask_back_uint8 = (mask_back > 0).astype(np.uint8) * 255 + + # Convert to BGR for OpenCV + orig_crop_bgr = cv2.cvtColor(orig_crop, cv2.COLOR_RGB2BGR) + gen_crop_bgr = cv2.cvtColor(gen_np, cv2.COLOR_RGB2BGR) + + # Determine clone mode based on average brightness of the defect patch + patch_mean = np.mean(gen_np) + if defect_plan.defect_type == "missing_screw" or defect_plan.defect_type == "extra_screw": + clone_mode = cv2.NORMAL_CLONE # Don't blend colors, just paste the dark hole + else: + clone_mode = cv2.NORMAL_CLONE if patch_mean < 30 else cv2.MIXED_CLONE + + center = (mask_back_uint8.shape[1] // 2, mask_back_uint8.shape[0] // 2) + blended_crop_bgr = cv2.seamlessClone( + gen_crop_bgr, orig_crop_bgr, mask_back_uint8, center, clone_mode + ) + blended_crop = cv2.cvtColor(blended_crop_bgr, cv2.COLOR_BGR2RGB) + + blended_image = image.copy() + blended_image[y1c:y2c, x1c:x2c] = blended_crop + + full_defect_mask = np.zeros((h, w), dtype=np.uint8) + full_defect_mask[y1c:y2c, x1c:x2c] = mask_back + + ys, xs = np.where(full_defect_mask > 0) + if len(ys) > 0 and len(xs) > 0: + mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] + else: + mask_bbox = target_bbox + + # ========================================================================= + # POST-PROCESS: Geometry warp for pin bending (Triac "double" defect) + # ========================================================================= + # Re-evaluate if we are targeting a leg + # is_leg_target = any(k in defect_plan.target_entity.lower() for k in ['leg', 'pin', 'lead']) + + # if is_leg_target: + # x1, y1, x2, y2 = mask_bbox + # pad = 20 + + # # Expand patch slightly to blend perfectly + # patch_y1, patch_y2 = max(0, y1-pad), min(h, y2+pad) + # patch_x1, patch_x2 = max(0, x1-pad), min(w, x2+pad) + + # # Extract the leg region + # leg_crop = blended_image[patch_y1:patch_y2, patch_x1:patch_x2] + + # # Determine which direction to pull the pin + # shift_x = 0 + # hint = "" + # if defect_plan.location_hint: + # hint = defect_plan.location_hint.lower() + # if not hint and defect_plan.target_subentity: + # hint = defect_plan.target_subentity.lower() + + # if 'right' in hint: + # shift_x = 18 # Pull right pin 18 pixels rightward + # elif 'left' in hint: + # shift_x = -18 # Pull left pin 18 pixels leftward + # elif 'double' in defect_plan.defect_type: + # # For double defect, usually left/right pins. Default to pulling left pin + # shift_x = -18 + + # # Define affine transformation to physically shift the pin pixels + # pts1 = np.float32([[0,0], [leg_crop.shape[1],0], [0,leg_crop.shape[0]]]) + # pts2 = np.float32([[shift_x,0], [leg_crop.shape[1]+shift_x,0], [shift_x,leg_crop.shape[0]]]) + + # M = cv2.getAffineTransform(pts1, pts2) + # warped_leg = cv2.warpAffine(leg_crop, M, (leg_crop.shape[1], leg_crop.shape[0])) + + # # Create mask for the warped region + # warp_mask = np.zeros(leg_crop.shape[:2], dtype=np.uint8) + # warp_mask[:, :] = 255 + + # # Blend warped leg back using seamlessClone + # blended_image_bgr = cv2.cvtColor(blended_image, cv2.COLOR_RGB2BGR) + # warped_leg_bgr = cv2.cvtColor(warped_leg, cv2.COLOR_RGB2BGR) + + # # Paste back into original spot with Poisson blending + # blended_crop_bgr = cv2.seamlessClone( + # warped_leg_bgr, + # blended_image_bgr[patch_y1:patch_y2, patch_x1:patch_x2], + # warp_mask, + # (leg_crop.shape[1]//2, leg_crop.shape[0]//2), + # cv2.NORMAL_CLONE + # ) + # blended_image_bgr[patch_y1:patch_y2, patch_x1:patch_x2] = blended_crop_bgr + # blended_image = cv2.cvtColor(blended_image_bgr, cv2.COLOR_BGR2RGB) + # ========================================================================= + + blended_with_green_box = create_visual_prompt_image(blended_image, mask_bbox) + + verification = self.verify(image, generated_image, full_defect_mask, defect_plan) + + if not verification['passed']: + defect_dir = self.output_dir / "failed" / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + else: + defect_dir = self.output_dir / f"{exp_id}_defect_{i}_{defect_plan.defect_type}" + + defect_dir.mkdir(parents=True, exist_ok=True) + + Image.fromarray(image).save(defect_dir / "real_clean_image.png") + + if skip_synthesis: + Image.fromarray(blended_image).save(defect_dir / "blended_factory_defect.png") + else: + generated_image.save(defect_dir / "blended_factory_defect.png") + + generated_image.save(defect_dir / "raw_defectfill_patch.png") + + mask_img = Image.fromarray((full_defect_mask > 0).astype(np.uint8) * 255) + mask_img.save(defect_dir / "defect_mask.png") + + if not skip_synthesis: + gen_conditions_to_save = {k: v for k, v in gen_conditions.items() if k != 'defect_plan'} + else: + gen_conditions_to_save = { + 'prompt': f"A {defect_plan.object_class} with {defect_plan.defect_type} (copy-paste compositing)", + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class + } + + metadata = { + 'experiment_id': exp_id, + 'product_description': product_description, + 'product_type': plan.product_type, + 'object_class': defect_plan.object_class, + 'defect_type': defect_plan.defect_type, + 'defect_plan': defect_plan.dict() if hasattr(defect_plan, 'dict') else vars(defect_plan), + 'generation_conditions': gen_conditions_to_save, + 'verification': verification, + 'timestamp': datetime.now().isoformat() + } + with open(defect_dir / "metadata.json", 'w') as f: + json.dump(metadata, f, indent=2, default=str) + + if not verification['passed']: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': False, + 'verification_passed': False, + 'error': f"VLM verification failed: {verification.get('explanation', 'no explanation')}" + }) + print(f"[Agent] Verification FAILED for {defect_plan.defect_type}.") + else: + results.append({ + 'defect_type': defect_plan.defect_type, + 'object_class': defect_plan.object_class, + 'success': True, + 'verification_passed': verification['passed'], + 'output_dir': str(defect_dir) + }) + print(f"[Agent] Defect {i+1} complete. Saved to {defect_dir}") + + except Exception as e: + print(f"[Agent] ERROR processing defect {i+1}: {str(e)}") + traceback.print_exc() + results.append({ + 'defect_type': getattr(defect_plan, 'defect_type', 'unknown'), + 'object_class': getattr(defect_plan, 'object_class', 'unknown'), + 'success': False, + 'error': str(e) + }) + + if hasattr(self, '_gsam_cache'): + keys_to_remove = [k for k in self._gsam_cache if k.startswith(str(image_path) + "::")] + for k in keys_to_remove: + self._gsam_cache.pop(k, None) + + elapsed = (datetime.now() - start_time).total_seconds() + print(f"\n{'='*60}") + print(f"[Agent] Pipeline complete in {elapsed:.1f}s") + print(f"[Agent] Results: {sum(1 for r in results if r['success'])}/{len(results)} succeeded") + + return { + 'experiment_id': exp_id, + 'product_type': plan.product_type, + 'results': results, + 'output_dir': str(self.output_dir), + 'elapsed_time': elapsed + } + + def cleanup(self): + if self.gsam_detector: + self.gsam_detector.cleanup() + self.gsam_detector = None + if self.defectfill_generator: + self.defectfill_generator.unload_models() + self.defectfill_generator = None + print("[Agent] All models cleaned up.") + + +# ============================================================================= +# CLI +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description='ArtiAgent โ€” DefectFill Edition (with VLM list selection)') + parser.add_argument('--product-desc', required=True, help='Product description (drives VLM selection)') + parser.add_argument('--image', required=True, help='Path to clean product image') + + # DefectFill checkpoint routing + parser.add_argument('--checkpoint-dir', required=True, + help='Root directory containing object_class/defect_type checkpoint subfolders') + parser.add_argument('--object-class', default=None, + help='Object class (optional; VLM selects from valid list if omitted)') + parser.add_argument('--defect-type', default=None, + help='Defect type (optional; VLM selects from valid list if omitted)') + + # Valid lists (auto-discovered from checkpoint-dir if not provided) + parser.add_argument('--valid-object-classes', default=None, + help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" ) + parser.add_argument('--valid-defect-types', default=None, + help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" ) + # Generation control + parser.add_argument('--output-dir', default='./defect_output', help='Output directory') + parser.add_argument('--caption', default=None, help='Optional image caption') + parser.add_argument('--max-defects', type=int, default=3, help='Max defects to generate') + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, help='Generation resolution') + parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps') + parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale') + + args = parser.parse_args() + + # Parse valid lists + valid_object_classes = None + valid_defect_types = None + if args.valid_object_classes: + valid_object_classes = json.loads(args.valid_object_classes) + if args.valid_defect_types: + valid_defect_types = json.loads(args.valid_defect_types) + + orchestrator = ArtiAgentOrchestrator( + device=args.device, + output_dir=args.output_dir, + vlm_model=args.vlm_model, + checkpoint_dir=args.checkpoint_dir, + object_class=args.object_class or "", + defect_type=args.defect_type or "", + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + image_size=args.image_size, + num_steps=args.num_steps, + guidance_scale=args.guidance_scale + ) + + result = orchestrator.run( + product_description=args.product_desc, + image_path=args.image, + caption=args.caption, + num_defects=args.num_defects, + defect_type=args.defect_type, + object_class=args.object_class + ) + + print(f"\nFinal output saved to: {result['output_dir']}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/batch_agent_orchestrator - Copy.py b/ArtiAgent - DefectFill/src/batch_agent_orchestrator - Copy.py new file mode 100644 index 0000000000000000000000000000000000000000..3e1f1a599822bfe9ebf533f8433639c4809484bd --- /dev/null +++ b/ArtiAgent - DefectFill/src/batch_agent_orchestrator - Copy.py @@ -0,0 +1,414 @@ +""" +Batch Agent Orchestrator โ€” DefectFill Edition (with VLM list selection) + +Auto-discovers valid object classes and defect types from the checkpoint directory. +VLM selects object_class from the discovered list based on product description. +If defect_type is not provided, VLM selects from the per-object-class list. + +Usage: + # Auto-discover checkpoints, VLM selects everything from product description + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage" \ + --output-dir "C:/AgentOutput" \ + --checkpoint-dir "C:/.../checkpoints" \ + --product-desc "VCSEL laser diode with emission aperture" \ + --device cuda + + # Provide explicit valid lists (overrides auto-discovery) + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage" \ + --output-dir "C:/AgentOutput" \ + --checkpoint-dir "C:/.../checkpoints" \ + --valid-object-classes '["xray_PCB","vcsel"]' \ + --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \ + --product-desc "VCSEL laser diode" \ + --device cuda + + # CSV manifest with per-image routing (overrides VLM selection for those images) + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage" \ + --output-dir "C:/AgentOutput" \ + --checkpoint-dir "C:/.../checkpoints" \ + --manifest "C:/products.csv" \ + --device cuda +""" + +import os +import sys +import csv +import json +import argparse +import traceback +from pathlib import Path +from datetime import datetime +from typing import Dict, List, Optional +from tqdm import tqdm + +import numpy as np +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from artiagent_orchestrator import ArtiAgentOrchestrator + + +def infer_product_from_path(image_path: Path) -> str: + """Infer product description from folder structure or filename.""" + parent = image_path.parent.name.lower() + if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']: + return parent.replace('_', ' ').replace('-', ' ') + stem = image_path.stem.lower() + for keyword in ['vcsel', 'lens', 'die', 'photodiode', 'sensor', 'chip', 'led', 'laser', 'optical']: + if keyword in stem: + return keyword + return "electronic component" + + +def infer_object_class_from_path(image_path: Path) -> str: + """Infer object_class from parent folder name.""" + parent = image_path.parent.name.lower() + if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']: + return parent + return "" + + +def discover_checkpoints(checkpoint_dir: str) -> Dict[str, List[str]]: + """ + Scan checkpoint directory to discover available object_class -> defect_type mappings. + Expected structure: + checkpoint_dir/ + xray_PCB/ + xray_die/ + checkpoints/checkpoint_final.pt + bubble/ + checkpoints/checkpoint_final.pt + """ + cp = Path(checkpoint_dir) + mapping = {} + if not cp.exists(): + print(f"[Batch] Warning: checkpoint-dir does not exist: {checkpoint_dir}") + return mapping + + for obj_dir in cp.iterdir(): + if not obj_dir.is_dir(): + continue + defect_types = [] + for defect_dir in obj_dir.iterdir(): + if not defect_dir.is_dir(): + continue + ckpt1 = defect_dir / "checkpoints" / "checkpoint_final.pt" + ckpt2 = defect_dir / "checkpoint_final.pt" + if ckpt1.exists() or ckpt2.exists(): + defect_types.append(defect_dir.name) + if defect_types: + mapping[obj_dir.name] = defect_types + + print(f"[Batch] Auto-discovered checkpoints: {json.dumps(mapping, indent=2)}") + return mapping + + +def load_manifest(manifest_path: str) -> Dict[str, Dict]: + """Load CSV manifest mapping image paths to object_class, defect_type, and product descriptions.""" + manifest = {} + with open(manifest_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + img_path = row.get('image_path', row.get('path', row.get('image', ''))).strip() + desc = row.get('product_description', row.get('description', row.get('product', ''))).strip() + obj_cls = row.get('object_class', row.get('object', '')).strip() + dfc_type = row.get('defect_type', row.get('defect', '')).strip() + if img_path: + manifest[Path(img_path).resolve()] = { + 'product_description': desc, + 'object_class': obj_cls, + 'defect_type': dfc_type + } + print(f"[Batch] Loaded manifest with {len(manifest)} entries") + return manifest + + +def discover_images(input_dir: str, extensions=('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')) -> List[Path]: + """Recursively discover all images in input directory.""" + input_path = Path(input_dir) + images = [] + for ext in extensions: + images.extend(input_path.rglob(f"*{ext}")) + images.extend(input_path.rglob(f"*{ext.upper()}")) + unique = sorted(set(images)) + print(f"[Batch] Discovered {len(unique)} images in {input_dir}") + return unique + + +def run_batch( + input_dir: str, + output_dir: str, + checkpoint_dir: str, + object_class: Optional[str] = None, + defect_type: Optional[str] = None, + product_desc: Optional[str] = None, + manifest_path: Optional[str] = None, + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + max_defects_per_image: int = 3, + device: str = 'cuda', + vlm_model: str = 'gemma3:12b', + image_size: int = 512, + num_steps: int = 50, + guidance_scale: float = 7.5, + resume: bool = False, + save_failed: bool = True +): + """Run agent orchestrator over all images in input directory.""" + + timestamp = datetime.now().strftime("%Y%m%d_%H%M") + output_path = Path(output_dir) / timestamp + output_path.mkdir(parents=True, exist_ok=True) + print(f"[Batch] Output folder: {output_path}") + + # ------------------------------------------------------------------ + # Auto-discover valid lists from checkpoint directory if not provided + # ------------------------------------------------------------------ + if valid_defect_types is None: + valid_defect_types = discover_checkpoints(checkpoint_dir) + if valid_object_classes is None: + valid_object_classes = list(valid_defect_types.keys()) + + print(f"[Batch] Valid object classes: {valid_object_classes}") + print(f"[Batch] Valid defect types mapping: {json.dumps(valid_defect_types, indent=2)}") + + # ------------------------------------------------------------------ + # Load manifest if provided + # ------------------------------------------------------------------ + manifest = {} + if manifest_path and os.path.exists(manifest_path): + manifest = load_manifest(manifest_path) + print(f"[Batch] Scenario A Active: CSV Manifest ({len(manifest)} entries)") + elif product_desc: + print(f"[Batch] Scenario B Active: Global Description = '{product_desc}'") + if object_class: + print(f"[Batch] Global object_class = '{object_class}'") + if defect_type: + print(f"[Batch] Global defect_type = '{defect_type}'") + else: + print("[Batch] Scenario C Active: Folder Name Inference (no description provided)") + + images = discover_images(input_dir) + if not images: + print("[Batch] No images found. Exiting.") + return + + progress_file = output_path / "batch_progress.json" + processed_ids = set() + if resume and progress_file.exists(): + with open(progress_file, 'r') as f: + progress = json.load(f) + processed_ids = set(progress.get('processed_paths', [])) + print(f"[Batch] Resuming: {len(processed_ids)} images already processed") + + # ------------------------------------------------------------------ + # Initialize orchestrator once with the valid lists + # ------------------------------------------------------------------ + orchestrator = ArtiAgentOrchestrator( + device=device, + output_dir=str(output_path), + vlm_model=vlm_model, + checkpoint_dir=checkpoint_dir, + object_class=object_class or "", + defect_type=defect_type or "", + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + image_size=image_size, + num_steps=num_steps, + guidance_scale=guidance_scale + ) + + stats = { + 'total': len(images), + 'processed': 0, + 'successful': 0, + 'failed': 0, + 'defects_generated': 0, + 'start_time': datetime.now().isoformat(), + 'processed_paths': [], + 'failed_images': [] + } + + if resume: + images = [img for img in images if str(img.resolve()) not in processed_ids] + + print(f"[Batch] Processing {len(images)} images...") + print(f"[Batch] Max defects per image: {max_defects_per_image}") + print("=" * 70) + + for img_path in tqdm(images, desc="Agent Batch Processing"): + img_key = str(img_path.resolve()) + + try: + # Resolve description, object_class, and defect_type per image + if img_key in manifest: + entry = manifest[img_key] + desc = entry.get('product_description') or product_desc or infer_product_from_path(img_path) + obj_cls = entry.get('object_class') or object_class or None + dfc_type = entry.get('defect_type') or defect_type or None + source = "manifest" + elif product_desc: + desc = product_desc + obj_cls = object_class or None + dfc_type = defect_type or None + source = "global" + else: + desc = infer_product_from_path(img_path) + obj_cls = object_class or infer_object_class_from_path(img_path) or None + dfc_type = defect_type or None + source = "inferred" + + # If object_class is still None, VLM will select from valid_object_classes + # If defect_type is still None, VLM will select from valid_defect_types for the chosen object_class + # If object_class is provided but not in valid list, warn + if obj_cls and valid_object_classes and obj_cls not in valid_object_classes: + print(f"[Batch] Warning: object_class '{obj_cls}' not in valid list {valid_object_classes}. " + f"VLM will select a valid one.") + obj_cls = None + + print(f"\n[Batch] Processing: {img_path.name} | desc source: {source}") + if obj_cls: + print(f"[Batch] object_class: '{obj_cls}' (user-provided)") + else: + print(f"[Batch] object_class: ") + if dfc_type: + print(f"[Batch] defect_type: '{dfc_type}' (user-provided)") + else: + print(f"[Batch] defect_type: ") + if source in ['inferred', 'global']: + print(f"[Batch] Using description: '{desc}'") + + result = orchestrator.run( + product_description=desc, + image_path=str(img_path), + max_defects=max_defects_per_image, + defect_type=dfc_type, + object_class=obj_cls + ) + + successful_defects = sum(1 for r in result['results'] if r['success']) + + stats['processed'] += 1 + stats['successful'] += 1 if successful_defects > 0 else 0 + stats['defects_generated'] += successful_defects + stats['processed_paths'].append(img_key) + + if successful_defects == 0: + stats['failed'] += 1 + stats['failed_images'].append({'path': img_key, 'reason': 'no_defects_generated'}) + + if stats['processed'] % 5 == 0: + with open(progress_file, 'w') as f: + json.dump(stats, f, indent=2) + + except Exception as e: + stats['failed'] += 1 + stats['failed_images'].append({'path': img_key, 'reason': str(e)}) + print(f"[Batch] FAILED: {img_path.name} -> {str(e)}") + if save_failed: + fail_dir = output_path / "_failed" / img_path.stem + fail_dir.mkdir(parents=True, exist_ok=True) + with open(fail_dir / "error.txt", 'w') as f: + f.write(traceback.format_exc()) + + with open(progress_file, 'w') as f: + json.dump(stats, f, indent=2) + + orchestrator.cleanup() + + elapsed = (datetime.now() - datetime.fromisoformat(stats['start_time'])).total_seconds() + hours = int(elapsed // 3600) + minutes = int((elapsed % 3600) // 60) + seconds = int(elapsed % 60) + + print("\n" + "=" * 70) + print("BATCH ORCHESTRATION COMPLETE") + print("=" * 70) + print(f"Total images: {stats['total']}") + print(f"Processed: {stats['processed']}") + print(f"Successful: {stats['successful']}") + print(f"Failed: {stats['failed']}") + print(f"Defects generated: {stats['defects_generated']}") + print(f"Total time: {hours}h {minutes}m {seconds}s") + print(f"Output directory: {output_path}") + print("=" * 70) + + +def main(): + parser = argparse.ArgumentParser(description='Batch Agent-Driven Defect Generation (DefectFill with VLM selection)') + + # Input / Output + parser.add_argument('--input-dir', required=True, help='Directory containing clean product images') + parser.add_argument('--output-dir', required=True, help='Output directory for all defect images') + + # DefectFill checkpoint routing + parser.add_argument('--checkpoint-dir', required=True, + help='Root directory containing object_class/defect_type checkpoint subfolders') + + # Description / routing sources + parser.add_argument('--product-desc', default=None, + help='Global product description applied to ALL images (drives VLM selection)') + parser.add_argument('--object-class', default=None, + help='Global object_class for ALL images (optional; VLM selects if omitted)') + parser.add_argument('--manifest', default=None, + help='CSV manifest with columns: image_path,object_class,defect_type,product_description') + + # Defect control + parser.add_argument('--defect-type', default=None, + help='Global defect type for ALL images (optional; VLM selects if omitted)') + parser.add_argument('--max-defects-per-image', type=int, default=3, + help='Maximum defects to generate per image (default: 3)') + + # Valid lists (auto-discovered from checkpoint-dir if not provided) + parser.add_argument('--valid-object-classes', default=None, + help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" ) + parser.add_argument('--valid-defect-types', default=None, + help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" ) + # Generation control + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, help='Generation resolution') + parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps') + parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale') + parser.add_argument('--resume', action='store_true', help='Resume from previous batch run') + parser.add_argument('--no-save-failed', action='store_true', help='Do not save failed case logs') + + args = parser.parse_args() + + # Parse valid lists from CLI + valid_object_classes = None + valid_defect_types = None + if args.valid_object_classes: + valid_object_classes = json.loads(args.valid_object_classes) + if args.valid_defect_types: + valid_defect_types = json.loads(args.valid_defect_types) + + run_batch( + input_dir=args.input_dir, + output_dir=args.output_dir, + checkpoint_dir=args.checkpoint_dir, + object_class=args.object_class, + defect_type=args.defect_type, + product_desc=args.product_desc, + manifest_path=args.manifest, + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + max_defects_per_image=args.max_defects_per_image, + device=args.device, + vlm_model=args.vlm_model, + image_size=args.image_size, + num_steps=args.num_steps, + guidance_scale=args.guidance_scale, + resume=args.resume, + save_failed=not args.no_save_failed + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/batch_agent_orchestrator.py b/ArtiAgent - DefectFill/src/batch_agent_orchestrator.py new file mode 100644 index 0000000000000000000000000000000000000000..6dd279a4f4b7eb9afbb380bc4fe33e1c056c126d --- /dev/null +++ b/ArtiAgent - DefectFill/src/batch_agent_orchestrator.py @@ -0,0 +1,414 @@ +""" +Batch Agent Orchestrator โ€” DefectFill Edition (with VLM list selection) + +Auto-discovers valid object classes and defect types from the checkpoint directory. +VLM selects object_class from the discovered list based on product description. +If defect_type is not provided, VLM selects from the per-object-class list. + +Usage: + # Auto-discover checkpoints, VLM selects everything from product description + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage" \ + --output-dir "C:/AgentOutput" \ + --checkpoint-dir "C:/.../checkpoints" \ + --product-desc "VCSEL laser diode with emission aperture" \ + --device cuda + + # Provide explicit valid lists (overrides auto-discovery) + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage" \ + --output-dir "C:/AgentOutput" \ + --checkpoint-dir "C:/.../checkpoints" \ + --valid-object-classes '["xray_PCB","vcsel"]' \ + --valid-defect-types '{"xray_PCB":["xray_die","bubble"],"vcsel":["scratch"]}' \ + --product-desc "VCSEL laser diode" \ + --device cuda + + # CSV manifest with per-image routing (overrides VLM selection for those images) + python batch_agent_orchestrator.py \ + --input-dir "C:/TestingImage" \ + --output-dir "C:/AgentOutput" \ + --checkpoint-dir "C:/.../checkpoints" \ + --manifest "C:/products.csv" \ + --device cuda +""" + +import os +import sys +import csv +import json +import argparse +import traceback +from pathlib import Path +from datetime import datetime +from typing import Dict, List, Optional +from tqdm import tqdm + +import numpy as np +from PIL import Image + +SCRIPT_DIR = Path(__file__).parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from artiagent_orchestrator import ArtiAgentOrchestrator + + +def infer_product_from_path(image_path: Path) -> str: + """Infer product description from folder structure or filename.""" + parent = image_path.parent.name.lower() + if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']: + return parent.replace('_', ' ').replace('-', ' ') + stem = image_path.stem.lower() + for keyword in ['vcsel', 'lens', 'die', 'photodiode', 'sensor', 'chip', 'led', 'laser', 'optical']: + if keyword in stem: + return keyword + return "electronic component" + + +def infer_object_class_from_path(image_path: Path) -> str: + """Infer object_class from parent folder name.""" + parent = image_path.parent.name.lower() + if parent and parent not in ['.', '', 'images', 'imgs', 'data', 'input']: + return parent + return "" + + +def discover_checkpoints(checkpoint_dir: str) -> Dict[str, List[str]]: + """ + Scan checkpoint directory to discover available object_class -> defect_type mappings. + Expected structure: + checkpoint_dir/ + xray_PCB/ + xray_die/ + checkpoints/checkpoint_final.pt + bubble/ + checkpoints/checkpoint_final.pt + """ + cp = Path(checkpoint_dir) + mapping = {} + if not cp.exists(): + print(f"[Batch] Warning: checkpoint-dir does not exist: {checkpoint_dir}") + return mapping + + for obj_dir in cp.iterdir(): + if not obj_dir.is_dir(): + continue + defect_types = [] + for defect_dir in obj_dir.iterdir(): + if not defect_dir.is_dir(): + continue + ckpt1 = defect_dir / "checkpoints" / "checkpoint_final.pt" + ckpt2 = defect_dir / "checkpoint_final.pt" + if ckpt1.exists() or ckpt2.exists(): + defect_types.append(defect_dir.name) + if defect_types: + mapping[obj_dir.name] = defect_types + + print(f"[Batch] Auto-discovered checkpoints: {json.dumps(mapping, indent=2)}") + return mapping + + +def load_manifest(manifest_path: str) -> Dict[str, Dict]: + """Load CSV manifest mapping image paths to object_class, defect_type, and product descriptions.""" + manifest = {} + with open(manifest_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + img_path = row.get('image_path', row.get('path', row.get('image', ''))).strip() + desc = row.get('product_description', row.get('description', row.get('product', ''))).strip() + obj_cls = row.get('object_class', row.get('object', '')).strip() + dfc_type = row.get('defect_type', row.get('defect', '')).strip() + if img_path: + manifest[Path(img_path).resolve()] = { + 'product_description': desc, + 'object_class': obj_cls, + 'defect_type': dfc_type + } + print(f"[Batch] Loaded manifest with {len(manifest)} entries") + return manifest + + +def discover_images(input_dir: str, extensions=('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')) -> List[Path]: + """Recursively discover all images in input directory.""" + input_path = Path(input_dir) + images = [] + for ext in extensions: + images.extend(input_path.rglob(f"*{ext}")) + images.extend(input_path.rglob(f"*{ext.upper()}")) + unique = sorted(set(images)) + print(f"[Batch] Discovered {len(unique)} images in {input_dir}") + return unique + + +def run_batch( + input_dir: str, + output_dir: str, + checkpoint_dir: str, + object_class: Optional[str] = None, + defect_type: Optional[str] = None, + product_desc: Optional[str] = None, + manifest_path: Optional[str] = None, + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + num_defects_per_image: int = 3, + device: str = 'cuda', + vlm_model: str = 'gemma3:12b', + image_size: int = 512, + num_steps: int = 50, + guidance_scale: float = 7.5, + resume: bool = False, + save_failed: bool = True +): + """Run agent orchestrator over all images in input directory.""" + + timestamp = datetime.now().strftime("%Y%m%d_%H%M") + output_path = Path(output_dir) / timestamp + output_path.mkdir(parents=True, exist_ok=True) + print(f"[Batch] Output folder: {output_path}") + + # ------------------------------------------------------------------ + # Auto-discover valid lists from checkpoint directory if not provided + # ------------------------------------------------------------------ + if valid_defect_types is None: + valid_defect_types = discover_checkpoints(checkpoint_dir) + if valid_object_classes is None: + valid_object_classes = list(valid_defect_types.keys()) + + print(f"[Batch] Valid object classes: {valid_object_classes}") + print(f"[Batch] Valid defect types mapping: {json.dumps(valid_defect_types, indent=2)}") + + # ------------------------------------------------------------------ + # Load manifest if provided + # ------------------------------------------------------------------ + manifest = {} + if manifest_path and os.path.exists(manifest_path): + manifest = load_manifest(manifest_path) + print(f"[Batch] Scenario A Active: CSV Manifest ({len(manifest)} entries)") + elif product_desc: + print(f"[Batch] Scenario B Active: Global Description = '{product_desc}'") + if object_class: + print(f"[Batch] Global object_class = '{object_class}'") + if defect_type: + print(f"[Batch] Global defect_type = '{defect_type}'") + else: + print("[Batch] Scenario C Active: Folder Name Inference (no description provided)") + + images = discover_images(input_dir) + if not images: + print("[Batch] No images found. Exiting.") + return + + progress_file = output_path / "batch_progress.json" + processed_ids = set() + if resume and progress_file.exists(): + with open(progress_file, 'r') as f: + progress = json.load(f) + processed_ids = set(progress.get('processed_paths', [])) + print(f"[Batch] Resuming: {len(processed_ids)} images already processed") + + # ------------------------------------------------------------------ + # Initialize orchestrator once with the valid lists + # ------------------------------------------------------------------ + orchestrator = ArtiAgentOrchestrator( + device=device, + output_dir=str(output_path), + vlm_model=vlm_model, + checkpoint_dir=checkpoint_dir, + object_class=object_class or "", + defect_type=defect_type or "", + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + image_size=image_size, + num_steps=num_steps, + guidance_scale=guidance_scale + ) + + stats = { + 'total': len(images), + 'processed': 0, + 'successful': 0, + 'failed': 0, + 'defects_generated': 0, + 'start_time': datetime.now().isoformat(), + 'processed_paths': [], + 'failed_images': [] + } + + if resume: + images = [img for img in images if str(img.resolve()) not in processed_ids] + + print(f"[Batch] Processing {len(images)} images...") + print(f"[Batch] Max defects per image: {num_defects_per_image}") + print("=" * 70) + + for img_path in tqdm(images, desc="Agent Batch Processing"): + img_key = str(img_path.resolve()) + + try: + # Resolve description, object_class, and defect_type per image + if img_key in manifest: + entry = manifest[img_key] + desc = entry.get('product_description') or product_desc or infer_product_from_path(img_path) + obj_cls = entry.get('object_class') or object_class or None + dfc_type = entry.get('defect_type') or defect_type or None + source = "manifest" + elif product_desc: + desc = product_desc + obj_cls = object_class or None + dfc_type = defect_type or None + source = "global" + else: + desc = infer_product_from_path(img_path) + obj_cls = object_class or infer_object_class_from_path(img_path) or None + dfc_type = defect_type or None + source = "inferred" + + # If object_class is still None, VLM will select from valid_object_classes + # If defect_type is still None, VLM will select from valid_defect_types for the chosen object_class + # If object_class is provided but not in valid list, warn + if obj_cls and valid_object_classes and obj_cls not in valid_object_classes: + print(f"[Batch] Warning: object_class '{obj_cls}' not in valid list {valid_object_classes}. " + f"VLM will select a valid one.") + obj_cls = None + + print(f"\n[Batch] Processing: {img_path.name} | desc source: {source}") + if obj_cls: + print(f"[Batch] object_class: '{obj_cls}' (user-provided)") + else: + print(f"[Batch] object_class: ") + if dfc_type: + print(f"[Batch] defect_type: '{dfc_type}' (user-provided)") + else: + print(f"[Batch] defect_type: ") + if source in ['inferred', 'global']: + print(f"[Batch] Using description: '{desc}'") + + result = orchestrator.run( + product_description=desc, + image_path=str(img_path), + num_defects=num_defects_per_image, + defect_type=dfc_type, + object_class=obj_cls + ) + + successful_defects = sum(1 for r in result['results'] if r['success']) + + stats['processed'] += 1 + stats['successful'] += 1 if successful_defects > 0 else 0 + stats['defects_generated'] += successful_defects + stats['processed_paths'].append(img_key) + + if successful_defects == 0: + stats['failed'] += 1 + stats['failed_images'].append({'path': img_key, 'reason': 'no_defects_generated'}) + + if stats['processed'] % 5 == 0: + with open(progress_file, 'w') as f: + json.dump(stats, f, indent=2) + + except Exception as e: + stats['failed'] += 1 + stats['failed_images'].append({'path': img_key, 'reason': str(e)}) + print(f"[Batch] FAILED: {img_path.name} -> {str(e)}") + if save_failed: + fail_dir = output_path / "_failed" / img_path.stem + fail_dir.mkdir(parents=True, exist_ok=True) + with open(fail_dir / "error.txt", 'w') as f: + f.write(traceback.format_exc()) + + with open(progress_file, 'w') as f: + json.dump(stats, f, indent=2) + + orchestrator.cleanup() + + elapsed = (datetime.now() - datetime.fromisoformat(stats['start_time'])).total_seconds() + hours = int(elapsed // 3600) + minutes = int((elapsed % 3600) // 60) + seconds = int(elapsed % 60) + + print("\n" + "=" * 70) + print("BATCH ORCHESTRATION COMPLETE") + print("=" * 70) + print(f"Total images: {stats['total']}") + print(f"Processed: {stats['processed']}") + print(f"Successful: {stats['successful']}") + print(f"Failed: {stats['failed']}") + print(f"Defects generated: {stats['defects_generated']}") + print(f"Total time: {hours}h {minutes}m {seconds}s") + print(f"Output directory: {output_path}") + print("=" * 70) + + +def main(): + parser = argparse.ArgumentParser(description='Batch Agent-Driven Defect Generation (DefectFill with VLM selection)') + + # Input / Output + parser.add_argument('--input-dir', required=True, help='Directory containing clean product images') + parser.add_argument('--output-dir', required=True, help='Output directory for all defect images') + + # DefectFill checkpoint routing + parser.add_argument('--checkpoint-dir', required=True, + help='Root directory containing object_class/defect_type checkpoint subfolders') + + # Description / routing sources + parser.add_argument('--product-desc', default=None, + help='Global product description applied to ALL images (drives VLM selection)') + parser.add_argument('--object-class', default=None, + help='Global object_class for ALL images (optional; VLM selects if omitted)') + parser.add_argument('--manifest', default=None, + help='CSV manifest with columns: image_path,object_class,defect_type,product_description') + + # Defect control + parser.add_argument('--defect-type', default=None, + help='Global defect type for ALL images (optional; VLM selects if omitted)') + parser.add_argument('--num-defects-per-image', type=int, default=3, + help='Number of defects to generate per image (default: 3)') + + # Valid lists (auto-discovered from checkpoint-dir if not provided) + parser.add_argument('--valid-object-classes', default=None, + help="JSON array of valid object classes, e.g., '[\"xray_PCB\",\"vcsel\"]'" ) + parser.add_argument('--valid-defect-types', default=None, + help="JSON dict mapping object_class to defect types, e.g., '{\"xray_PCB\":[\"xray_die\",\"bubble\"]}'" ) + # Generation control + parser.add_argument('--device', default='cuda', help='Device (cuda/cpu)') + parser.add_argument('--vlm-model', default='gemma3:12b', help='Local VLM model') + parser.add_argument('--image-size', type=int, default=512, help='Generation resolution') + parser.add_argument('--num-steps', type=int, default=50, help='Denoising steps') + parser.add_argument('--guidance-scale', type=float, default=7.5, help='CFG scale') + parser.add_argument('--resume', action='store_true', help='Resume from previous batch run') + parser.add_argument('--no-save-failed', action='store_true', help='Do not save failed case logs') + + args = parser.parse_args() + + # Parse valid lists from CLI + valid_object_classes = None + valid_defect_types = None + if args.valid_object_classes: + valid_object_classes = json.loads(args.valid_object_classes) + if args.valid_defect_types: + valid_defect_types = json.loads(args.valid_defect_types) + + run_batch( + input_dir=args.input_dir, + output_dir=args.output_dir, + checkpoint_dir=args.checkpoint_dir, + object_class=args.object_class, + defect_type=args.defect_type, + product_desc=args.product_desc, + manifest_path=args.manifest, + valid_object_classes=valid_object_classes, + valid_defect_types=valid_defect_types, + num_defects_per_image=args.num_defects_per_image, + device=args.device, + vlm_model=args.vlm_model, + image_size=args.image_size, + num_steps=args.num_steps, + guidance_scale=args.guidance_scale, + resume=args.resume, + save_failed=not args.no_save_failed + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/detect_similar_product.py b/ArtiAgent - DefectFill/src/detect_similar_product.py new file mode 100644 index 0000000000000000000000000000000000000000..ca1a5b8ad125506a60e4f53f68a23d2c33db3f75 --- /dev/null +++ b/ArtiAgent - DefectFill/src/detect_similar_product.py @@ -0,0 +1,83 @@ +import os +from pathlib import Path +import torch +import torch.nn.functional as F +from PIL import Image +from torchvision import transforms +from flask import Flask, request, jsonify, send_file + +# 1. Load DINOv2 for fast feature extraction +device = "cuda" if torch.cuda.is_available() else "cpu" +dinov2 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14").to(device) +dinov2.eval() + +img_transforms = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), +]) + +# Folder structure containing 1 reference image per TRAINED product line +# Example: trained_gallery/vcsel_cap_v1.png, trained_gallery/pcb_rev2.png +# TRAINED_GALLERY_DIR = "storage/trained_products_gallery" + +def get_image_embedding(image_path: str) -> torch.Tensor: + """Extracts a normalized 384-dim feature vector using DINOv2.""" + img = Image.open(image_path).convert("RGB") + tensor = img_transforms(img).unsqueeze(0).to(device) + with torch.no_grad(): + embedding = dinov2(tensor) + return F.normalize(embedding, p=2, dim=1) + +def check_product_training_status(target_image_path: str, gallery_dir: Path, similarity_threshold: float = 0.85) -> dict: + """Checks if the target image matches any trained product in the gallery directory.""" + if not gallery_dir.exists() or not os.listdir(gallery_dir): + return { + "is_trained": False, + "matched_product": None, + "confidence": 0.0, + "message": "No trained product models found in system registry." + } + + # if not os.path.exists(TRAINED_GALLERY_DIR) or not os.listdir(TRAINED_GALLERY_DIR): + # return { + # "is_trained": False, + # "matched_product": None, + # "confidence": 0.0, + # "message": "No trained product models found in system registry." + # } + + target_emb = get_image_embedding(target_image_path) + + best_score = 0.0 + best_product_name = None + + # Compare incoming image against all trained product templates + for fname in os.listdir(gallery_dir): + if not fname.lower().endswith(('.png', '.jpg', '.jpeg')): + continue + + ref_path = os.path.join(gallery_dir, fname) + ref_emb = get_image_embedding(ref_path) + + # Calculate Cosine Similarity + similarity = (target_emb @ ref_emb.T).item() + + if similarity > best_score: + best_score = similarity + best_product_name = os.path.splitext(fname)[0] # Extract product name + + if best_score >= similarity_threshold: + return { + "is_trained": True, + "matched_product": best_product_name, + "confidence": round(best_score, 4), + "message": f"Product recognized as '{best_product_name}'. Ready for DefectFill generation." + } + else: + return { + "is_trained": False, + "matched_product": None, + "confidence": round(best_score, 4), + "message": f"Product not recognized (Best match confidence: {round(best_score * 100, 1)}%). Please upload at least 1 clean reference images to train this new product type first." + } \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/flux/__init__.py b/ArtiAgent - DefectFill/src/flux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..43c365a49d6980e88acba10ef3069f110a59644a --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/__init__.py @@ -0,0 +1,11 @@ +try: + from ._version import version as __version__ # type: ignore + from ._version import version_tuple +except ImportError: + __version__ = "unknown (no version information available)" + version_tuple = (0, 0, "unknown", "noinfo") + +from pathlib import Path + +PACKAGE = __package__.replace("_", "-") +PACKAGE_ROOT = Path(__file__).parent diff --git a/ArtiAgent - DefectFill/src/flux/__main__.py b/ArtiAgent - DefectFill/src/flux/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..d5cf0fd2444d4cda4053fa74dad3371556b886e5 --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/__main__.py @@ -0,0 +1,4 @@ +from .cli import app + +if __name__ == "__main__": + app() diff --git a/ArtiAgent - DefectFill/src/flux/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46cf619948bf6c42c3e8f33cb0710cf5f75492c2 Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/__pycache__/artifacts_util.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/__pycache__/artifacts_util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb6eccc55fd52291945088cd96b2e2e04ba57333 Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/__pycache__/artifacts_util.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/__pycache__/math.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/__pycache__/math.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3752064f736ab300fcf0f6bd9e72a412c3416df0 Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/__pycache__/math.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/__pycache__/model.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86e79a513113889165d28e2121f97243b014a50c Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/__pycache__/model.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/__pycache__/sampling.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/__pycache__/sampling.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edf1f9c6090dbb91b2f7f54ccf86f1afeb62ffef Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/__pycache__/sampling.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/__pycache__/util.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/__pycache__/util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3e2b9ff8fdf2cad9c44a1008c1943dac8ab8073 Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/__pycache__/util.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/api.py b/ArtiAgent - DefectFill/src/flux/api.py new file mode 100644 index 0000000000000000000000000000000000000000..b08202adb35d2ffae320bb9b47f567e538837836 --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/api.py @@ -0,0 +1,194 @@ +import io +import os +import time +from pathlib import Path + +import requests +from PIL import Image + +API_ENDPOINT = "https://api.bfl.ml" + + +class ApiException(Exception): + def __init__(self, status_code: int, detail: str | list[dict] | None = None): + super().__init__() + self.detail = detail + self.status_code = status_code + + def __str__(self) -> str: + return self.__repr__() + + def __repr__(self) -> str: + if self.detail is None: + message = None + elif isinstance(self.detail, str): + message = self.detail + else: + message = "[" + ",".join(d["msg"] for d in self.detail) + "]" + return f"ApiException({self.status_code=}, {message=}, detail={self.detail})" + + +class ImageRequest: + def __init__( + self, + prompt: str, + width: int = 1024, + height: int = 1024, + name: str = "flux.1-pro", + num_steps: int = 50, + prompt_upsampling: bool = False, + seed: int | None = None, + validate: bool = True, + launch: bool = True, + api_key: str | None = None, + ): + """ + Manages an image generation request to the API. + + Args: + prompt: Prompt to sample + width: Width of the image in pixel + height: Height of the image in pixel + name: Name of the model + num_steps: Number of network evaluations + prompt_upsampling: Use prompt upsampling + seed: Fix the generation seed + validate: Run input validation + launch: Directly launches request + api_key: Your API key if not provided by the environment + + Raises: + ValueError: For invalid input + ApiException: For errors raised from the API + """ + if validate: + if name not in ["flux.1-pro"]: + raise ValueError(f"Invalid model {name}") + elif width % 32 != 0: + raise ValueError(f"width must be divisible by 32, got {width}") + elif not (256 <= width <= 1440): + raise ValueError(f"width must be between 256 and 1440, got {width}") + elif height % 32 != 0: + raise ValueError(f"height must be divisible by 32, got {height}") + elif not (256 <= height <= 1440): + raise ValueError(f"height must be between 256 and 1440, got {height}") + elif not (1 <= num_steps <= 50): + raise ValueError(f"steps must be between 1 and 50, got {num_steps}") + + self.request_json = { + "prompt": prompt, + "width": width, + "height": height, + "variant": name, + "steps": num_steps, + "prompt_upsampling": prompt_upsampling, + } + if seed is not None: + self.request_json["seed"] = seed + + self.request_id: str | None = None + self.result: dict | None = None + self._image_bytes: bytes | None = None + self._url: str | None = None + if api_key is None: + self.api_key = os.environ.get("BFL_API_KEY") + else: + self.api_key = api_key + + if launch: + self.request() + + def request(self): + """ + Request to generate the image. + """ + if self.request_id is not None: + return + response = requests.post( + f"{API_ENDPOINT}/v1/image", + headers={ + "accept": "application/json", + "x-key": self.api_key, + "Content-Type": "application/json", + }, + json=self.request_json, + ) + result = response.json() + if response.status_code != 200: + raise ApiException(status_code=response.status_code, detail=result.get("detail")) + self.request_id = response.json()["id"] + + def retrieve(self) -> dict: + """ + Wait for the generation to finish and retrieve response. + """ + if self.request_id is None: + self.request() + while self.result is None: + response = requests.get( + f"{API_ENDPOINT}/v1/get_result", + headers={ + "accept": "application/json", + "x-key": self.api_key, + }, + params={ + "id": self.request_id, + }, + ) + result = response.json() + if "status" not in result: + raise ApiException(status_code=response.status_code, detail=result.get("detail")) + elif result["status"] == "Ready": + self.result = result["result"] + elif result["status"] == "Pending": + time.sleep(0.5) + else: + raise ApiException(status_code=200, detail=f"API returned status '{result['status']}'") + return self.result + + @property + def bytes(self) -> bytes: + """ + Generated image as bytes. + """ + if self._image_bytes is None: + response = requests.get(self.url) + if response.status_code == 200: + self._image_bytes = response.content + else: + raise ApiException(status_code=response.status_code) + return self._image_bytes + + @property + def url(self) -> str: + """ + Public url to retrieve the image from + """ + if self._url is None: + result = self.retrieve() + self._url = result["sample"] + return self._url + + @property + def image(self) -> Image.Image: + """ + Load the image as a PIL Image + """ + return Image.open(io.BytesIO(self.bytes)) + + def save(self, path: str): + """ + Save the generated image to a local path + """ + suffix = Path(self.url).suffix + if not path.endswith(suffix): + path = path + suffix + Path(path).resolve().parent.mkdir(parents=True, exist_ok=True) + with open(path, "wb") as file: + file.write(self.bytes) + + +if __name__ == "__main__": + from fire import Fire + + Fire(ImageRequest) diff --git a/ArtiAgent - DefectFill/src/flux/artifacts_util.py b/ArtiAgent - DefectFill/src/flux/artifacts_util.py new file mode 100644 index 0000000000000000000000000000000000000000..98caf48d4614d0de4eee997a6b3fee88fd35c848 --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/artifacts_util.py @@ -0,0 +1,335 @@ +import torch +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as patches +import os + +def patch_coor_to_ind(x, y, w, txt_len): + return y * w + x + txt_len + +def patch_ind_to_coor(ind, w, txt_len, return_shape=False): + ind -= txt_len + y = ind // w + x = ind % w + return [y, x] + +def patch_indices_to_coords(indices, w, txt_len=512): + """ + Convert patch indices back to patch coordinates. + + Args: + indices: List of patch indices + w: Patch width + txt_len: Text length offset + + Returns: + List of (y, x) patch coordinates + """ + return [patch_ind_to_coor(ind, w, txt_len) for ind in indices] + +def bbox_to_patch_indices(bbox_coordinates, h, w, patch_size=16, txt_len=512): + xmin, xmax, ymin, ymax = bbox_coordinates + patch_xmin = xmin // patch_size + patch_xmax = (xmax - 1) // patch_size + patch_ymin = ymin // patch_size + patch_ymax = (ymax - 1) // patch_size + indices = [ + patch_coor_to_ind(px, py, w, txt_len) + for py in range(patch_ymin, patch_ymax + 1) + for px in range(patch_xmin, patch_xmax + 1) + ] + return indices + +def bbox_to_patch_coords(bbox_coordinates, patch_size=16, return_shape=False): + xmin, xmax, ymin, ymax = bbox_coordinates + patch_xmin = xmin // patch_size + patch_xmax = (xmax - 1) // patch_size + patch_ymin = ymin // patch_size + patch_ymax = (ymax - 1) // patch_size + coords = [ (py, px) + for py in range(patch_ymin, patch_ymax + 1) + for px in range(patch_xmin, patch_xmax + 1) + ] + if return_shape: + return coords, patch_ymax-patch_ymin+1, patch_xmax-patch_xmin+1 + return coords + +# New shape-based functions +def mask_to_patch_indices(mask, patch_size=16, txt_len=512): + """ + Convert a binary mask to patch indices. + + Args: + mask: Binary mask (numpy array) where 1 indicates the region of interest + patch_size: Size of each patch (default 16 for Flux) + txt_len: Text length offset (default 512) + + Returns: + List of patch indices + """ + h, w = mask.shape + patch_h, patch_w = h // patch_size, w // patch_size + + # Downsample mask to patch resolution + patch_mask = np.zeros((patch_h, patch_w), dtype=bool) + + for py in range(patch_h): + for px in range(patch_w): + # Get the patch region in the original mask + y_start, y_end = py * patch_size, (py + 1) * patch_size + x_start, x_end = px * patch_size, (px + 1) * patch_size + + # If any part of the patch overlaps with the mask, include it + patch_region = mask[y_start:y_end, x_start:x_end] + if np.any(patch_region): + patch_mask[py, px] = True + + # Convert patch coordinates to indices + indices = [] + for py in range(patch_h): + for px in range(patch_w): + if patch_mask[py, px]: + indices.append(patch_coor_to_ind(px, py, patch_w, txt_len)) + + return indices + +def mask_to_patch_coords(mask, patch_size=16): + """ + Convert a binary mask to patch coordinates. + + Args: + mask: Binary mask (numpy array) where 1 indicates the region of interest + patch_size: Size of each patch (default 16 for Flux) + + Returns: + List of (py, px) patch coordinates + """ + h, w = mask.shape + patch_h, patch_w = h // patch_size, w // patch_size + + # Downsample mask to patch resolution + patch_mask = np.zeros((patch_h, patch_w), dtype=bool) + + for py in range(patch_h): + for px in range(patch_w): + # Get the patch region in the original mask + y_start, y_end = py * patch_size, (py + 1) * patch_size + x_start, x_end = px * patch_size, (px + 1) * patch_size + + # If any part of the patch overlaps with the mask, include it + patch_region = mask[y_start:y_end, x_start:x_end] + if np.any(patch_region): + patch_mask[py, px] = True + + # Convert patch coordinates to list + coords = [] + for py in range(patch_h): + for px in range(patch_w): + if patch_mask[py, px]: + coords.append((py, px)) + + return coords + +def get_closest_patch_ind(h, w, bbox_coordinates, patch_size=16, txt_len=512): + # Get only valid coordinates (value == 1) + + large_array = np.ones((h,w), dtype=int) + small_grid_coords, bbox_h, bbox_w = bbox_to_patch_coords(bbox_coordinates, patch_size=patch_size, return_shape=True) + + for y, x in small_grid_coords: + large_array[y,x] = 0 + + valid_coords = np.argwhere(large_array == 1) + + result = np.empty((bbox_h, bbox_w), dtype=object) + min_h, min_w = small_grid_coords[0] + + for idx, coord in enumerate(small_grid_coords): + y, x = coord + distances = np.abs(valid_coords[:, 0] - y) + np.abs(valid_coords[:, 1] - x) + min_idx = np.argmin(distances) + closest_coord = tuple(valid_coords[min_idx]) + result[y-min_h,x-min_w] = closest_coord + + return [patch_coor_to_ind(x,y,w,txt_len) for y, x in result.flatten()] + +def get_neighbors_patch_ind(h, w, bbox_coordinates, img_ids, patch_size=16, txt_len=512): + # Get the patch coordinates and shape of the bbox + small_grid_coords, bbox_h, bbox_w = bbox_to_patch_coords(bbox_coordinates, patch_size=patch_size, return_shape=True) + small_grid_coords = np.array(small_grid_coords) + + indices_array = img_ids.cpu().numpy().squeeze().copy() + indices_array = indices_array.reshape((h*w,-1)) + + # indices_array = indices_array.reshape((h*w,-1)) + + # Compute bbox center in patch coordinates + min_h, min_w = np.min(small_grid_coords, axis=0) + max_h, max_w = np.max(small_grid_coords, axis=0) + center_y = (min_h + max_h) // 2 + center_x = (min_w + max_w) // 2 + + min_h_p, min_w_p = max(min_h-2, 0), max(min_w-2, 0) + max_h_p, max_w_p = min(h, max_h+2), min(w, max_w+2) + + # Compute shortest distance from center to bbox edge + radius_y = min(center_y - min_h, max_h - center_y) + radius_x = min(center_x - min_w, max_w - center_x) + radius = min(radius_y, radius_x) + + # Get all valid coordinates in the grid + yy, xx = np.meshgrid(np.arange(h), np.arange(w), indexing='ij') + all_coords = np.stack([yy.ravel(), xx.ravel()], axis=1) + + # Exclude bbox coordinates + bbox_set = set(map(tuple, small_grid_coords)) + filtered_coords = [tuple(coord) for coord in all_coords if tuple(coord) not in bbox_set] + + # result_coords=[] + # for center_y, center_x in small_grid_coords: + # neighbors = [] + # for dy in range(-radius-1, radius + 2): + # for dx in range(-radius-1, radius + 2): + # if abs(dy) + abs(dx) <= radius+2: + # ny, nx = center_y + dy, center_x + dx + # if 0 <= ny < h and 0 <= nx < w: + # if (ny, nx) not in bbox_set: + # neighbors.append((ny, nx)) + # result_coords.append(neighbors) + + # return small_grid_coords.tolist(), result_coords + + result=[] + for center_y, center_x in small_grid_coords: + neighbors = [] + for dy in range(-radius-1, radius + 2): + for dx in range(-radius-1, radius + 2): + if abs(dy) + abs(dx) <= radius+1: + ny, nx = center_y + dy, center_x + dx + if min_h_p <= ny < max_h_p and min_w_p <= nx < max_w_p: + if (ny, nx) not in bbox_set: + neighbors.append((ny, nx)) + # if len(neighbors) == 0: + # import pdb;pdb.set_trace() + neighbors_ind = indices_array[[patch_coor_to_ind(x,y,w,0) for (y,x) in neighbors],:] + result.append(neighbors_ind.mean(0)) + + return torch.from_numpy(np.array(result)).unsqueeze(0) + +def perturb_pe(h, w, bbox_coordinates, img_ids, patch_size=16, txt_len=512): + patch_ids = bbox_to_patch_indices(bbox_coordinates, h, w, patch_size, txt_len=0) + indices_array = torch.from_numpy(img_ids.cpu().numpy().squeeze().copy()[patch_ids, :]) + noise = torch.randn_like(indices_array) + # Mask for non-zero elements + nonzero_mask = indices_array != 0 + # Clone indices_array to preserve original + perturbed_indices_array = indices_array.clone() + # Apply noise only to non-zero elements + perturbed_indices_array[nonzero_mask] += noise[nonzero_mask] + # perturbed_indices_array = indices_array + torch.randn_like(indices_array) * 0.3 + return perturbed_indices_array.unsqueeze(0) + +def shuffle_pe(h, w, patch_ids, patch_size=16, txt_len=512, intensity=3): + bbox_coords = [patch_ind_to_coor(ind, w, txt_len) for ind in patch_ids] + # bbox_coords = bbox_to_patch_coords(bbox_coordinates, patch_size=patch_size) + + shuffled_coords = [] + + for y, x in bbox_coords: + dx = torch.randint(-intensity, intensity + 1, (1,)).item() + dy = torch.randint(-intensity, intensity + 1, (1,)).item() + # dx = 0 + # dy = 0 + + new_x = max(0, min(x + dx, w - 1)) + new_y = max(0, min(y + dy, h - 1)) + + shuffled_coords.append((new_y, new_x)) + + return [patch_coor_to_ind(x,y,w,txt_len) for y, x in shuffled_coords] + +def sample_closest_patch_ind(h, w, patch_indices, reference_patch_indices, patch_size=16, txt_len=512): + """ + Sample closest patches for arbitrary shape with randomization. + + Args: + h, w: Patch grid dimensions + patch_indices: List of patch indices defining the shape + reference_patch_indices: List of patch indices to use as reference/candidates + patch_size: Size of each patch + txt_len: Text length offset + + Returns: + List of sampled closest patch indices + """ + shape_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in patch_indices]) + + # Convert reference patch indices to coordinates + reference_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in reference_patch_indices]) + + result = [] + for y, x in shape_coords: + if 0 <= y < h and 0 <= x < w: + # Calculate distances only to reference patch coordinates + distances = np.abs(reference_coords[:, 0] - y) + np.abs(reference_coords[:, 1] - x) + inv_d = 1.0 / (distances + 1e-8) + inv_d = np.pow(inv_d, 2) + p_weight = inv_d / np.sum(inv_d) + idx = np.random.choice(len(distances), p=p_weight) + closest_coord = tuple(reference_coords[idx]) + result.append(patch_coor_to_ind(closest_coord[1], closest_coord[0], w, txt_len)) + + return result + +def get_closest_patch_coords(target_coords, reference_coords): + """ + Map each target coordinate to its closest reference coordinate. + + Args: + target_coords: List of (y, x) coordinates that need to be mapped + reference_coords: List of (y, x) coordinates to use as reference/candidates + + Returns: + List of closest reference coordinates for each target coordinate + """ + target_coords = np.array(target_coords) + reference_coords = np.array(reference_coords) + + result = [] + for ty, tx in target_coords: + # Calculate Manhattan distances to all reference coordinates + distances = np.abs(reference_coords[:, 0] - ty) + np.abs(reference_coords[:, 1] - tx) + min_idx = np.argmin(distances) + closest_coord = tuple(reference_coords[min_idx]) + result.append(closest_coord) + + return result + + +def get_closest_patch_inds(h, w, target_patch_indices, reference_patch_indices, txt_len=512): + """ + Map each target patch index to the closest reference patch index using Manhattan distance. + + Args: + h, w: Patch grid dimensions (not used directly but kept for API symmetry) + target_patch_indices: List of patch indices to map + reference_patch_indices: List of candidate reference patch indices + txt_len: Text length offset used in index<->coord conversions + + Returns: + List of closest reference patch indices corresponding to each target patch index + """ + if len(target_patch_indices) == 0 or len(reference_patch_indices) == 0: + return [] + + # Convert indices to (y, x) coordinates + target_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in target_patch_indices]) + reference_coords = np.array([patch_ind_to_coor(ind, w, txt_len) for ind in reference_patch_indices]) + + result = [] + for ty, tx in target_coords: + distances = np.abs(reference_coords[:, 0] - ty) + np.abs(reference_coords[:, 1] - tx) + min_idx = int(np.argmin(distances)) + result.append(reference_patch_indices[min_idx]) + + return result diff --git a/ArtiAgent - DefectFill/src/flux/math.py b/ArtiAgent - DefectFill/src/flux/math.py new file mode 100644 index 0000000000000000000000000000000000000000..a1f0a62c0d63a5425551e5bc411ea86fc47c53ee --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/math.py @@ -0,0 +1,54 @@ +import math +import torch +from einops import rearrange +from torch import Tensor + + +def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor: + q, k = apply_rope(q, k, pe) + + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "B H L D -> B L (H D)") + + return x + +def attention_masked(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, patch_ids: list[int], mask: Tensor, return_weight:bool=False) -> Tensor: + q, k = apply_rope(q, k, pe) + # x = torch.nn.functional.scaled_dot_product_attention(q,k,v,attn_mask=attn_mask) + if return_weight: + x, m = scaled_dot_product_attention_masked(q, k, v, patch_ids, mask, return_weight) + x = rearrange(x, "B H L D -> B L (H D)") + return x, m + else: + x = scaled_dot_product_attention_masked(q, k, v, patch_ids, mask, return_weight) + x = rearrange(x, "B H L D -> B L (H D)") + return x + + # return x, m + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + assert dim % 2 == 0 + scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = torch.einsum("...n,d->...nd", pos, omega) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2) + return out.float() + + +def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) + + +def scaled_dot_product_attention_masked(query, key, value, patch_ids, attn_mask, return_weight=False): + scale_factor = 1 / math.sqrt(query.size(-1)) + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_mask + attn_weight = torch.softmax(attn_weight, dim=-1) + if return_weight: + return attn_weight @ value, attn_weight + return attn_weight @ value diff --git a/ArtiAgent - DefectFill/src/flux/model.py b/ArtiAgent - DefectFill/src/flux/model.py new file mode 100644 index 0000000000000000000000000000000000000000..582990668e435b82472d4eb6c6f0fab541d116fa --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/model.py @@ -0,0 +1,249 @@ +from dataclasses import dataclass + +import torch +from torch import Tensor, nn +import numpy as np + +from flux.modules.layers import (DoubleStreamBlock, EmbedND, LastLayer, + MLPEmbedder, SingleStreamBlock, + timestep_embedding) + + +@dataclass +class FluxParams: + in_channels: int + out_channels: int + vec_in_dim: int + context_in_dim: int + hidden_size: int + mlp_ratio: float + num_heads: int + depth: int + depth_single_blocks: int + axes_dim: list[int] + theta: int + qkv_bias: bool + guidance_embed: bool + + +class Flux(nn.Module): + """ + Transformer model for flow matching on sequences. + """ + + def __init__(self, params: FluxParams): + super().__init__() + + self.params = params + self.in_channels = params.in_channels + self.out_channels = params.out_channels + if params.hidden_size % params.num_heads != 0: + raise ValueError( + f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}" + ) + pe_dim = params.hidden_size // params.num_heads + if sum(params.axes_dim) != pe_dim: + raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}") + self.hidden_size = params.hidden_size + self.num_heads = params.num_heads + self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim) + self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) + self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size) + self.guidance_in = ( + MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + ) + self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size) + + self.double_blocks = nn.ModuleList( + [ + DoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + ) + for _ in range(params.depth) + ] + ) + + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio) + for _ in range(params.depth_single_blocks) + ] + ) + + self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels) + self._sequential_offload = False + + def enable_sequential_cpu_offload(self): + self._sequential_offload = True + + def forward( + self, + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor | None = None, + info = None, + ref_img: Tensor | None = None, # โ† NEW + ref_img_ids: Tensor | None = None, # โ† NEW + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + # Ensure inputs match the model's dtype (NF4 can silently upcast to float32) + target_dtype = self.img_in.weight.dtype + img = img.to(target_dtype) + txt = txt.to(target_dtype) + if y.dtype != target_dtype: + y = y.to(target_dtype) + + # running on sequences img + img = self.img_in(img) + original_img_seq_len = img.shape[1] # โ† REMEMBER: how many original img tokens + + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + if guidance is None: + raise ValueError("Didn't get guidance strength for guidance distilled model.") + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # NEW: Concatenate reference tokens into image stream + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + if ref_img is not None and ref_img_ids is not None: + ref = self.img_in(ref_img) # project reference patches same way + img = torch.cat([img, ref], dim=1) + img_ids = torch.cat([img_ids, ref_img_ids], dim=1) + + if ref_img is not None: + print(f"[Flux.forward] Attending to {ref_img.shape[1]} ref tokens + {original_img_seq_len} img tokens") + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + inject_pe = pe.clone() + if not info['inverse']: + # Defensive clamp: GSAM indices may exceed seq_len for non-16-divisible images + seq_len = pe.shape[2] + + # Initialize accumulated lists for tracking all processed IDs + accumulated_target_ids = [] + accumulated_ref_ids = [] + for artifact_data in info['artifact_data']: + if artifact_data['artifact_type'] == 'addition' and info['addition']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + + # Accumulate IDs + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + elif artifact_data['artifact_type'] == 'removal' and info['removal']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + + # ref_ids = get_closest_patch_inds(info['patch_h'], info['patch_w'], target_ids, ref_ids) + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + # Accumulate IDs (after target_ids modification) + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + elif artifact_data['artifact_type'] == 'distortion' and info['distortion']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + + if len(ref_ids) == 0: + # For distortion with no reference patches, shuffle target patches + ref_ids = target_ids.copy() + np.random.shuffle(ref_ids) + # Ensure target_ids and ref_ids are different for distortion + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + # Accumulate IDs (after any ref_ids modification) + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + elif artifact_data['artifact_type'] == 'fusion' and info['fusion']: + ref_ids = artifact_data['reference_patch_indices'].copy() + target_ids = artifact_data['target_patch_indices'].copy() + ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids] + target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids] + + # np.random.shuffle(ref_ids) + if len(target_ids) > 0 and len(ref_ids) > 0: + inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:] + + # Accumulate IDs + if info['inject']: + accumulated_target_ids.extend(target_ids) + accumulated_ref_ids.extend(ref_ids) + + info['patch_ids'] = accumulated_target_ids + info['patch_ref_ids'] = accumulated_ref_ids + info['timesteps'] = timesteps + + + if self._sequential_offload: + for block in self.double_blocks: + block = block.to(img.device) + img, txt = block(img=img, txt=txt, vec=vec, pe=inject_pe, info=info) + block = block.cpu() + torch.cuda.empty_cache() + else: + for block in self.double_blocks: + img, txt = block(img=img, txt=txt, vec=vec, pe=inject_pe, info=info) + + cnt = 0 + img = torch.cat((txt, img), 1) + info['type'] = 'single' + if self._sequential_offload: + for block in self.single_blocks: + block = block.to(img.device) + info['id'] = cnt + if cnt < 19: + img, info = block(img, vec=vec, pe=inject_pe, info=info) + else: + img, info = block(img, vec=vec, pe=pe, info=info) + block = block.cpu() + torch.cuda.empty_cache() + cnt += 1 + else: + for block in self.single_blocks: + info['id'] = cnt + if cnt < 19: + img, info = block(img, vec=vec, pe=inject_pe, info=info) + else: + img, info = block(img, vec=vec, pe=pe, info=info) + cnt += 1 + + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # MODIFIED: Extract only ORIGINAL img tokens + # Before: img = img[:, txt.shape[1] :, ...] (gets img + ref) + # After: img = img[:, txt.shape[1] : txt.shape[1] + original_img_seq_len, ...] + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + img = img[:, txt.shape[1] : txt.shape[1] + original_img_seq_len, ...] + + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + return img, info \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/flux/modules/__pycache__/autoencoder.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/modules/__pycache__/autoencoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c340ba212c90e5afb1a6e4ca2bcb5bbd35b4eba Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/modules/__pycache__/autoencoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/modules/__pycache__/conditioner.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/modules/__pycache__/conditioner.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..658ee8c803d824ca7c62beef87d4bb4264b04c73 Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/modules/__pycache__/conditioner.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/modules/__pycache__/layers.cpython-310.pyc b/ArtiAgent - DefectFill/src/flux/modules/__pycache__/layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2a6b4c42c2cfffa108a5eeae740e00fb71fb517 Binary files /dev/null and b/ArtiAgent - DefectFill/src/flux/modules/__pycache__/layers.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/flux/modules/autoencoder.py b/ArtiAgent - DefectFill/src/flux/modules/autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..86bdec01bd09c872721fe267fe1bd83d32d5fdec --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/modules/autoencoder.py @@ -0,0 +1,313 @@ +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn + + +@dataclass +class AutoEncoderParams: + resolution: int + in_channels: int + ch: int + out_ch: int + ch_mult: list[int] + num_res_blocks: int + z_channels: int + scale_factor: float + shift_factor: float + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, h, w = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + + return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + # no asymmetric padding in torch conv, must do it ourselves + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: Tensor): + pad = (0, 1, 0, 1) + x = nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + return x + + +class Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor): + x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + x = self.conv(x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + ch: int, + ch_mult: list[int], + num_res_blocks: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + # downsampling + self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) + + curr_res = resolution + in_ch_mult = (1,) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + block_in = self.ch + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + ch: int, + out_ch: int, + ch_mult: list[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.ffactor = 2 ** (self.num_resolutions - 1) + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = ch * ch_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + + # z to block_in + self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: Tensor) -> Tensor: + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class DiagonalGaussian(nn.Module): + def __init__(self, sample: bool = True, chunk_dim: int = 1): + super().__init__() + self.sample = sample + self.chunk_dim = chunk_dim + + def forward(self, z: Tensor) -> Tensor: + mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim) + # import pdb;pdb.set_trace() + if self.sample: + std = torch.exp(0.5 * logvar) + return mean #+ std * torch.randn_like(mean) + else: + return mean + + +class AutoEncoder(nn.Module): + def __init__(self, params: AutoEncoderParams): + super().__init__() + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian() + + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: Tensor) -> Tensor: + z = self.reg(self.encoder(x)) + z = self.scale_factor * (z - self.shift_factor) + return z + + def decode(self, z: Tensor) -> Tensor: + z = z / self.scale_factor + self.shift_factor + return self.decoder(z) + + def forward(self, x: Tensor) -> Tensor: + return self.decode(self.encode(x)) diff --git a/ArtiAgent - DefectFill/src/flux/modules/conditioner.py b/ArtiAgent - DefectFill/src/flux/modules/conditioner.py new file mode 100644 index 0000000000000000000000000000000000000000..98dbffd132a61cb74f1cc50ad3405d2cd58f3268 --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/modules/conditioner.py @@ -0,0 +1,64 @@ +import torch +from torch import Tensor, nn +from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer, BitsAndBytesConfig) + + +class HFEmbedder(nn.Module): + def __init__(self, version: str, max_length: int, is_clip, **hf_kwargs): + super().__init__() + self.is_clip = is_clip + self.max_length = max_length + self.output_key = "pooler_output" if self.is_clip else "last_hidden_state" + + # Safely remove 'load_in_8bit' and 'device_map' from hf_kwargs so they don't get passed to __init__ + self.is_8bit = hf_kwargs.pop("load_in_8bit", False) + device_map = hf_kwargs.pop("device_map", "cuda") + + if self.is_clip: + self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs) + else: + self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length) + if self.is_8bit: + # Use BitsAndBytesConfig for modern transformers + # Remove torch_dtype conflict if present in kwargs + hf_kwargs.pop("torch_dtype", None) + q_config = BitsAndBytesConfig(load_in_8bit=True) + # Remove torch_dtype conflict if present in hf_kwargs + hf_kwargs.pop("torch_dtype", None) + + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained( + version, + quantization_config=q_config, + device_map=hf_kwargs.pop("device_map", "cuda"), + **hf_kwargs + ) + else: + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs) + + self.hf_module = self.hf_module.eval().requires_grad_(False) + + def to(self, *args, **kwargs): + # If loaded in 8-bit, bitsandbytes handles device placement automatically. + # Calling .to() on an 8-bit model will crash, so we skip it. + if self.is_8bit: + return self + return super().to(*args, **kwargs) + + def forward(self, text: list[str]) -> Tensor: + batch_encoding = self.tokenizer( + text, + truncation=True, + max_length=self.max_length, + return_length=False, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + + outputs = self.hf_module( + input_ids=batch_encoding["input_ids"].to(self.hf_module.device), + attention_mask=None, + output_hidden_states=False, + ) + return outputs[self.output_key].to(torch.bfloat16) diff --git a/ArtiAgent - DefectFill/src/flux/modules/layers.py b/ArtiAgent - DefectFill/src/flux/modules/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..3ecc8b7600e0e02eb74d163965cee5f0349acbf5 --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/modules/layers.py @@ -0,0 +1,288 @@ +import math +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn +import random + +from flux.math import attention, attention_masked, rope + +import os + +class EmbedND(nn.Module): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + n_axes = ids.shape[-1] + emb = torch.cat( + [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], + dim=-3, + ) + + return emb.unsqueeze(1) + + +def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + t = time_factor * t + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + t.device + ) + + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(t) + return embedding + + +class MLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class RMSNorm(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor): + x_dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=x_dtype) * self.scale.to(dtype=x_dtype) + + +class QKNorm(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = RMSNorm(dim) + self.key_norm = RMSNorm(dim) + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: + q = self.query_norm(q) + k = self.key_norm(k) + return q.to(v), k.to(v) + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.norm = QKNorm(head_dim) + self.proj = nn.Linear(dim, dim) + + def forward(self, x: Tensor, pe: Tensor, patch_ids=None) -> Tensor: + qkv = self.qkv(x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + if patch_ids is None: + x = attention(q, k, v, pe=pe) + else: + x = attention_masked(q, k, v, pe=pe, patch_ids=patch_ids) + + x = self.proj(x) + return x + + +@dataclass +class ModulationOut: + shift: Tensor + scale: Tensor + gate: Tensor + + +class Modulation(nn.Module): + def __init__(self, dim: int, double: bool): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) + + def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1) + + return ( + ModulationOut(*out[:3]), + ModulationOut(*out[3:]) if self.is_double else None, + ) + + +class DoubleStreamBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False): + super().__init__() + + mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.num_heads = num_heads + self.hidden_size = hidden_size + self.img_mod = Modulation(hidden_size, double=True) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + self.txt_mod = Modulation(hidden_size, double=True) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, info) -> tuple[Tensor, Tensor]: + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + # prepare image for attention + img_modulated = self.img_norm1(img) + img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift + img_qkv = self.img_attn.qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) + + # prepare txt for attention + txt_modulated = self.txt_norm1(txt) + txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift + txt_qkv = self.txt_attn.qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) + + # run actual attention + q = torch.cat((txt_q, img_q), dim=2) #[8, 24, 512, 128] + [8, 24, 900, 128] -> [8, 24, 1412, 128] + k = torch.cat((txt_k, img_k), dim=2) + v = torch.cat((txt_v, img_v), dim=2) + + attn = attention(q, k, v, pe=pe) + + txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + + # calculate the txt bloks + txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) + return img, txt + + +class SingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + qk_scale: float | None = None, + ): + super().__init__() + self.hidden_dim = hidden_size + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.scale = qk_scale or head_dim**-0.5 + + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + # qkv and mlp_in + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) + # proj and mlp_out + self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) + + self.norm = QKNorm(head_dim) + + self.hidden_size = hidden_size + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + self.mlp_act = nn.GELU(approximate="tanh") + self.modulation = Modulation(hidden_size, double=False) + + def forward(self, x: Tensor, vec: Tensor, pe: Tensor, info) -> Tensor: + mod, _ = self.modulation(vec) + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + + # Note: If the memory of your device is not enough, you may consider uncommenting the following code. + # if info['inject'] and info['id'] > 19: + # store_path = os.path.join(info['feature_path'], str(info['t']) + '_' + str(info['second_order']) + '_' + str(info['id']) + '_' + info['type'] + '_' + 'V' + '.pth') + # if info['inverse']: + # torch.save(v, store_path) + # if not info['inverse']: + # v = torch.load(store_path, weights_only=True) + + # Save the features in the memory + + if info['inject'] and info['id'] > 19: + feature_name = str(info['t']) + '_' + str(info['second_order']) + '_' + str(info['id']) + '_' + info['type'] + '_' + 'V' + if info['inverse']: + info['feature'][feature_name] = v.clone() + else: + # Try to load feature, but continue if it doesn't exist + if feature_name in info['feature']: + v = info['feature'][feature_name] + + num_patches = v.size(2) + mask = torch.ones(num_patches, dtype=torch.bool, device=v.device) + mask[info['patch_ids']] = False + mask[:512] = False + keep_indices = mask.nonzero(as_tuple=True)[0] + if feature_name in info['feature']: + backup = info['feature'][feature_name] + v[:, :, keep_indices, :] = backup[:, :, keep_indices, :] + if len(info['patch_ref_ids']) != 0: + v[:, :, info['patch_ids'], :] = v[:, :, info['patch_ref_ids'], :] + + attn = attention(q, k, v, pe=pe) + # compute activation in mlp stream, cat again and run second linear layer + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output, info + + +class LastLayer(nn.Module): + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x: Tensor, vec: Tensor) -> Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + x = self.linear(x) + return x diff --git a/ArtiAgent - DefectFill/src/flux/sampling.py b/ArtiAgent - DefectFill/src/flux/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..f6bb29287c2bf83d19bd268d398bbd4e6b65ecac --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/sampling.py @@ -0,0 +1,390 @@ +import math +from typing import Callable + +import torch +from einops import rearrange, repeat +from torch import Tensor + +from .model import Flux +from .modules.conditioner import HFEmbedder + +def prepare(t5: HFEmbedder, clip: HFEmbedder, img: Tensor, prompt: str | list[str], + info=None) -> dict[str, Tensor]: + """ + Prepare inputs for the flux model with support for patch indices. + + Args: + t5, clip: Text encoders + img: Input image tensor + prompt: Text prompt(s) + info: Additional information dictionary, must contain 'artifact_data'. + + Returns: + Dictionary containing prepared inputs + """ + bs, c, h, w = img.shape + if bs == 1 and not isinstance(prompt, str): + bs = len(prompt) + img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + if img.shape[0] == 1 and bs > 1: + img = repeat(img, "1 ... -> bs ...", bs=bs) + + img_ids = torch.zeros(h // 2, w // 2, 3) + img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None] + img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :] + img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) + if isinstance(prompt, str): + prompt = [prompt] + txt = t5(prompt) + if txt.shape[0] == 1 and bs > 1: + txt = repeat(txt, "1 ... -> bs ...", bs=bs) + txt_ids = torch.zeros(bs, txt.shape[1], 3) + + vec = clip(prompt) + if vec.shape[0] == 1 and bs > 1: + vec = repeat(vec, "1 ... -> bs ...", bs=bs) + + patch_h, patch_w = h // 2, w // 2 + + # Add patch dimensions to info for model to use + info['patch_h'] = patch_h + info['patch_w'] = patch_w + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # NEW: Patchify reference latents for RAG visual conditioning + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + if info is not None and info.get('reference_latents'): + ref_list = info['reference_latents'] # list of [B, 16, H, W] tensors + ref_tokens_all = [] + ref_ids_all = [] + + for ref_lat in ref_list: + ref_bs, ref_c, ref_h, ref_w = ref_lat.shape + # Patchify same way as img: [B, 16, H, W] -> [B, (H/2)*(W/2), 64] + ref_p = rearrange(ref_lat, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + if ref_p.shape[0] == 1 and bs > 1: + ref_p = repeat(ref_p, "1 ... -> bs ...", bs=bs) + + # Position IDs matching img_ids pattern + ref_ids = torch.zeros(ref_h // 2, ref_w // 2, 3) + ref_ids[..., 1] = ref_ids[..., 1] + torch.arange(ref_h // 2)[:, None] + ref_ids[..., 2] = ref_ids[..., 2] + torch.arange(ref_w // 2)[None, :] + ref_ids = repeat(ref_ids, "h w c -> b (h w) c", b=ref_p.shape[0]) + + ref_tokens_all.append(ref_p) + ref_ids_all.append(ref_ids) + + # Store in info dict (will be added to return dict below) + info['ref_img'] = torch.cat(ref_tokens_all, dim=1).to(img.device, dtype=torch.bfloat16) + info['ref_img_ids'] = torch.cat(ref_ids_all, dim=1).to(img.device) + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + result = { + "img": img, + "img_ids": img_ids.to(img.device), + "txt": txt.to(device=img.device, dtype=torch.bfloat16), # <--- Cast to bfloat16 + "txt_ids": txt_ids.to(img.device), + "vec": vec.to(device=img.device, dtype=torch.bfloat16), # <--- Cast to bfloat16 + } + + # Add reference tensors if they were computed above + if info is not None and 'ref_img' in info: + result["ref_img"] = info['ref_img'] + result["ref_img_ids"] = info['ref_img_ids'] + + if "ref_img" in result: + print(f"[FLUX prepare] ref_img shape: {result['ref_img'].shape}") + + return result, (patch_h, patch_w) + + +def time_shift(mu: float, sigma: float, t: Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + +def get_lin_function( + x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15 +) -> Callable[[float], float]: + m = (y2 - y1) / (x2 - x1) + b = y1 - m * x1 + return lambda x: m * x + b + + +def get_schedule( + num_steps: int, + image_seq_len: int, + base_shift: float = 0.5, + max_shift: float = 1.15, + shift: bool = True, +) -> list[float]: + # extra step for zero + timesteps = torch.linspace(1, 0, num_steps + 1) + + # shifting the schedule to favor high timesteps for higher signal images + if shift: + # estimate mu based on linear estimation between two points + mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len) + timesteps = time_shift(mu, 1.0, timesteps) + + return timesteps.tolist() + +def denoise_first_order( + model: Flux, + # model input + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + vec: Tensor, + # sampling parameters + timesteps: list[float], + inverse, + info, + percentage_of_steps = 1.0, + guidance: float = 5.0, + ref_img: Tensor | None = None, # โ† ADD + ref_img_ids: Tensor | None = None, # โ† ADD +): + # this is ignored for schnell + inject_list = [True] * int(info['inject_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['inject_step'])) + attn_mask_list = [True] * int(info['attn_mask_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['attn_mask_step'])) + + # PE step lists for each artifact type + pe_step_addition_list = [True] * int(info['pe_step_addition']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_addition'])) + pe_step_removal_list = [True] * int(info['pe_step_removal']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_removal'])) + pe_step_distortion_list = [True] * int(info['pe_step_distortion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_distortion'])) + pe_step_fusion_list = [True] * int(info['pe_step_fusion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_fusion'])) + if inverse: + timesteps = timesteps[::-1] + inject_list = inject_list[::-1] + if percentage_of_steps != 1: + end_timestep_idx = int(len(timesteps) * percentage_of_steps) + if inverse: + timesteps = timesteps[:end_timestep_idx] + # inject_list = inject_list[:end_timestep_idx - 1] + else: + timesteps = timesteps[len(timesteps) - end_timestep_idx:] + # inject_list = inject_list[len(inject_list) - end_timestep_idx + 1:] + + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) + for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) + info['t'] = t_prev if inverse else t_curr + info['inverse'] = inverse + info['second_order'] = False + info['inject'] = inject_list[i] + info['attn_mask'] = attn_mask_list[i] + info['addition'] = pe_step_addition_list[i] + info['removal'] = pe_step_removal_list[i] + info['distortion'] = pe_step_distortion_list[i] + info['fusion'] = pe_step_fusion_list[i] + + pred, info = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec, + guidance=guidance_vec, + info=info, + # โ•โ•โ• ADD THESE TWO LINES โ•โ•โ• + ref_img=ref_img, + ref_img_ids=ref_img_ids, + ) + + img = img + (t_prev - t_curr) * pred + return img, info + +def denoise_fireflow( + model: Flux, + # model input + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + vec: Tensor, + # sampling parameters + timesteps: list[float], + inverse, + info, + percentage_of_steps = 1.0, + guidance: float = 5.0, + ref_img: Tensor | None = None, # โ† ADD + ref_img_ids: Tensor | None = None, # โ† ADD + ): + # this is ignored for schnell + inject_list = [True] * int(info['inject_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['inject_step'])) + attn_mask_list = [True] * int(info['attn_mask_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['attn_mask_step'])) + + # PE step lists for each artifact type + pe_step_addition_list = [True] * int(info['pe_step_addition']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_addition'])) + pe_step_removal_list = [True] * int(info['pe_step_removal']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_removal'])) + pe_step_distortion_list = [True] * int(info['pe_step_distortion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_distortion'])) + pe_step_fusion_list = [True] * int(info['pe_step_fusion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_fusion'])) + + if inverse: + timesteps = timesteps[::-1] + inject_list = inject_list[::-1] + if percentage_of_steps != 1: + end_timestep_idx = int(len(timesteps) * percentage_of_steps) + if inverse: + timesteps = timesteps[:end_timestep_idx] + # inject_list = inject_list[:end_timestep_idx - 1] + else: + timesteps = timesteps[len(timesteps) - end_timestep_idx:] + # inject_list = inject_list[len(inject_list) - end_timestep_idx + 1:] + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) + + step_list = [] + next_step_velocity = None + for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) + info['t'] = t_prev if inverse else t_curr + info['inverse'] = inverse + info['second_order'] = False + info['inject'] = inject_list[i] + info['attn_mask'] = attn_mask_list[i] + info['addition'] = pe_step_addition_list[i] + info['removal'] = pe_step_removal_list[i] + info['distortion'] = pe_step_distortion_list[i] + info['fusion'] = pe_step_fusion_list[i] + + if next_step_velocity is None: + pred, info = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec, + guidance=guidance_vec, + info=info, + # โ•โ•โ• ADD THESE TWO LINES โ•โ•โ• + ref_img=ref_img, + ref_img_ids=ref_img_ids, + ) + else: + pred = next_step_velocity + + img_mid = img + (t_prev - t_curr) / 2 * pred + + t_vec_mid = torch.full((img.shape[0],), t_curr + (t_prev - t_curr) / 2, dtype=img.dtype, device=img.device) + info['second_order'] = True + pred_mid, info = model( + img=img_mid, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec_mid, + guidance=guidance_vec, + info=info, + ref_img=ref_img, + ref_img_ids=ref_img_ids + ) + next_step_velocity = pred_mid + + img = img + (t_prev - t_curr) * pred_mid + + return img, info + + +def denoise( + model: Flux, + # model input + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + vec: Tensor, + # sampling parameters + timesteps: list[float], + inverse, + info, + percentage_of_steps = 1.0, + guidance: float = 4.0, + ref_img: Tensor | None = None, # โ† ADD + ref_img_ids: Tensor | None = None, # โ† ADD +): + # this is ignored for schnell + inject_list = [True] * int(info['inject_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['inject_step'])) + attn_mask_list = [True] * int(info['attn_mask_step']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['attn_mask_step'])) + + # PE step lists for each artifact type + pe_step_addition_list = [True] * int(info['pe_step_addition']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_addition'])) + pe_step_removal_list = [True] * int(info['pe_step_removal']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_removal'])) + pe_step_distortion_list = [True] * int(info['pe_step_distortion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_distortion'])) + pe_step_fusion_list = [True] * int(info['pe_step_fusion']) + [False] * (int(len(timesteps) * percentage_of_steps) -1 - int(info['pe_step_fusion'])) + + if inverse: + timesteps = timesteps[::-1] + inject_list = inject_list[::-1] + + if percentage_of_steps != 1: + end_timestep_idx = int(len(timesteps) * percentage_of_steps) + if inverse: + timesteps = timesteps[:end_timestep_idx] + # inject_list = inject_list[:end_timestep_idx - 1] + else: + timesteps = timesteps[len(timesteps) - end_timestep_idx:] + # inject_list = inject_list[len(inject_list) - end_timestep_idx + 1:] + + guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) + for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) + info['t'] = t_prev if inverse else t_curr + info['inverse'] = inverse + info['second_order'] = False + info['inject'] = inject_list[i] + info['attn_mask'] = attn_mask_list[i] + info['addition'] = pe_step_addition_list[i] + info['removal'] = pe_step_removal_list[i] + info['distortion'] = pe_step_distortion_list[i] + info['fusion'] = pe_step_fusion_list[i] + + pred, info = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec, + guidance=guidance_vec, + info=info, + # โ•โ•โ• ADD THESE TWO LINES โ•โ•โ• + ref_img=ref_img, + ref_img_ids=ref_img_ids + ) + + + img_mid = img + (t_prev - t_curr) / 2 * pred + + t_vec_mid = torch.full((img.shape[0],), (t_curr + (t_prev - t_curr) / 2), dtype=img.dtype, device=img.device) + info['second_order'] = True + pred_mid, info = model( + img=img_mid, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + timesteps=t_vec_mid, + guidance=guidance_vec, + info=info + ) + + first_order = (pred_mid - pred) / ((t_prev - t_curr) / 2) + img = img + (t_prev - t_curr) * pred + 0.5 * (t_prev - t_curr) ** 2 * first_order + + return img, info + + +def unpack(x: Tensor, height: int, width: int) -> Tensor: + return rearrange( + x, + "b (h w) (c ph pw) -> b c (h ph) (w pw)", + h=math.ceil(height / 16), + w=math.ceil(width / 16), + ph=2, + pw=2, + ) diff --git a/ArtiAgent - DefectFill/src/flux/util.py b/ArtiAgent - DefectFill/src/flux/util.py new file mode 100644 index 0000000000000000000000000000000000000000..c3f8e8092462a72bcb75af58e716f1b87a6751d8 --- /dev/null +++ b/ArtiAgent - DefectFill/src/flux/util.py @@ -0,0 +1,346 @@ +import os +from dataclasses import dataclass + +import torch +from einops import rearrange +from huggingface_hub import hf_hub_download +# from imwatermark import WatermarkEncoder +from safetensors.torch import load_file as load_sft + +from flux.model import Flux, FluxParams +from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams +from flux.modules.conditioner import HFEmbedder +from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel, + T5Tokenizer, BitsAndBytesConfig) # <--- Added BitsAndBytesConfig + + +@dataclass +class ModelSpec: + params: FluxParams + ae_params: AutoEncoderParams + ckpt_path: str | None + ae_path: str | None + repo_id: str | None + repo_flow: str | None + repo_ae: str | None + +configs = { + "flux-dev": ModelSpec( + repo_id="black-forest-labs/FLUX.1-dev", + repo_flow="flux1-dev.safetensors", + repo_ae=None, + ckpt_path=os.getenv("FLUX_DEV"), + params=FluxParams( + in_channels=64, + out_channels=64, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=True, + ), + ae_path=os.getenv("AE"), + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), + "flux-fill-dev": ModelSpec( + repo_id="black-forest-labs/FLUX.1-Fill-dev", + repo_flow="flux1-fill-dev.safetensors", + repo_ae="ae.safetensors", + ckpt_path=os.getenv("FLUX_FILL_DEV"), + params=FluxParams( + in_channels=64, + out_channels=384, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=True, + ), + ae_path=os.getenv("AE"), + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), + "flux-kontext-dev": ModelSpec( + repo_id="black-forest-labs/FLUX.1-Kontext-dev", + repo_flow="flux1-kontext-dev.safetensors", + repo_ae="ae.safetensors", + ckpt_path=os.getenv("FLUX_FILL_DEV"), + params=FluxParams( + in_channels=64, + out_channels=64, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10_000, + qkv_bias=True, + guidance_embed=True, + ), + ae_path=os.getenv("AE"), + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), + "flux-schnell": ModelSpec( + repo_id="black-forest-labs/FLUX.1-schnell", + repo_flow="flux1-schnell.safetensors", + repo_ae="black-forest-labs/FLUX.1-schnell", + ckpt_path=os.getenv("FLUX_SCHNELL"), + params=FluxParams( + in_channels=64, # ArtiAgent custom input dimension logic handled in load_flow_model + out_channels=64, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=3072, + mlp_ratio=4.0, + num_heads=24, + depth=19, + depth_single_blocks=38, + axes_dim=[16, 56, 56], + theta=10000.0, + qkv_bias=True, + guidance_embed=False, + ), + ae_path="ae.safetensors", + ae_params=AutoEncoderParams( + resolution=256, + in_channels=3, + ch=128, + out_ch=3, + ch_mult=[1, 2, 4, 4], + num_res_blocks=2, + z_channels=16, + scale_factor=0.3611, + shift_factor=0.1159, + ), + ), +} + + +def print_load_warning(missing: list[str], unexpected: list[str]) -> None: + if len(missing) > 0 and len(unexpected) > 0: + print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing)) + print("\n" + "-" * 79 + "\n") + print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected)) + elif len(missing) > 0: + print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing)) + elif len(unexpected) > 0: + print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected)) + +def _replace_linear_with_4bit(module, compute_dtype=torch.bfloat16): + """Recursively replace all nn.Linear with bitsandbytes 4-bit layers""" + import bitsandbytes as bnb + for name, child in module.named_children(): + if name == "img_in": + continue # Skip img_in to preserve ArtiAgent's custom shape handling + if isinstance(child, torch.nn.Linear): + has_bias = child.bias is not None + new_layer = bnb.nn.Linear4bit( + child.in_features, + child.out_features, + bias=has_bias, + compute_dtype=compute_dtype, + compress_statistics=True, + quant_type="nf4", + ) + new_layer.weight = bnb.nn.Params4bit( + child.weight.data, + requires_grad=False, + quant_type="nf4", + ) + if has_bias: + new_layer.bias = torch.nn.Parameter(child.bias.data) + setattr(module, name, new_layer) + else: + _replace_linear_with_4bit(child, compute_dtype) + +def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True): + # Loading Flux + print("Init model") + + ckpt_path = configs[name].ckpt_path + if ( + ckpt_path is None + and configs[name].repo_id is not None + and configs[name].repo_flow is not None + and hf_download + ): + ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow) + + # Initialize model directly on CPU or target device (avoids meta-tensor shape replacement) + target_device = torch.device(device) + model = Flux(configs[name].params).to(dtype=torch.bfloat16) + + if ckpt_path is not None: + print("Loading checkpoint") + # load_sft doesn't support torch.device + sd = load_sft(ckpt_path, device="cpu") + + # --- ADD THIS LINE TO STRIP FP8 / COMFYUI KEY PREFIXES --- + sd = {k.replace("model.diffusion_model.", ""): v for k, v in sd.items()} + # --------------------------------------------------------- + + # --- FIX: HANDLE EXPANDED IMG_IN (384 channels vs 64 channels) --- + img_in_weight = sd.pop("img_in.weight", None) + img_in_bias = sd.pop("img_in.bias", None) + + # Load all standard layers safely + missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) + print_load_warning(missing, unexpected) + + # Copy base 64 channels into ArtiAgent's expanded 384-channel input layer + # In src/flux/util.py inside load_flow_model(): + + if img_in_weight is not None: + with torch.no_grad(): + w = img_in_weight.to(device=device, dtype=torch.bfloat16) + + # Check if model.img_in weight expects 384 channels while checkpoint has 64 + if model.img_in.weight.shape[1] != w.shape[1]: + # Slice model.img_in.weight to match the 64-channel input tensor + model.img_in.weight = torch.nn.Parameter(model.img_in.weight[:, :w.shape[1]]) + + model.img_in.weight.copy_(w) + + if img_in_bias is not None and getattr(model.img_in, "bias", None) is not None: + with torch.no_grad(): + b = img_in_bias.to(device=device, dtype=torch.bfloat16) + model.img_in.bias.copy_(b) + + # Quantize all Linear layers to NF4 on CPU before moving to GPU + print("Quantizing model to NF4 (this may take a minute)...") + _replace_linear_with_4bit(model, compute_dtype=torch.bfloat16) + print("NF4 quantization complete.") + + # Move model to target CUDA device + model = model.to(target_device) + return model + + +def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder: + # Force T5 onto CPU; sampling.py already moves the encoded txt tensor to GPU + return HFEmbedder( + "google/t5-v1_1-xxl", + max_length=max_length, + is_clip=False, + torch_dtype=torch.bfloat16, + device_map="cpu" + ) + + +def load_clip(device: str | torch.device = "cuda") -> HFEmbedder: + # Keep on CPU; sampling.py moves vec to GPU after encoding + return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, is_clip=True, torch_dtype=torch.bfloat16) + + +def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder: + ckpt_path = configs[name].ae_path + + # If ckpt_path is just a filename and doesn't exist locally, download it + if ckpt_path is not None and not os.path.exists(ckpt_path) and hf_download: + repo_id = configs[name].repo_ae or configs[name].repo_id + ckpt_path = hf_hub_download(repo_id, ckpt_path) + elif ckpt_path is None and configs[name].repo_id is not None and hf_download: + repo_id = configs[name].repo_ae or configs[name].repo_id + ckpt_path = hf_hub_download(repo_id, "ae.safetensors") + + # Loading the autoencoder + print("Init AE") + + # Initialize directly on CPU to avoid meta-tensor initialization issues + ae = AutoEncoder(configs[name].ae_params) + + if ckpt_path is not None: + sd = load_sft(ckpt_path, device=str(device)) + missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True) + print_load_warning(missing, unexpected) + + ae = ae.to(device) + return ae + + +# class WatermarkEmbedder: +# def __init__(self, watermark): +# self.watermark = watermark +# self.num_bits = len(WATERMARK_BITS) +# self.encoder = WatermarkEncoder() +# self.encoder.set_watermark("bits", self.watermark) + +# def __call__(self, image: torch.Tensor) -> torch.Tensor: +# """ +# Adds a predefined watermark to the input image + +# Args: +# image: ([N,] B, RGB, H, W) in range [-1, 1] + +# Returns: +# same as input but watermarked +# """ +# image = 0.5 * image + 0.5 +# squeeze = len(image.shape) == 4 +# if squeeze: +# image = image[None, ...] +# n = image.shape[0] +# image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1] +# # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255] +# # watermarking libary expects input as cv2 BGR format +# for k in range(image_np.shape[0]): +# image_np[k] = self.encoder.encode(image_np[k], "dwtDct") +# image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to( +# image.device +# ) +# image = torch.clamp(image / 255, min=0.0, max=1.0) +# if squeeze: +# image = image[0] +# image = 2 * image - 1 +# return image + + +# # A fixed 48-bit message that was chosen at random +# WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110 +# # bin(x)[2:] gives bits of x as str, use int to convert them to 0/1 +# WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]] +# embed_watermark = WatermarkEmbedder(WATERMARK_BITS) diff --git a/ArtiAgent - DefectFill/src/loop-file.py b/ArtiAgent - DefectFill/src/loop-file.py new file mode 100644 index 0000000000000000000000000000000000000000..31ca15d3a3cf0dd076e6aab0cda5aa404141cba0 --- /dev/null +++ b/ArtiAgent - DefectFill/src/loop-file.py @@ -0,0 +1,41 @@ +import subprocess +import sys + +# Common base arguments across all runs +BASE_CMD = [ + sys.executable, # Uses the currently active Python interpreter + "batch_agent_orchestrator.py", + "--input-dir", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\TestingImage\03_xray_solder", + "--output-dir", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\ArtiAgent - DefectFill\agent_output\03_xray_solder", + "--checkpoint-dir", r"C:\Users\admin_mtds\OneDrive\Desktop\ChinKuan\ArtiAgent - DefectFill\engine\DefectFill\checkpoints", + "--product-desc", "X-ray PCB with solder joints and die components", + "--max-defects-per-image", "3", + "--guidance-scale", "8.0", + "--device", "cuda" +] + +# Variations for each step in a single loop +# TASKS = [ +# ["--defect-type", "tiger-strip"], +# ["--defect-type", "bubble"], +# [] # Default/No defect-type specified +# ] + +TOTAL_LOOPS = 20 + +if __name__ == "__main__": + for loop_num in range(1, TOTAL_LOOPS + 1): + print(f"\n" + "=" * 50) + print(f"๐Ÿš€ STARTING LOOP {loop_num} OF {TOTAL_LOOPS}") + print("=" * 50) + print(f"\n--> Running command for Loop {loop_num}...") + + # Run the base command and wait for it to finish + result = subprocess.run(BASE_CMD) + + # Check if the command failed + if result.returncode != 0: + print(f"โŒ Error encountered on Loop {loop_num}. Execution stopped.") + sys.exit(result.returncode) + + print(f"\n๐ŸŽ‰ Success! Completed all {TOTAL_LOOPS} loops.") \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/model.py b/ArtiAgent - DefectFill/src/model.py new file mode 100644 index 0000000000000000000000000000000000000000..480fb6783ec95ffc60149c6c708e230ffa1b976d --- /dev/null +++ b/ArtiAgent - DefectFill/src/model.py @@ -0,0 +1,390 @@ +import os +# Hugging Face Mirror Configuration +# Option 1: hf-mirror.com (Available in some regions) +# Option 2: Use ModelScope as an alternative +USE_MODELSCOPE = False # Set to True for ModelScope, False for HuggingFace + +import torch +import torch.nn as nn +from diffusers import StableDiffusionInpaintPipeline, DDIMScheduler, UNet2DConditionModel +from transformers import CLIPTextModel +from peft import LoraConfig, get_peft_model +import lpips +import torch.nn.functional as F +from typing import Dict, List, Optional, Tuple +import math +from diffusers.models.attention_processor import Attention, AttnProcessor + + +class AttentionStoreProcessor(AttnProcessor): + """Attention Processor used to store cross-attention maps for steering""" + def __init__(self, model=None, layer_name=""): + super().__init__() + self.model = model # Reference to the main model instance + self.layer_name = layer_name # Store layer name directly + + def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None): + batch_size, sequence_length, _ = hidden_states.shape + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + + query = attn.to_q(hidden_states) + + is_cross_attention = encoder_hidden_states is not None + + if not is_cross_attention: + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + else: + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.head_to_batch_dim(query) + key = attn.head_to_batch_dim(key) + value = attn.head_to_batch_dim(value) + + attention_scores = torch.matmul(query, key.transpose(-1, -2)) * attn.scale + attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1) + + # Fast Direct Lookup (No recursive loop!) + if is_cross_attention and self.model is not None and "up_blocks" in self.layer_name: + try: + num_heads = attn.heads + total_elements = attention_probs.numel() + query_len = hidden_states.shape[1] + key_len = encoder_hidden_states.shape[1] if encoder_hidden_states is not None else query_len + + expected_size = batch_size * num_heads * query_len * key_len + if total_elements == expected_size: + reshaped_probs = attention_probs.reshape(batch_size, num_heads, query_len, key_len) + if not hasattr(self.model, "attention_maps"): + self.model.attention_maps = {} + self.model.attention_maps[self.layer_name] = reshaped_probs.detach().clone() + except Exception as e: + pass + + hidden_states = torch.matmul(attention_probs, value) + hidden_states = attn.batch_to_head_dim(hidden_states) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states + +class DefectFillModel(nn.Module): + def __init__(self, device="cuda", lora_rank=8, lora_alpha=16, seed=42, placeholder_token=""): + super().__init__() + torch.manual_seed(seed) + self.device = device + + # Base Model ID + hf_model_id = "sd2-community/stable-diffusion-2-inpainting" + + # Select model source based on configuration + if USE_MODELSCOPE: + try: + from modelscope import snapshot_download + print(f"[ModelScope] Downloading model: {hf_model_id}") + local_model_path = snapshot_download(hf_model_id) + print(f"[ModelScope] Model downloaded to: {local_model_path}") + + self.pipeline = StableDiffusionInpaintPipeline.from_pretrained( + local_model_path, + torch_dtype=torch.float16 + ).to(device) + + self.scheduler = DDIMScheduler.from_pretrained( + local_model_path, + subfolder="scheduler" + ) + except ImportError: + print("[Warning] modelscope not installed. Try: pip install modelscope") + print("[Info] Attempting HuggingFace fallback...") + self.pipeline = StableDiffusionInpaintPipeline.from_pretrained( + hf_model_id, torch_dtype=torch.float16 + ).to(device) + self.scheduler = DDIMScheduler.from_pretrained(hf_model_id, subfolder="scheduler") + else: + self.pipeline = StableDiffusionInpaintPipeline.from_pretrained( + hf_model_id, torch_dtype=torch.float16 + ).to(device) + self.scheduler = DDIMScheduler.from_pretrained(hf_model_id, subfolder="scheduler") + + self.pipeline.set_progress_bar_config(disable=True) + self.scheduler.set_timesteps(30) + + # ========== Textual Inversion: Add learnable defect token [V*] ========== + self.placeholder_token = placeholder_token + + # Add new token to tokenizer + num_added_tokens = self.pipeline.tokenizer.add_tokens([self.placeholder_token]) + if num_added_tokens == 0: + print(f"[Warning] Token {self.placeholder_token} already exists in tokenizer") + else: + print(f"[Textual Inversion] Added {num_added_tokens} new token: {self.placeholder_token}") + + # Resize text encoder embeddings + self.pipeline.text_encoder.resize_token_embeddings(len(self.pipeline.tokenizer)) + + # Get ID for the new token + self.placeholder_token_id = self.pipeline.tokenizer.convert_tokens_to_ids(self.placeholder_token) + print(f"[Textual Inversion] placeholder_token_id = {self.placeholder_token_id}") + + # Initialize new token with the embedding of "defect" + initializer_token = "defect" + initializer_token_ids = self.pipeline.tokenizer.encode(initializer_token, add_special_tokens=False) + if len(initializer_token_ids) > 0: + initializer_token_id = initializer_token_ids[0] + token_embeds = self.pipeline.text_encoder.get_input_embeddings().weight.data + token_embeds[self.placeholder_token_id] = token_embeds[initializer_token_id].clone() + print(f"[Textual Inversion] Initialized '{self.placeholder_token}' using '{initializer_token}' (id={initializer_token_id})") + + # LoRA Configuration + unet_lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + init_lora_weights="gaussian" + ) + + text_encoder_lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + target_modules=["q_proj", "k_proj", "v_proj", "out_proj"], + init_lora_weights="gaussian" + ) + + # Apply LoRA adapters + self.pipeline.unet = get_peft_model(self.pipeline.unet, unet_lora_config) + self.pipeline.text_encoder = get_peft_model(self.pipeline.text_encoder, text_encoder_lora_config) + + # Freeze VAE parameters + for param in self.pipeline.vae.parameters(): + param.requires_grad = False + + # VGG model for LPIPS loss + self.lpips_model = lpips.LPIPS(net='vgg', spatial=True).to(device) + + self.attention_maps = {} + self.register_attention_processor() + self.defect_token_indices = [] + + def register_attention_processor(self): + """Replace standard UNet attention processors with custom ones""" + self.attention_maps = {} + for name, module in self.pipeline.unet.named_modules(): + if isinstance(module, Attention) and "attn2" in name: # Target Cross-Attention only + # Pass 'name' directly into the processor + module.processor = AttentionStoreProcessor(model=self, layer_name=name) + + def get_attention_loss(self, mask_latents: torch.Tensor) -> torch.Tensor: + """ + Calculates Attention Loss - forces token attention maps to align with the defect mask. + """ + if not self.attention_maps: + return torch.tensor(0.0, device=mask_latents.device) + + if len(mask_latents.shape) == 3: + mask_latents = mask_latents.unsqueeze(1) + + batch_size = mask_latents.shape[0] + attention_loss = torch.tensor(0.0, device=mask_latents.device) + + # Use only decoder (up_blocks) attention maps + decoder_attention_maps = { + name: attn_map for name, attn_map in self.attention_maps.items() + if "up_blocks" in name + } + + if not decoder_attention_maps: + return torch.tensor(0.0, device=mask_latents.device) + + for b in range(batch_size): + token_idx = self.defect_token_indices[b] if b < len(self.defect_token_indices) else -1 + if token_idx < 0: + continue + + mask = mask_latents[b].squeeze(0) # (H, W) + resized_attention_maps = [] + + for name, attn_map in decoder_attention_maps.items(): + try: + if b < attn_map.shape[0]: + # Average attention across all heads for the specific token + defect_attn = attn_map[b, :, :, token_idx].mean(dim=0) + + seq_len = defect_attn.shape[0] + h = int(math.sqrt(seq_len)) + if h * h == seq_len: + defect_attn = defect_attn.reshape(h, h) + resized_attn = F.interpolate( + defect_attn.unsqueeze(0).unsqueeze(0), + size=mask.shape, + mode='bilinear', + align_corners=False + ).squeeze() + resized_attention_maps.append(resized_attn) + except Exception: + continue + + if resized_attention_maps: + avg_attn_map = torch.stack(resized_attention_maps).mean(dim=0) + # L2 Loss: ||AttentionMap - Mask||^2 + sample_loss = F.mse_loss(avg_attn_map, mask) + attention_loss += sample_loss + + return attention_loss / batch_size if batch_size > 0 else attention_loss + + def get_text_embeddings(self, prompts, enable_grad=True): + """Encodes prompts and locates the precise index of the token""" + if not hasattr(self, 'pipeline') or self.pipeline is None: + raise ValueError("Pipeline not initialized") + + if isinstance(prompts, str): + prompts = [prompts] + + text_inputs = self.pipeline.tokenizer( + prompts, + padding="max_length", + max_length=self.pipeline.tokenizer.model_max_length, + truncation=True, + return_tensors="pt" + ).to(self.pipeline.device) + + input_ids = text_inputs.input_ids + + # Locate the token position in each prompt + self.defect_token_indices = [] + for ids in input_ids: + positions = (ids == self.placeholder_token_id).nonzero(as_tuple=True)[0] + self.defect_token_indices.append(positions[0].item() if len(positions) > 0 else -1) + + if enable_grad: + text_embeddings = self.pipeline.text_encoder(input_ids)[0] + else: + with torch.no_grad(): + text_embeddings = self.pipeline.text_encoder(input_ids)[0] + + return text_embeddings + + def forward( + self, + noisy_latents: torch.Tensor, + masked_image_latents: torch.Tensor, + mask_latents: torch.Tensor, + timesteps: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """ + Training Forward Pass - Implements 9-channel input. + Input format: [noisy_latents(4), masked_background(4), mask(1)] + """ + self.attention_maps = {} + concat_latents = torch.cat([noisy_latents, masked_image_latents, mask_latents], dim=1) + + noise_pred = self.pipeline.unet( + concat_latents, + timesteps, + encoder_hidden_states=encoder_hidden_states, + ).sample + + attention_loss = self.get_attention_loss(mask_latents) + + return { + "noise_pred": noise_pred, + "attention_loss": attention_loss + } + + @staticmethod + def compute_masked_mse(noise_pred: torch.Tensor, noise: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Helper to calculate MSE loss only within the masked area""" + weighted_loss = mask * ((noise_pred - noise) ** 2) + return torch.sum(weighted_loss) / (torch.sum(mask) + 1e-8) + + def compute_defect_loss(self, noise_pred: torch.Tensor, noise: torch.Tensor, mask_latents: torch.Tensor) -> torch.Tensor: + """L_def loss: MSE restricted to the defect mask region""" + return self.compute_masked_mse(noise_pred, noise, mask_latents) + + def compute_object_loss(self, noise_pred: torch.Tensor, noise: torch.Tensor, mask_latents: torch.Tensor, alpha: float = 0.3) -> torch.Tensor: + """L_obj loss: Uses weighted mask M' = M + alpha*(1-M) to preserve object context""" + weighted_mask = mask_latents + alpha * (1 - mask_latents) + return self.compute_masked_mse(noise_pred, noise, weighted_mask) + + def generate( + self, + image: torch.Tensor, + mask: torch.Tensor, + prompt: str, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + """ + Complete Inference Pipeline: + 1. 9-channel input configuration + 2. Classifier-Free Guidance (CFG) + 3. Iterative background preservation: x_t = M * x_t_pred + (1-M) * x_t_background + """ + device = image.device + dtype = image.dtype + batch_size = image.shape[0] + + # Normalize image to [-1, 1] if needed + if image.min() >= 0 and image.max() <= 1: + image = 2 * image - 1 + + if len(mask.shape) == 3: mask = mask.unsqueeze(1) + if mask.max() > 1: mask = mask / 255.0 + + with torch.no_grad(): + # Encode clean image and create masked background latent b = E(I * (1-M)) + latents_clean = self.pipeline.vae.encode(image).latent_dist.sample() + latents_clean = latents_clean * self.pipeline.vae.config.scaling_factor + + masked_image = image * (1 - mask) + masked_image_latents = self.pipeline.vae.encode(masked_image).latent_dist.sample() + masked_image_latents = masked_image_latents * self.pipeline.vae.config.scaling_factor + + mask_latents = F.interpolate(mask, size=latents_clean.shape[-2:], mode='nearest') + + # Text embeddings for CFG + text_embeddings = self.get_text_embeddings([prompt] * batch_size, enable_grad=False) + uncond_embeddings = self.get_text_embeddings([""] * batch_size, enable_grad=False) + text_embeddings_cfg = torch.cat([uncond_embeddings, text_embeddings]) + + self.scheduler.set_timesteps(num_inference_steps) + latents = torch.randn(latents_clean.shape, generator=generator, device=device, dtype=dtype) + + # Denoising loop + for t in self.scheduler.timesteps: + # Generate background noise for current timestep (for background preservation) + noise_for_bg = torch.randn(latents_clean.shape, generator=generator, device=device, dtype=dtype) + latents_background = self.scheduler.add_noise(latents_clean, noise_for_bg, t) + + # Prepare inputs for CFG + latent_input = torch.cat([latents] * 2) + masked_input = torch.cat([masked_image_latents] * 2) + mask_input = torch.cat([mask_latents] * 2) + concat_input = torch.cat([latent_input, masked_input, mask_input], dim=1) + + timestep_tensor = torch.tensor([t] * (batch_size * 2), device=device, dtype=torch.long) + + noise_pred = self.pipeline.unet( + concat_input, + timestep_tensor, + encoder_hidden_states=text_embeddings_cfg + ).sample + + # Perform CFG + noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + + latents = self.scheduler.step(noise_pred, t, latents).prev_sample + + # ========== KEY STEP: Iterative Background Preservation ========== + latents = mask_latents * latents + (1 - mask_latents) * latents_background + + # Decode latents to pixels + latents = latents / self.pipeline.vae.config.scaling_factor + with torch.no_grad(): + images = self.pipeline.vae.decode(latents).sample + + return (images + 1) / 2 # Convert back to [0, 1] range \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/README.md b/ArtiAgent - DefectFill/src/pipeline/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ae88fed6ed971ca929c7cf2eb66d0ebf04e606b9 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/README.md @@ -0,0 +1,319 @@ +# Image Artifacts Pipeline + +A modular pipeline for generating synthetic image artifacts using GSAM (Grounded Segment Anything Model) for part detection and FLUX diffusion model for artifact generation. + +## Overview + +This pipeline provides clean, modular Python components for the two-stage artifact generation workflow: + +``` +๐Ÿ“ฆ pipeline/ +โ”œโ”€โ”€ ๐Ÿ“„ __init__.py # Package initialization +โ”œโ”€โ”€ ๐Ÿ“„ data_loader.py # Dataset handling (COCO, ImageNet, Custom) +โ”œโ”€โ”€ ๐Ÿ“„ gsam_detector.py # GSAM model integration (GroundingDINO + SAM) +โ”œโ”€โ”€ ๐Ÿ“„ instance_processor.py # Instance filtering and bbox operations +โ”œโ”€โ”€ ๐Ÿ“„ flux_generator.py # FLUX model operations +โ”œโ”€โ”€ ๐Ÿ“„ visualization.py # Image visualization utilities +โ”œโ”€โ”€ ๐Ÿ“„ prompts.py # OpenAI API utilities for vocabulary generation +โ””โ”€โ”€ ๐Ÿ“„ README.md # This file +``` + +## Architecture + +``` +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ COCO/ImageNet โ”‚ +โ”‚ Custom Dataset โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ Data Loader โ”‚ +โ”‚ (data_loader) โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ GSAM Detector โ”‚โ—„โ”€โ”€โ”€โ”€โ”ค OpenAI API โ”‚ +โ”‚ (GroundingDINO โ”‚ โ”‚ (Vocabulary)โ”‚ +โ”‚ + SAM/SAM-HQ) โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚Instance Processorโ”‚ +โ”‚ (Filter, Sample,โ”‚ +โ”‚ Create Patches)โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ FLUX Generator โ”‚ +โ”‚ (Diffusion with โ”‚ +โ”‚ Patch Guidance) โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ Visualizer โ”‚ +โ”‚ (Results/Masks) โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +``` + +## Components + +### 1. Data Loaders + +Handles dataset loading and image sampling for COCO, ImageNet, and custom directories. + +```python +from pipeline import COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader + +# COCO Dataset +coco_loader = COCODataLoader( + dataset_path="/path/to/coco/annotations", + image_path="/path/to/coco/images" +) +cat_ids = coco_loader.get_category_ids(['person']) +img_info, img_array, caption = coco_loader.sample_image_by_category(cat_ids) + +# ImageNet Dataset +imagenet_loader = ImageNetDataLoader( + dataset_path="/path/to/imagenet", + split="train" +) + +# Custom Directory +custom_loader = CustomDirectoryDataLoader( + directory_path="/path/to/images" +) +``` + +### 2. GSAMDetector + +Integrates GroundingDINO and SAM/SAM-HQ for part detection with OpenAI-generated vocabulary. + +```python +from pipeline import GSAMDetector + +detector = GSAMDetector( + grounding_config="GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py", + grounding_checkpoint="weight/groundingdino_swint_ogc.pth", + sam_checkpoint="weight/sam_vit_h_4b8939.pth", + sam_version="vit_h", + use_sam_hq=False, + box_threshold=0.3, + text_threshold=0.25, + device='cuda' +) + +# Detect parts in image +entity_predictions, subentity_predictions, vis_output = detector.detect_parts( + image=img_array, + entities=['person'], + subentities=['person head', 'person arm', 'person leg'], + entity_subentity_mapping={'person': ['person head', 'person arm', 'person leg']}, + min_area_ratio=0.005, + max_area_ratio=0.5 +) +``` + +### 3. InstanceProcessor + +Handles instance filtering, sampling, bounding box operations, and patch annotation creation. + +```python +from pipeline import InstanceProcessor + +# Sample instance by confidence score +sampled_instance = InstanceProcessor.sample_instance_by_score( + predictions, + min_area_ratio=0.01, + max_area_ratio=0.5 +) + +# Generate bbox suggestions for addition artifacts +suggested_bbox = InstanceProcessor.generate_bbox_suggestion( + predictions=predictions, + reference_bbox=reference_bbox, + class_name=class_name, + vocab=vocab, + max_ref_overlap=0.3, + min_entity_overlap=0.1 +) + +# Create annotation with patch indices +annotation_data = InstanceProcessor.create_annotation_dict( + instance=sampled_instance, + img_shape=image.shape, + artifact_type='distortion', + patch_size=16 +) +``` + +### 4. FluxGenerator + +Manages FLUX diffusion model operations for artifact generation with patch-based guidance. + +```python +from pipeline import FluxGenerator, FluxConfig + +# Configure FLUX model +config = FluxConfig( + name='flux-dev', + guidance=5.0, + num_steps=25, + pe_step=0.5, # Position encoding step + seed=42 +) + +# Artifact-type-specific PE steps +config_advanced = FluxConfig( + name='flux-dev', + guidance=5.0, + num_steps=25, + pe_step={ + 'addition': 0.3, + 'removal': 0.3, + 'distortion': 0.5 + }, + seed=42 +) + +generator = FluxGenerator(device='cuda', config=config_advanced) + +# Generate artifact image +generated_image = generator.generate_with_artifacts( + source_prompt="a photo of a person", + target_prompt="a photo of a person", + bbox=target_bbox, + bbox_ref=reference_bbox, + artifact_type='distortion', + source_img=image +) +``` + +### 5. ImageVisualizer + +Provides visualization utilities for debugging and quality assurance. + +```python +from pipeline import ImageVisualizer + +visualizer = ImageVisualizer() + +# Show single image with caption +visualizer.show_image(image, caption, title="Original", base_dir="output/") + +# Show comparison +visualizer.show_comparison( + original_image, + generated_image, + artifact_data, + caption="Distortion Artifact", + base_dir="output/" +) + +# Show bounding box overlay +visualizer.show_bbox_overlay( + image, + target_bbox, + base_dir="output/", + filename="bbox_overlay.png" +) + +# Show patch masks +visualizer.show_patch_masks( + image, + reference_patches, + target_patches, + base_dir="output/" +) +``` + +## Artifact Types + +The pipeline supports three types of image artifacts: + +### 1. Distortion +Modifies the appearance of existing parts while keeping them in place. +- Uses reference patches to guide where distortion is applied +- Configurable distortion kernels (jitter, swirl, voronoi) + +### 2. Removal +Removes detected parts from images naturally. +- Uses reference patches to identify removal areas +- FLUX inpainting fills removed areas contextually + +### 3. Addition +Adds new instances of detected parts in suitable locations. +- Uses reference patches as source templates +- Generates target patches using IoU-based intelligent placement +- Maintains visual consistency with surrounding context + +## Configuration + +### Detection Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `box_threshold` | `0.3` | Detection confidence threshold | +| `text_threshold` | `0.25` | Text-image matching threshold | +| `min_area_ratio` | `0.005` | Minimum part size (0.5% of image) | +| `max_area_ratio` | `0.5` | Maximum part size (50% of image) | +| `nms_threshold` | `0.5` | Non-maximum suppression threshold | + +### Generation Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `guidance` | `5.0` | Guidance scale for FLUX | +| `num_steps` | `25` | Number of diffusion steps | +| `pe_step` | `0.5` | Position encoding step size | +| `inject` | `25` | Injection step in diffusion | +| `seed` | `42` | Random seed for reproducibility | + +## Dependencies + +- `torch` - PyTorch for deep learning +- `openai` - For vocabulary generation +- `pycocotools` - For COCO dataset handling +- `supervision` - For detection utilities +- `groundingdino` - For grounded detection +- `segment_anything` - For segmentation +- `matplotlib` - For visualization +- `PIL` - For image processing +- `numpy` - For numerical operations + +## Installation + +1. Install GroundingDINO and SAM following their respective installation guides +2. Download model weights and place them in `src/weight/` directory +3. Set up OpenAI API key: `export OPENAI_API_KEY='your-key'` +4. Install required dependencies + +See the main README for detailed installation instructions. + +## Usage + +See `batch_gsam_segmentation.py` and `batch_flux_generation.py` for complete batch processing examples. + +## Performance Tips + +1. **GPU Memory**: Use SAM `vit_b` for limited GPU memory, `vit_h` for best quality +2. **Filtering**: Adjust area ratios to balance quality vs. quantity +3. **Speed**: Lower `num_steps` (15-20) for faster generation +4. **Quality**: Higher `num_steps` (25-35) for better results + +## Troubleshooting + +### Common Issues + +1. **GSAM setup issues**: Ensure GroundingDINO and SAM weights are downloaded +2. **OpenAI API errors**: Check API key and rate limits +3. **COCO dataset errors**: Verify dataset paths and structure +4. **GPU memory issues**: Use smaller SAM model or reduce batch size + +## License + +This pipeline integrates multiple open-source components, each with their own licenses. See the main repository LICENSE and model_licenses/ directory for details. diff --git a/ArtiAgent - DefectFill/src/pipeline/__init__.py b/ArtiAgent - DefectFill/src/pipeline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5fe885cb6ab84b2cbc29196b9ff9137fe157ff35 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/__init__.py @@ -0,0 +1,38 @@ +""" +Image Artifacts Pipeline + +This module provides detection and processing pipeline components for generating +image artifacts using various segmentation models. +""" + +# Import GSAMDetector if gsam dependencies exist +try: + from .gsam_detector import GSAMDetector +except ImportError: + GSAMDetector = None + +# Import Flux components if flux dependencies exist +try: + from .flux_generator import FluxGenerator, FluxConfig +except ImportError: + FluxGenerator = None + FluxConfig = None + + +try: + from .data_loader import COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader +except ImportError: + COCODataLoader, ImageNetDataLoader, CustomDirectoryDataLoader = None, None, None + +from .instance_processor import InstanceProcessor + + +__all__ = [ + 'GSAMDetector', + 'FluxGenerator', + 'FluxConfig', + 'COCODataLoader', + 'ImageNetDataLoader', + 'CustomDirectoryDataLoader', + 'InstanceProcessor', +] \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d78504750b49843a5f2db8319acdbcf7ad3cae23 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/__init__.cpython-311.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4d9bf0093a7c1cc2d9b9876e45aa4ed59a4b452 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/__init__.cpython-311.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/data_loader.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/data_loader.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a5ac44e13dfd72bfa22f1e97874e64304100073e Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/data_loader.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/data_loader.cpython-311.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/data_loader.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..473fe9fd556d225594c2b8a68d5aedc5dc10b6c7 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/data_loader.cpython-311.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/defect_rag.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/defect_rag.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80f51cd7cdc94258049b55efd3857eea9b917825 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/defect_rag.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/defectdiffu_generator.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/defectdiffu_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9ab2fd99095fa93719a614d764ac0152cada728 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/defectdiffu_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/defectfill_generator.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/defectfill_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af0dc3a01a35c63934cd4eeecb9ba9e0dcaa0683 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/defectfill_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/domain_router.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/domain_router.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7bc285e60d99a509d2e9eee3fa41da4ee63f7a5e Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/domain_router.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/flux_generator.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/flux_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a73586f58dace3605d0fda0e13ee2730641ae625 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/flux_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/flux_generator.cpython-311.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/flux_generator.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8624a69492ad38e34059020ab04204f90011c20 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/flux_generator.cpython-311.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/gsam_detector.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/gsam_detector.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..950d58e33947887a80e95fa979b850816d132097 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/gsam_detector.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/gsam_detector.cpython-311.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/gsam_detector.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad080bb5978f041e8611c69460caa534674ce6bf Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/gsam_detector.cpython-311.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/instance_processor.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/instance_processor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e1e3cc6986fd5de2a1b77ff3dc672c3d40485cf Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/instance_processor.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/instance_processor.cpython-311.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/instance_processor.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9486a6a6b5c41bffe9d70fe5389a2134c35e88b Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/instance_processor.cpython-311.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/local_vlm_client.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/local_vlm_client.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c207a38c827deb328374f2e63c14b40e3fe9c1a Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/local_vlm_client.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/__pycache__/prompts.cpython-310.pyc b/ArtiAgent - DefectFill/src/pipeline/__pycache__/prompts.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e650c33e4e31da4fce6d67b9d457180c7bf9027 Binary files /dev/null and b/ArtiAgent - DefectFill/src/pipeline/__pycache__/prompts.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/pipeline/data_loader.py b/ArtiAgent - DefectFill/src/pipeline/data_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..05ad8e165ecb1c753161f67452cf91bb56c43402 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/data_loader.py @@ -0,0 +1,732 @@ +import os +import numpy as np +from pycocotools.coco import COCO +from typing import List, Dict, Tuple, Optional, Union +import pathlib +import json +import glob +from PIL import Image +import logging +from typing import Any + + +def preprocess_image_for_flux(image_path_or_pil: Union[str, Image.Image]) -> np.ndarray: + """ + Shared image preprocessing function for flux model compatibility + + Args: + image_path_or_pil: Either a file path to image or PIL Image object + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + # Load image with PIL if path provided + if isinstance(image_path_or_pil, str): + img = Image.open(image_path_or_pil) + else: + img = image_path_or_pil + + if img.mode != 'RGB': + img = img.convert('RGB') + + # Rescale if shortest side is less than 480 + width, height = img.size + if min(width, height) < 480: + scale_factor = 480 / min(width, height) + new_width = int(width * scale_factor) + new_height = int(height * scale_factor) + img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + img_array = np.array(img) + + # Ensure dimensions are divisible by 16 for flux model compatibility + shape = img_array.shape + new_h = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16 + new_w = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16 + + # Crop image to new dimensions + img_array = img_array[:new_h, :new_w, :] + + return img_array + + +class COCODataLoader: + """Handler for COCO dataset loading and image sampling""" + + def __init__(self, dataset_path: str, image_path: str): + """ + Initialize COCO data loader + + Args: + dataset_path: Path to COCO annotations directory + image_path: Path to COCO images directory + """ + self.dataset_path = dataset_path + self.image_path = image_path + + # Load COCO annotations + self.caption_file = os.path.join(dataset_path, "captions_train2017.json") + self.class_file = os.path.join(dataset_path, "instances_train2017.json") + + self.coco_cap = COCO(self.caption_file) + self.coco_class = COCO(self.class_file) + + # Get all image IDs + self.image_ids = self.coco_cap.getImgIds() + + def get_category_ids(self, super_categories: List[str]) -> List[int]: + """ + Get category IDs for given super categories + + Args: + super_categories: List of super category names (e.g., ['person', 'animal']) + + Returns: + List of category IDs + """ + cat_ids = self.coco_class.getCatIds(supNms=super_categories) + return cat_ids + + def get_category_names(self, cat_ids: List[int]) -> List[str]: + """Get category names from category IDs""" + cats = self.coco_class.loadCats(cat_ids) + return [cat['name'] for cat in cats] + + def sample_image_by_category(self, cat_ids: List[int]) -> Tuple[Dict, np.ndarray, str]: + """ + Sample a random image containing objects from specified categories + + Args: + cat_ids: List of category IDs to sample from + + Returns: + Tuple of (image_info, image_array, caption) with image dimensions adjusted to be divisible by 16 + """ + # Get images containing specified categories + img_ids = self.coco_class.getImgIds(catIds=cat_ids[0]) # Use first category for sampling + + # Sample random image + sampled_id = img_ids[np.random.randint(0, len(img_ids))] + + # Load image info and array + img_info = self.coco_class.loadImgs(sampled_id)[0] + + # Load and preprocess image + img_path = os.path.join(self.image_path, img_info['file_name']) + img_array = preprocess_image_for_flux(img_path) + + # Get caption + ann_ids = self.coco_cap.getAnnIds(imgIds=img_info['id']) + anns = self.coco_cap.loadAnns(ann_ids) + caption = anns[0]['caption'] if anns else "" + + return img_info, img_array, caption + + def get_image_categories(self, img_info: Dict) -> List[str]: + """ + Get all category names present in an image + + Args: + img_info: Image information dictionary + + Returns: + List of unique category names in the image + """ + # Get category information for the image + ann_ids_class = self.coco_class.getAnnIds(imgIds=img_info['id']) + anns_class = self.coco_class.loadAnns(ann_ids_class) + + # Extract category IDs from annotations + cat_ids_in_image = [ann['category_id'] for ann in anns_class] + + # Get category names + categories_in_image = [] + for cat_id in cat_ids_in_image: + cat_info = self.coco_class.loadCats([cat_id])[0] + categories_in_image.append(cat_info['name']) + + # Remove duplicates and return + return list(set(categories_in_image)) + + def load_image_by_info(self, img_info: Dict) -> np.ndarray: + """ + Load image array from image info dictionary + + Args: + img_info: COCO image info dictionary + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + img_path = os.path.join(self.image_path, img_info['file_name']) + return preprocess_image_for_flux(img_path) + + def get_image_caption(self, img_info: Dict) -> str: + """ + Get caption for a specific image + + Args: + img_info: COCO image info dictionary + + Returns: + Image caption string + """ + ann_ids = self.coco_cap.getAnnIds(imgIds=img_info['id']) + anns = self.coco_cap.loadAnns(ann_ids) + caption = anns[0]['caption'] if anns else "" + return caption + + def create_category_directories(self, category_names: List[str], base_path: str = 'data/coco_2017_extracted'): + """Create directories for each category""" + for category in category_names: + pathlib.Path(f'{base_path}/{category}').mkdir(parents=True, exist_ok=True) + + +class ImageNetDataLoader: + """Handler for ImageNet dataset loading and image sampling""" + + def __init__(self, dataset_path: str, split: str = 'train'): + """ + Initialize ImageNet data loader + + Args: + dataset_path: Path to ImageNet dataset directory + split: Dataset split ('train' or 'val') + """ + self.dataset_path = dataset_path + self.split = split + self.split_path = os.path.join(dataset_path, split) + + # Load class mapping if available + self.class_mapping = self._load_class_mapping() + + # Get all synset directories + self.synsets = [d for d in os.listdir(self.split_path) + if os.path.isdir(os.path.join(self.split_path, d))] + + # Build image index + self._build_image_index() + + def _load_class_mapping(self) -> Dict[str, str]: + """ + Load class mapping from synset IDs to human-readable names + + Returns: + Dictionary mapping synset IDs to class names + """ + mapping_files = [ + os.path.join(self.dataset_path, 'imagenet_class_index.json'), + os.path.join(self.dataset_path, 'synset_words.txt'), + os.path.join(self.dataset_path, 'LOC_synset_mapping.txt') + ] + + class_mapping = {} + + # Try loading from JSON format first + for mapping_file in mapping_files: + if os.path.exists(mapping_file): + if mapping_file.endswith('.json'): + with open(mapping_file, 'r') as f: + data = json.load(f) + for idx, (synset, name) in data.items(): + class_mapping[synset] = name + break + elif mapping_file.endswith('.txt'): + with open(mapping_file, 'r') as f: + for line in f: + parts = line.strip().split('\t') + if len(parts) >= 2: + synset = parts[0] + name = parts[1] + class_mapping[synset] = name + break + + return class_mapping + + def _build_image_index(self): + """Build index of all images in the dataset""" + self.image_index = {} + + for synset in self.synsets: + synset_path = os.path.join(self.split_path, synset) + image_files = [] + + # Support common image formats + for ext in ['*.JPEG', '*.jpg', '*.jpeg', '*.png', '*.bmp']: + image_files.extend(glob.glob(os.path.join(synset_path, ext))) + + self.image_index[synset] = image_files + + def get_class_names(self) -> List[str]: + """ + Get all available class names + + Returns: + List of class names (human-readable if mapping available, else synset IDs) + """ + if self.class_mapping: + return [self.class_mapping.get(synset, synset) for synset in self.synsets] + else: + return self.synsets + + def get_synsets(self) -> List[str]: + """Get all available synset IDs""" + return self.synsets + + def sample_image_by_class(self, class_names: List[str] = None, synsets: List[str] = None) -> Tuple[Dict, np.ndarray, str]: + """ + Sample a random image from specified classes or synsets + + Args: + class_names: List of human-readable class names to sample from + synsets: List of synset IDs to sample from (takes precedence over class_names) + + Returns: + Tuple of (image_info, image_array, class_name) + """ + # Determine synsets to sample from + if synsets: + target_synsets = [s for s in synsets if s in self.synsets] + elif class_names: + # Convert class names to synsets + target_synsets = [] + for class_name in class_names: + for synset, mapped_name in self.class_mapping.items(): + if mapped_name.lower() == class_name.lower() and synset in self.synsets: + target_synsets.append(synset) + else: + # Sample from all available synsets + target_synsets = self.synsets + + if not target_synsets: + raise ValueError("No matching synsets found for the specified classes") + + # Sample random synset + sampled_synset = np.random.choice(target_synsets) + + # Sample random image from the synset + if not self.image_index[sampled_synset]: + raise ValueError(f"No images found for synset {sampled_synset}") + + sampled_image_path = np.random.choice(self.image_index[sampled_synset]) + + # Load image + img_array = self._load_and_preprocess_image(sampled_image_path) + + # Create image info + img_info = { + 'file_name': os.path.basename(sampled_image_path), + 'file_path': sampled_image_path, + 'synset': sampled_synset, + 'class_name': self.class_mapping.get(sampled_synset, sampled_synset), + 'height': img_array.shape[0], + 'width': img_array.shape[1] + } + + class_name = self.class_mapping.get(sampled_synset, sampled_synset) + + return img_info, img_array, class_name + + def load_image_by_path(self, image_path: str) -> np.ndarray: + """ + Load image from file path with preprocessing + + Args: + image_path: Path to image file + + Returns: + Preprocessed image array + """ + return self._load_and_preprocess_image(image_path) + + def _load_and_preprocess_image(self, image_path: str) -> np.ndarray: + """ + Load and preprocess image for flux model compatibility + + Args: + image_path: Path to image file + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + return preprocess_image_for_flux(image_path) + + def get_images_by_synset(self, synset: str) -> List[str]: + """ + Get all image paths for a specific synset + + Args: + synset: Synset ID + + Returns: + List of image paths + """ + return self.image_index.get(synset, []) + + def get_synset_stats(self) -> Dict[str, int]: + """ + Get statistics about number of images per synset + + Returns: + Dictionary mapping synset IDs to image counts + """ + return {synset: len(images) for synset, images in self.image_index.items()} + + def create_class_directories(self, class_names: List[str], base_path: str = 'data/imagenet_extracted'): + """ + Create directories for each class + + Args: + class_names: List of class names or synsets + base_path: Base directory to create class folders in + """ + for class_name in class_names: + # Use synset as folder name if it exists, otherwise use class name + if class_name in self.synsets: + folder_name = class_name + else: + # Find synset for class name + folder_name = class_name + for synset, mapped_name in self.class_mapping.items(): + if mapped_name.lower() == class_name.lower(): + folder_name = synset + break + + pathlib.Path(f'{base_path}/{folder_name}').mkdir(parents=True, exist_ok=True) + + +class CustomDirectoryDataLoader: + """Handler for custom directory structure with images directly in a single directory""" + + def __init__(self, dataset_path: str): + """ + Initialize custom directory data loader + + Args: + dataset_path: Path to directory containing images directly + Expected structure: dataset_path/*.jpg, dataset_path/*.png, etc. + """ + self.dataset_path = dataset_path + + if not os.path.exists(dataset_path): + raise ValueError(f"Dataset path does not exist: {dataset_path}") + + # Build image index from directory + self._build_image_index() + + if not self.image_paths: + raise ValueError(f"No images found in {dataset_path}") + + def _build_image_index(self): + """Build index of all images in the directory""" + self.image_paths = [] + + # Support common image formats + for ext in ['*.jpg', '*.jpeg', '*.JPG', '*.JPEG', '*.png', '*.PNG', + '*.bmp', '*.BMP', '*.tiff', '*.TIFF', '*.tif', '*.TIF']: + self.image_paths.extend(glob.glob(os.path.join(self.dataset_path, ext))) + + self.image_paths.sort() # Sort for consistent ordering + + def get_image_count(self) -> int: + """ + Get total number of images in the directory + + Returns: + Number of images + """ + return len(self.image_paths) + + def get_all_image_paths(self) -> List[str]: + """ + Get all image paths in the directory + + Returns: + List of image paths + """ + return self.image_paths.copy() + + def sample_random_image(self) -> Tuple[Dict, np.ndarray]: + """ + Sample a random image from the directory + + Returns: + Tuple of (image_info, image_array) + """ + if not self.image_paths: + raise ValueError("No images available to sample") + + # Sample random image path + sampled_image_path = np.random.choice(self.image_paths) + + # Load and preprocess image + img_array = self._load_and_preprocess_image(sampled_image_path) + + # Create image info + img_info = { + 'file_name': os.path.basename(sampled_image_path), + 'file_path': sampled_image_path, + 'height': img_array.shape[0], + 'width': img_array.shape[1] + } + + return img_info, img_array + + def sample_multiple_images(self, num_samples: int = 1) -> List[Tuple[Dict, np.ndarray]]: + """ + Sample multiple images from the directory + + Args: + num_samples: Number of images to sample + + Returns: + List of tuples (image_info, image_array) + """ + if num_samples > len(self.image_paths): + raise ValueError(f"Requested {num_samples} samples but only {len(self.image_paths)} images available") + + # Sample without replacement + sampled_paths = np.random.choice(self.image_paths, size=num_samples, replace=False) + + results = [] + for image_path in sampled_paths: + img_array = self._load_and_preprocess_image(image_path) + img_info = { + 'file_name': os.path.basename(image_path), + 'file_path': image_path, + 'height': img_array.shape[0], + 'width': img_array.shape[1] + } + results.append((img_info, img_array)) + + return results + + def load_image_by_path(self, image_path: str) -> np.ndarray: + """ + Load image from file path with preprocessing + + Args: + image_path: Path to image file + + Returns: + Preprocessed image array + """ + return self._load_and_preprocess_image(image_path) + + def _load_and_preprocess_image(self, image_path: str) -> np.ndarray: + """ + Load and preprocess image for flux model compatibility + + Args: + image_path: Path to image file + + Returns: + Image array with dimensions adjusted to be divisible by 16 + """ + return preprocess_image_for_flux(image_path) + + def load_image_by_info(self, img_info: Dict) -> np.ndarray: + """ + Load image by image info dictionary + + Args: + img_info: Dictionary containing 'file_path' key + + Returns: + Preprocessed image array + """ + image_path = img_info.get('file_path') + if not image_path: + raise ValueError("Image info must contain 'file_path' key") + return self._load_and_preprocess_image(image_path) + + +def _get_coco_image_list( + data_loader: COCODataLoader, + categories: List[str], + max_images: Optional[int] = None, + max_instances_per_image: Optional[int] = None +) -> List[Dict[str, Any]]: + """ + Get image list for COCO dataset with optional filtering. + + Args: + data_loader: COCO data loader instance + categories: List of categories to process + max_images: Maximum number of images to process + max_instances_per_image: Maximum instances per image for filtering + + Returns: + List of image information dictionaries + """ + cat_ids = data_loader.get_category_ids(categories) + image_list = [] + image_ids_seen = set() + + # Count instances per image if filtering is requested + instance_counts = {} + if max_instances_per_image is not None: + print("Counting instances per image...") + from collections import defaultdict + instance_counts = defaultdict(int) + for ann in data_loader.coco_class.dataset['annotations']: + image_id = ann['image_id'] + instance_counts[image_id] += 1 + + for cat_id in cat_ids: + img_ids = data_loader.coco_class.getImgIds(catIds=[cat_id]) + for img_id in img_ids: + if img_id not in image_ids_seen: + # Filter by instance count if specified + if max_instances_per_image is not None: + if instance_counts[img_id] >= max_instances_per_image: + continue + + img_info = data_loader.coco_class.loadImgs([img_id])[0] + image_list.append(img_info) + image_ids_seen.add(img_id) + + print("number of images", len(image_list)) + return image_list + + +def _get_imagenet_image_list( + data_loader: ImageNetDataLoader, + categories: List[str], + max_images: Optional[int] = None +) -> List[Dict[str, Any]]: + """ + Get image list for ImageNet dataset. + + Args: + data_loader: ImageNet data loader instance + categories: List of categories to process + max_images: Maximum number of images to process + + Returns: + List of image information dictionaries + """ + # Determine target synsets + target_synsets = [] + for class_name in categories: + for synset, mapped_name in data_loader.class_mapping.items(): + if mapped_name.lower() == class_name.lower() and synset in data_loader.synsets: + target_synsets.append(synset) + + if not target_synsets: + target_synsets = data_loader.synsets + + image_list = [] + for synset in target_synsets: + image_paths = data_loader.get_images_by_synset(synset) + for img_path in image_paths: + img_info = { + 'id': hash(img_path) % 1000000, # Generate unique ID + 'file_name': os.path.basename(img_path), + 'file_path': img_path, + 'synset': synset, + 'class_name': data_loader.class_mapping.get(synset, synset) + } + image_list.append(img_info) + + if max_images and len(image_list) >= max_images: + break + if max_images and len(image_list) >= max_images: + break + + return image_list + + +def _get_custom_image_list( + data_loader: CustomDirectoryDataLoader, + categories: List[str], + max_images: Optional[int] = None, + logger: logging.Logger = None +) -> List[Dict[str, Any]]: + """ + Get image list for custom dataset. + + Args: + data_loader: Custom directory data loader instance + categories: List of categories (ignored for flat directory structure) + max_images: Maximum number of images to process + logger: Logger instance + + Returns: + List of image information dictionaries + """ + # Get all available image paths from the directory + all_image_paths = data_loader.get_all_image_paths() + if logger: + logger.info(f"Found {len(all_image_paths)} images in custom dataset directory") + + # Limit images if max_images is specified + if max_images and max_images < len(all_image_paths): + all_image_paths = all_image_paths[:max_images] + if logger: + logger.info(f"Limited to first {max_images} images") + + # Create image info list + image_list = [] + for img_path in all_image_paths: + img_info = { + 'id': hash(img_path) % 1000000, # Generate unique ID + 'file_name': os.path.basename(img_path), + 'file_path': img_path + } + image_list.append(img_info) + + return image_list + + +def _get_image_list( + dataset_type: str, + data_loader: Any, + categories: List[str], + max_images: Optional[int] = None, + max_instances_per_image: Optional[int] = None, + logger: logging.Logger = None +) -> List[Dict[str, Any]]: + """ + Get image list based on dataset type. + + Args: + dataset_type: Type of dataset + data_loader: Data loader instance + categories: List of categories to process + max_images: Maximum number of images to process + max_instances_per_image: Maximum number of instances per image + logger: Logger instance + + Returns: + List of image information dictionaries + """ + if dataset_type == "coco": + return _get_coco_image_list(data_loader, categories, max_images, max_instances_per_image) + elif dataset_type == "imagenet": + return _get_imagenet_image_list(data_loader, categories, max_images) + elif dataset_type == "custom": + return _get_custom_image_list(data_loader, categories, max_images, logger) + else: + raise ValueError(f"Unsupported dataset type: {dataset_type}") + + +def _initialize_data_loader(dataset_type: str, config: Dict[str, Any]) -> Any: + """ + Initialize the appropriate data loader based on dataset type. + + Args: + dataset_type: Type of dataset ('coco', 'imagenet', 'custom') + config: Configuration dictionary + + Returns: + Initialized data loader instance + """ + if dataset_type == "coco": + return COCODataLoader(config['dataset_path'], config['image_path']) + elif dataset_type == "imagenet": + return ImageNetDataLoader(config['dataset_path'], config['imagenet_split']) + elif dataset_type == "custom": + return CustomDirectoryDataLoader(config['dataset_path']) + else: + raise ValueError(f"Unsupported dataset type: {dataset_type}") \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/defect_rag.py b/ArtiAgent - DefectFill/src/pipeline/defect_rag.py new file mode 100644 index 0000000000000000000000000000000000000000..1d080b7445ffa37c5cbcf2d04d5d5c00bfe6aecc --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/defect_rag.py @@ -0,0 +1,181 @@ +# ============================================ +# DefectRAG โ€” DefectDiffu Edition +# ============================================ +# Retrieves in-context defect examples from ChromaDB and composes +# DefectDiffu text prompts (c_d) from retrieved knowledge. + +import chromadb +from sentence_transformers import SentenceTransformer +import json +from pathlib import Path +from typing import List, Dict, Optional + + +class DefectRAG: + """ + Retrieval-Augmented Generation for defect examples. + Now includes helper to build DefectDiffu defect prompts (c_d) from RAG results. + """ + + def __init__( + self, + db_path: str = "data/defect_db", + collection: str = "defect_patches", + model: str = "all-MiniLM-L6-v2" + ): + self.db_path = db_path + self.collection_name = collection + self._client = None + self._collection = None + self._encoder = None + + @property + def client(self): + if self._client is None: + self._client = chromadb.PersistentClient(path=self.db_path) + return self._client + + @property + def collection(self): + if self._collection is None: + self._collection = self.client.get_collection(self.collection_name) + return self._collection + + @property + def encoder(self): + if self._encoder is None: + self._encoder = SentenceTransformer('all-MiniLM-L6-v2') + return self._encoder + + def _build_where_filter( + self, + commercial_only: bool = True, + domain_filter: Optional[str] = None + ) -> Optional[Dict]: + conditions = [] + if commercial_only: + conditions.append({"commercial_ok": True}) + if domain_filter and domain_filter != "general": + conditions.append({"domain": domain_filter}) + if len(conditions) == 0: + return None + elif len(conditions) == 1: + return conditions[0] + else: + return {"$and": conditions} + + def retrieve( + self, + defect_plan, + k: int = 3, + domain_filter: Optional[str] = None, + commercial_only: bool = True + ) -> List[Dict]: + """Retrieve top-k matching defect examples.""" + query_parts = [ + getattr(defect_plan, 'artifact_type', defect_plan.defect_type), + defect_plan.description, + "on", + defect_plan.target_entity + ] + query = " ".join(query_parts) + query_emb = self.encoder.encode(query) + where_filter = self._build_where_filter(commercial_only, domain_filter) + + results = self.collection.query( + query_embeddings=[query_emb.tolist()], + n_results=k, + where=where_filter + ) + + examples = [] + for i, meta in enumerate(results['metadatas'][0]): + paths = json.loads(meta['paths']) if isinstance(meta.get('paths'), str) else meta.get('paths', {}) + examples.append({ + 'paths': paths, + 'caption': meta.get('caption', ''), + 'domain': meta.get('domain', 'unknown'), + 'license': meta.get('license', 'unknown'), + 'source': meta.get('source', 'unknown'), + 'defect_name': meta.get('defect_name', 'unknown'), + 'score': results['distances'][0][i] if results.get('distances') else None, + 'metadata': {k: v for k, v in meta.items() if k not in {'paths', 'caption', 'domain', 'license', 'source', 'defect_name'}} + }) + return examples + + def retrieve_by_text( + self, + text: str, + k: int = 3, + domain_filter: Optional[str] = None + ) -> List[Dict]: + """Direct text search (for debugging/testing).""" + query_emb = self.encoder.encode(text) + where_filter = self._build_where_filter(True, domain_filter) + results = self.collection.query( + query_embeddings=[query_emb.tolist()], + n_results=k, + where=where_filter + ) + examples = [] + for i, meta in enumerate(results['metadatas'][0]): + paths = json.loads(meta['paths']) if isinstance(meta.get('paths'), str) else meta.get('paths', {}) + examples.append({ + 'paths': paths, + 'caption': meta.get('caption', ''), + 'domain': meta.get('domain', 'unknown'), + 'score': results['distances'][0][i] if results.get('distances') else None + }) + return examples + + def compose_defect_prompt( + self, + base_description: str, + examples: List[Dict], + max_examples: int = 1 + ) -> str: + """ + Build a DefectDiffu defect prompt (c_d) by enriching the base description + with captions from retrieved RAG examples. + + Example output: + "A photo of a small transparent bubble trapped under glass, similar to + a spherical air pocket with dark meniscus ring" + """ + if not examples: + return f"A photo of {base_description}" + + captions = [ex.get('caption', '') for ex in examples[:max_examples] if ex.get('caption')] + if captions: + enriched = f"{base_description}, similar to {captions[0]}" + return f"A photo of {enriched}" + return f"A photo of {base_description}" + + def get_stats(self) -> Dict: + """Get DB statistics.""" + count = self.collection.count() + results = self.collection.get() + domains = {} + licenses = {} + sources = {} + for meta in results["metadatas"]: + domains[meta.get("domain", "unknown")] = domains.get(meta.get("domain"), 0) + 1 + licenses[meta.get("license", "unknown")] = licenses.get(meta.get("license"), 0) + 1 + sources[meta.get("source", "unknown")] = sources.get(meta.get("source"), 0) + 1 + return { + "total_entries": count, + "domains": domains, + "licenses": licenses, + "sources": sources + } + + +# Singleton instance +_rag_instance = None + +def get_rag(db_path="data/defect_db") -> DefectRAG: + """Get or create singleton RAG instance.""" + global _rag_instance + if _rag_instance is None: + _rag_instance = DefectRAG(db_path=db_path) + return _rag_instance diff --git a/ArtiAgent - DefectFill/src/pipeline/defectdiffu_generator.py b/ArtiAgent - DefectFill/src/pipeline/defectdiffu_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..5b034e7d8934fbdbeb53bf4d1e78ca4973530710 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/defectdiffu_generator.py @@ -0,0 +1,361 @@ +""" +DefectDiffu Generator Wrapper + +Replaces FLUX/RF-Solver-Edit with DefectDiffu's text-guided disentangled +architecture for manufacturing defect generation. + +DefectDiffu (ECCV 2024) uses: + - Three disentangled text prompts: c_p (product/bg), c_d (defect), c_f (fusion) + - Double-free strategy with perturbation scales w_p and w_d + - Automatic mask extraction from defect-block cross-attention maps + - DiT backbone + Stable Diffusion VAE + +INTEGRATION NOTE: + This file contains the INTERFACE. You must plug in the actual DefectDiffu + model forward pass from the official repo at the marked TODO sections. + Repo: https://github.com/FFDD-diffusion/DefectDiffu +""" + +import os +import torch +import torch.nn.functional as F +import numpy as np +from typing import Dict, Tuple, Optional, List +from dataclasses import dataclass +from PIL import Image +import warnings + +# === DefectDiffu actual imports (must be in PYTHONPATH) === +from diffusers.models import AutoencoderKL + +import sys +from pathlib import Path + +# Automatically locate and add the DefectDiffu engine directory to sys.path +CURRENT_DIR = Path(__file__).resolve().parent +DEFECTDIFFU_DIR = CURRENT_DIR.parent.parent / "engine" / "DefectDiffu" + +if DEFECTDIFFU_DIR.exists() and str(DEFECTDIFFU_DIR) not in sys.path: + sys.path.insert(0, str(DEFECTDIFFU_DIR)) +import clip.clip as clip + +from models_add_cross_concate import DiT +from diffusion import create_diffusion + + +# ========================================================================= +# Mask binarization helpers (copied from test.py) +# ========================================================================= + +def rgb_to_gray(tensor): + r, g, b = tensor[:, 0], tensor[:, 1], tensor[:, 2] + gray = 0.299 * r + 0.587 * g + 0.114 * b + return gray + + +def iterative_thresholding_batch(gray_tensor): + gray_np = gray_tensor.detach().cpu().numpy() + binarized = np.zeros_like(gray_np, dtype=np.uint8) + + for i in range(gray_np.shape[0]): + img = gray_np[i] + T = img.mean() + prev_T = -1 + + while abs(T - prev_T) > 1e-4: + prev_T = T + G1 = img[img >= T] + G2 = img[img < T] + m1 = G1.mean() if G1.size > 0 else 0 + m2 = G2.mean() if G2.size > 0 else 0 + T = (m1 + m2) / 2 + + binarized[i] = (img >= T).astype(np.uint8) + + return torch.from_numpy(binarized).to(gray_tensor.device) + + +def binarize_tensor_iterative(x): + gray = rgb_to_gray(x) + binary = iterative_thresholding_batch(gray) + return binary.unsqueeze(1) + +@dataclass +class DefectDiffuConfig: + """Configuration for DefectDiffu inference.""" + ckpt_path: str # Path to trained DefectDiffu checkpoint + vae_path: str # Path to SD VAE (stabilityai/sd-vae-ft-mse) + dit_model: str = "DiT-XL/2" # DiT variant (DiT-XL/2, DiT-L/2, etc.) + image_size: int = 512 # Must match training resolution + num_steps: int = 50 # DDPM/DDIM inference steps + cfg_scale: float = 1.0 # Classifier-free guidance (if used) + device: str = "cuda" + offload: bool = False # CPU offload for low-VRAM GPUs + seed: int = 42 + + +class DefectDiffuGenerator: + """ + Wrapper around DefectDiffu for the agentic pipeline. + + Unlike FLUX (which edits an existing image via inversion-injection), + DefectDiffu generates a NEW image from noise conditioned on three text + prompts. The input "clean image" is used only for planning/verification, + not as a pixel-level source for editing. + """ + + def __init__(self, config: DefectDiffuConfig): + self.config = config + self.device = torch.device(config.device) + self._models_loaded = False + + # Placeholders โ€” populated in _load_models() + self.dit = None + self.vae = None + self.text_encoder = None + self.tokenizer = None + self.scheduler = None + + self._load_models() + + # ------------------------------------------------------------------ + # TODO: Replace the methods below with actual DefectDiffu code + # ------------------------------------------------------------------ + + def _load_models(self): + """Load DiT, VAE, text encoder, and scheduler.""" + print(f"[DefectDiffu] Loading checkpoint: {self.config.ckpt_path}") + print(f"[DefectDiffu] VAE: {self.config.vae_path}") + + # 1. CLIP RN50 (must match training) + self.model_clip, _ = clip.load('RN50', self.device) + self.model_clip.eval() + + # 2. DiT architecture (must match train.py exactly) + latent_size = self.config.image_size // 8 + self.dit = DiT( + depth=28, + hidden_size=1152, + patch_size=2, + num_heads=16, + input_size=latent_size, + num_classes=1000 + ).to(self.device) + + print(f"[DefectDiffu] Loading DiT weights from: {self.config.ckpt_path}") + checkpoint = torch.load(self.config.ckpt_path, map_location=self.device) + if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: + self.dit.load_state_dict(checkpoint['model_state_dict']) + else: + self.dit.load_state_dict(checkpoint) + self.dit.eval() + + # 3. Stable Diffusion VAE + self.vae = AutoencoderKL.from_pretrained(self.config.vae_path).to(self.device) + self.vae.eval() + + # 4. Diffusion sampler (respacing = num_steps) + self.diffusion = create_diffusion(timestep_respacing=str(self.config.num_steps)) + + self._models_loaded = True + print("[DefectDiffu] All models loaded successfully.") + + def _encode_text(self, prompt: str) -> torch.Tensor: + """Encode a text prompt into CLIP RN50 text embeddings.""" + with torch.no_grad(): + tokens = clip.tokenize([prompt]).to(self.device) + emb = self.model_clip.encode_text(tokens) + emb = emb / emb.norm(dim=-1, keepdim=True) + emb = emb.float() + return emb + + def _extract_mask_from_attention( + self, + mask_latent: torch.Tensor + ) -> np.ndarray: + """ + Decode mask latent through VAE and binarize using iterative thresholding. + Matches test.py post-processing. + """ + with torch.no_grad(): + mask_decoded = self.vae.decode(mask_latent / 0.18215).sample # [1, 3, H, W] + + # Binarize with iterative thresholding (Otsu-like) + mask_binary = binarize_tensor_iterative(mask_decoded) # [1, 1, H, W] + mask_bool = mask_binary[0, 0].cpu().numpy() > 0 + + return mask_bool + + def _denoise_with_double_free( + self, + z: torch.Tensor, + emb_p: torch.Tensor, + emb_d: torch.Tensor, + emb_f: torch.Tensor, + emb_good: torch.Tensor, + emb_null_good: torch.Tensor, + w_d: float, + w_p: float + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Run DefectDiffu inference via p_sample_loop with dual-branch CFG. + Matches test.py exactly. + + Returns: + (img_latent, mask_latent) in VAE latent space, shape [1, 4, H, W] + """ + # Build paired conditioning: defect_class vs good_class + y_defect_class = [emb_d, emb_p, emb_f] + y_good_class = [emb_good, emb_p, emb_null_good] + y = [y_defect_class, y_good_class] + + # Duplicate latent for CFG (concatenated batch) + z_cfg = torch.cat([z, z], dim=0) + + model_kwargs = dict(y=y, cfg_scale=float(w_d)) + + with torch.no_grad(): + samples, cross = self.diffusion.p_sample_loop( + self.dit.forward_with_cfg_2, + z_cfg.shape, + z_cfg, + clip_denoised=False, + model_kwargs=model_kwargs, + progress=False, + device=self.device + ) + + # Unchunk: first half is the defect-conditioned output + img_latent, _ = samples.chunk(2, dim=0) + mask_latent, _ = cross.chunk(2, dim=0) + + return img_latent, mask_latent + + # ------------------------------------------------------------------ + # Public API โ€” used by the orchestrator + # ------------------------------------------------------------------ + + @torch.no_grad() + def generate( + self, + c_p: str, + c_d: str, + c_f: str, + w_d: float = 1.0, + w_p: float = 1.0, + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray]: + """ + Generate a synthetic defect image and its binary mask. + + Args: + c_p: Background/product prompt (e.g. "A photo of metal nut") + c_d: Defect prompt (e.g. "A photo of scratch") + c_f: Fusion prompt (e.g. "A photo of metal nut with scratch") + w_d: Defect strength perturbation scale (0.0 = no defect, 2.0 = severe) + w_p: Product consistency scale (usually 1.0, increase for stronger product fidelity) + seed: Random seed + + Returns: + (pil_image, binary_mask) where mask is bool array [H, W] + """ + if seed is None: + seed = self.config.seed + torch.manual_seed(seed) + np.random.seed(seed) + + print(f"[DefectDiffu] Generating: w_d={w_d}, w_p={w_p}") + print(f"[DefectDiffu] c_p: {c_p}") + print(f"[DefectDiffu] c_d: {c_d}") + print(f"[DefectDiffu] c_f: {c_f}") + + # 1. Parse product name from c_p for the null-good prompt + product_name = c_p.replace("A photo of ", "").strip() + + # 2. Encode all five text conditions (must match training format) + emb_d = self._encode_text(c_d) # "a photo of scratch" + emb_p = self._encode_text(c_p) # "a photo of vcsel" + emb_f = self._encode_text(c_f) # "a photo of scratch vcsel" + emb_good = self._encode_text("a photo of good") + emb_null_good = self._encode_text(f"a photo of good {product_name}") + + # 3. Initialize latent noise + latent_h = self.config.image_size // 8 + latent_w = self.config.image_size // 8 + z = torch.randn(1, 4, latent_h, latent_w, device=self.device) + + # 4. DefectDiffu double-free denoising + img_latent, mask_latent = self._denoise_with_double_free( + z, emb_p, emb_d, emb_f, emb_good, emb_null_good, w_d, w_p + ) + + # 5. Decode image latent โ†’ RGB + with torch.no_grad(): + img_tensor = self.vae.decode(img_latent / 0.18215).sample # [1, 3, H, W] + img_tensor = (img_tensor + 1) / 2 # [-1, 1] โ†’ [0, 1] + img_tensor = img_tensor.clamp(0, 1) + + img_np = (img_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) + pil_image = Image.fromarray(img_np) + + # 6. Extract defect mask from mask latent + binary_mask = self._extract_mask_from_attention(mask_latent) + + print(f"[DefectDiffu] Generation complete. Mask coverage: {binary_mask.mean():.3f}") + return pil_image, binary_mask + + @torch.no_grad() + def generate_from_plan( + self, + product_description: str, + defect_description: str, + severity: str = "medium", + w_d: Optional[float] = None, + w_p: float = 1.0, + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray, Dict]: + """ + Convenience wrapper that builds the three DefectDiffu prompts from + product/defect descriptions and maps severity to w_d. + """ + # Map severity to defect strength + severity_to_wd = {"low": 0.6, "minor": 0.6, + "medium": 1.0, "moderate": 1.0, + "high": 1.5, "severe": 1.5} + if w_d is None: + w_d = severity_to_wd.get(severity.lower(), 1.0) + + c_p = f"A photo of {product_description}" + c_d = f"A photo of {defect_description}" + c_f = f"A photo of {product_description} with {defect_description}" + + img, mask = self.generate(c_p, c_d, c_f, w_d=w_d, w_p=w_p, seed=seed) + + meta = { + "c_p": c_p, "c_d": c_d, "c_f": c_f, + "w_d": w_d, "w_p": w_p, "seed": seed or self.config.seed + } + return img, mask, meta + + def unload_models(self): + """Free GPU memory.""" + self.dit = None + self.vae = None + self.model_clip = None + self.diffusion = None + if self.device.type == "cuda": + torch.cuda.empty_cache() + print("[DefectDiffu] Models unloaded.") + + +def get_defectdiffu_generator( + ckpt_path: str, + vae_path: str, + device: str = "cuda", + **kwargs +) -> DefectDiffuGenerator: + """Factory function for easy instantiation.""" + config = DefectDiffuConfig(ckpt_path=ckpt_path, vae_path=vae_path, + device=device, **kwargs) + return DefectDiffuGenerator(config) + diff --git a/ArtiAgent - DefectFill/src/pipeline/defectfill_generator - Copy.py b/ArtiAgent - DefectFill/src/pipeline/defectfill_generator - Copy.py new file mode 100644 index 0000000000000000000000000000000000000000..4d852b1d1482241cb52db4e0046992bc8cb77d5f --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/defectfill_generator - Copy.py @@ -0,0 +1,248 @@ +""" +DefectFill Generator Wrapper โ€” ArtiAgent Industrial Inpainting Edition + +Replaces FLUX / DefectDiffu with DefectFill (Inpainting Diffusion + LoRA) +for targeted industrial defect generation (e.g., X-ray PCB, die delamination, Head-In-Pillow). + +DefectFill uses: + - Pre-trained Stable Diffusion Inpainting backbone (SD 2.0 / SD 1.5) + - Fine-tuned LoRA adapters for specific industrial defect classes + - Direct ROI patch + binary mask inpainting to generate photorealistic defect textures +""" + +import os +import torch +import numpy as np +from typing import Dict, Tuple, Optional, Union +from dataclasses import dataclass +from PIL import Image +import warnings + +from diffusers import StableDiffusionInpaintPipeline, DPMSolverMultistepScheduler + + +@dataclass +class DefectFillConfig: + """Configuration for DefectFill inpainting inference.""" + ckpt_path: str = "" # Path to trained LoRA checkpoint directory or file (.safetensors / .bin) + base_model: str = "sd2-community/stable-diffusion-2-inpainting" # Base model repo or local path + image_size: int = 512 # Standard training/inference resolution (512 or 768) + num_steps: int = 50 # Denoising inference steps + guidance_scale: float = 7.5 # Classifier-Free Guidance (CFG) scale + device: str = "cuda" + torch_dtype: torch.dtype = torch.float16 + seed: int = 42 + + +class DefectFillGenerator: + """ + Wrapper around DefectFill Inpainting Diffusion for the agentic pipeline. + + Unlike text-to-image models, DefectFill takes a clean crop patch and a binary + mask, modifying ONLY the masked region to introduce realistic defect textures + while preserving non-defect surroundings. + """ + + def __init__(self, config: DefectFillConfig): + self.config = config + self.device = torch.device(config.device) + self.dtype = config.torch_dtype if config.device != "cpu" else torch.float32 + + self.pipe = None + self._models_loaded = False + + self._load_models() + + def _load_models(self): + """Load Stable Diffusion Inpainting pipeline and apply DefectFill LoRA weights.""" + print(f"[DefectFill] Loading base inpainting model: {self.config.base_model}") + + # 1. Initialize SD Inpainting Pipeline + self.pipe = StableDiffusionInpaintPipeline.from_pretrained( + self.config.base_model, + torch_dtype=self.dtype, + safety_checker=None + ) + + # 2. Set fast, robust DPMSolverMultistep scheduler + self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config) + + # 3. Load DefectFill LoRA Weights (if provided) + if self.config.ckpt_path and os.path.exists(self.config.ckpt_path): + print(f"[DefectFill] Loading LoRA weights from: {self.config.ckpt_path}") + try: + if os.path.isdir(self.config.ckpt_path): + self.pipe.load_lora_weights(self.config.ckpt_path) + else: + checkpoint_dir = os.path.dirname(self.config.ckpt_path) + weight_name = os.path.basename(self.config.ckpt_path) + self.pipe.load_lora_weights(checkpoint_dir, weight_name=weight_name) + print("[DefectFill] LoRA weights loaded successfully.") + except Exception as e: + print(f"[DefectFill] Warning: Failed to load LoRA weights directly via diffusers ({e}).") + print("[DefectFill] Running base inpainting model without adapter.") + else: + print("[DefectFill] No checkpoint path provided or file not found. Operating in base inpainting mode.") + + # 4. Move pipeline to target device + self.pipe.to(self.device) + + # 5. Enable memory optimizations if available + if self.device.type == "cuda": + try: + self.pipe.enable_attention_slicing() + except Exception: + pass + + self._models_loaded = True + print("[DefectFill] Model initialization complete.") + + def _format_image(self, input_img: Union[np.ndarray, Image.Image, str]) -> Image.Image: + """Convert input image into a 3-channel RGB PIL Image of target size.""" + if isinstance(input_img, str): + img = Image.open(input_img).convert("RGB") + elif isinstance(input_img, np.ndarray): + img = Image.fromarray(input_img).convert("RGB") + elif isinstance(input_img, Image.Image): + img = input_img.convert("RGB") + else: + raise TypeError(f"Unsupported image input type: {type(input_img)}") + + if img.size != (self.config.image_size, self.config.image_size): + img = img.resize((self.config.image_size, self.config.image_size), Image.LANCZOS) + return img + + def _format_mask(self, input_mask: Union[np.ndarray, Image.Image]) -> Tuple[Image.Image, np.ndarray]: + """Convert binary mask input into PIL 'L' mode image and boolean numpy array.""" + if isinstance(input_mask, np.ndarray): + mask_np = input_mask + elif isinstance(input_mask, Image.Image): + mask_np = np.array(input_mask) + else: + raise TypeError(f"Unsupported mask input type: {type(input_mask)}") + + # Binarize mask (0 = keep background, 255 = inpaint defect) + mask_binary = (mask_np > 0).astype(np.uint8) * 255 + mask_pil = Image.fromarray(mask_binary).convert("L") + + if mask_pil.size != (self.config.image_size, self.config.image_size): + mask_pil = mask_pil.resize((self.config.image_size, self.config.image_size), Image.NEAREST) + mask_binary = np.array(mask_pil) > 0 + + return mask_pil, (mask_binary > 0) + + @torch.no_grad() + def inpaint( + self, + image: Union[np.ndarray, Image.Image, str], + mask: Union[np.ndarray, Image.Image], + prompt: str, + defect_type: Optional[str] = None, + seed: Optional[int] = None, + guidance_scale: Optional[float] = None, + num_steps: Optional[int] = None + ) -> Image.Image: + """ + Inpaint a defect onto the provided image patch using the binary mask. + + Args: + image: Clean background image/patch. + mask: Binary mask where white (255) defines the target defect ROI. + prompt: Text prompt describing the target defect (e.g., "a photo of xray_PCB with xray_die defect"). + defect_type: Optional defect class name. + seed: Random seed for reproducible generation. + guidance_scale: CFG scale (defaults to config value if None). + num_steps: Denoising steps (defaults to config value if None). + + Returns: + PIL.Image.Image: Inpainted image containing the generated defect texture. + """ + if not self._models_loaded: + self._load_models() + + pil_image = self._format_image(image) + pil_mask, _ = self._format_mask(mask) + + num_inference_steps = num_steps if num_steps is not None else self.config.num_steps + guidance = guidance_scale if guidance_scale is not None else self.config.guidance_scale + current_seed = seed if seed is not None else self.config.seed + + generator = torch.Generator(device=self.device).manual_seed(current_seed) + + print(f"[DefectFill] Running Inpainting: Prompt='{prompt}', Steps={num_inference_steps}, Seed={current_seed}") + + # Run diffusion inpainting + output = self.pipe( + prompt=prompt, + image=pil_image, + mask_image=pil_mask, + height=self.config.image_size, + width=self.config.image_size, + num_inference_steps=num_inference_steps, + guidance_scale=guidance, + generator=generator + ).images[0] + + return output + + @torch.no_grad() + def generate( + self, + image: Union[np.ndarray, Image.Image, str], + mask: Union[np.ndarray, Image.Image], + prompt: str, + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray]: + """Convenience method returning both inpainted image and boolean mask array.""" + pil_mask, bool_mask = self._format_mask(mask) + inpainted_image = self.inpaint(image=image, mask=pil_mask, prompt=prompt, seed=seed) + return inpainted_image, bool_mask + + @torch.no_grad() + def generate_from_plan( + self, + product_description: str, + defect_description: str, + image_patch: Union[np.ndarray, Image.Image], + mask_patch: Union[np.ndarray, Image.Image], + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray, Dict]: + """Agentic pipeline helper wrapper.""" + prompt = f"a photo of {product_description} with {defect_description} defect" + inpainted_patch, bool_mask = self.generate( + image=image_patch, + mask=mask_patch, + prompt=prompt, + seed=seed + ) + meta = { + "prompt": prompt, + "product_description": product_description, + "defect_description": defect_description, + "seed": seed or self.config.seed + } + return inpainted_patch, bool_mask, meta + + def unload_models(self): + """Free GPU memory.""" + self.pipe = None + self._models_loaded = False + if self.device.type == "cuda": + torch.cuda.empty_cache() + print("[DefectFill] Models unloaded from memory.") + + +def get_defectfill_generator( + ckpt_path: str = "", + base_model: str = "sd2-community/stable-diffusion-2-inpainting", + device: str = "cuda", + **kwargs +) -> DefectFillGenerator: + """Factory function for instantiating DefectFillGenerator.""" + config = DefectFillConfig( + ckpt_path=ckpt_path, + base_model=base_model, + device=device, + **kwargs + ) + return DefectFillGenerator(config) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/defectfill_generator.py b/ArtiAgent - DefectFill/src/pipeline/defectfill_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..86d2ffc4f3a3de4432a92767ae5e6760470720ef --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/defectfill_generator.py @@ -0,0 +1,264 @@ +""" +DefectFill Generator Wrapper โ€” ArtiAgent Industrial Inpainting Edition +Uses the actual DefectFillModel (model.py) for checkpoint compatibility. +""" + +import os +import sys +import torch +import numpy as np +from typing import Dict, Tuple, Optional, Union +from dataclasses import dataclass +from PIL import Image + +# Ensure model.py and utils.py are importable +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +if SCRIPT_DIR not in sys.path: + sys.path.insert(0, SCRIPT_DIR) + +from model import DefectFillModel +from utils import load_checkpoint + + +@dataclass +class DefectFillConfig: + ckpt_path: str = "" + image_size: int = 512 + num_steps: int = 50 + guidance_scale: float = 7.5 + device: str = "cuda" + torch_dtype: torch.dtype = torch.float16 + seed: int = 42 + lora_rank: int = 8 + lora_alpha: int = 16 + + +class DefectFillGenerator: + """ + Wrapper around the REAL DefectFillModel (model.py). + Compatible with checkpoints saved via utils.save_checkpoint(). + """ + + def __init__(self, config: DefectFillConfig): + self.config = config + self.device = torch.device(config.device) + self.dtype = config.torch_dtype if config.device != "cpu" else torch.float32 + + self.model: Optional[DefectFillModel] = None + self.ckpt_path = config.ckpt_path + self._models_loaded = False + + self._load_models() + + def _infer_lora_rank_from_checkpoint(self, ckpt_path: str) -> int: + """Peek at the checkpoint and infer LoRA rank from any lora_A tensor (handles nested dicts).""" + if not ckpt_path or not os.path.exists(ckpt_path): + return self.config.lora_rank + try: + state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False) + + def _find_rank(d, depth=0): + if depth > 4: + return None + for k, v in d.items(): + if isinstance(v, torch.Tensor) and "lora_A" in k: + return v.shape[0] # lora_A shape is (rank, in_features) + elif isinstance(v, dict): + found = _find_rank(v, depth + 1) + if found is not None: + return found + return None + + inferred_rank = _find_rank(state_dict) + if inferred_rank is not None: + print(f"[DefectFill] Inferred LoRA rank {inferred_rank} from checkpoint") + return inferred_rank + else: + print("[DefectFill] Could not find lora_A in checkpoint; using config lora_rank") + except Exception as e: + print(f"[DefectFill] Could not infer LoRA rank from checkpoint: {e}") + return self.config.lora_rank + + def _load_models(self): + print(f"[DefectFill] Initializing DefectFillModel on {self.device}...") + + # Auto-detect rank so checkpoints trained with r=16 load correctly + effective_rank = self._infer_lora_rank_from_checkpoint(self.ckpt_path) + if effective_rank != self.config.lora_rank: + print(f"[DefectFill] Checkpoint rank mismatch detected. " + f"Overriding lora_rank {self.config.lora_rank} -> {effective_rank}") + + self.model = DefectFillModel( + device=self.device, + lora_rank=effective_rank, + lora_alpha=self.config.lora_alpha, + seed=self.config.seed + ) + + # Move VAE to correct dtype + self.model.pipeline.vae.to(dtype=self.dtype) + + # Load checkpoint if provided + if self.ckpt_path and os.path.exists(self.ckpt_path): + print(f"[DefectFill] Loading checkpoint: {self.ckpt_path}") + load_checkpoint(self.model, None, self.ckpt_path) + else: + print(f"[DefectFill] Warning: No checkpoint found at {self.ckpt_path}. Using base model.") + + self.model.pipeline.unet.eval() + self.model.pipeline.text_encoder.eval() + + self._models_loaded = True + print("[DefectFill] Model initialization complete.") + + def _format_image(self, input_img: Union[np.ndarray, Image.Image, str]) -> torch.Tensor: + """Convert input to torch tensor [1, 3, H, W] in [-1, 1].""" + if isinstance(input_img, str): + img = Image.open(input_img).convert("RGB") + elif isinstance(input_img, np.ndarray): + img = Image.fromarray(input_img).convert("RGB") + elif isinstance(input_img, Image.Image): + img = input_img.convert("RGB") + else: + raise TypeError(f"Unsupported image input type: {type(input_img)}") + + if img.size != (self.config.image_size, self.config.image_size): + img = img.resize((self.config.image_size, self.config.image_size), Image.LANCZOS) + + img_np = np.array(img).astype(np.float32) / 255.0 # [0,1] + img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0) # [1,3,H,W] + img_tensor = img_tensor * 2.0 - 1.0 # [-1,1] + return img_tensor.to(device=self.device, dtype=self.dtype) + + def _format_mask(self, input_mask: Union[np.ndarray, Image.Image]) -> torch.Tensor: + """Convert mask to torch tensor [1, 1, H, W] in [0, 1].""" + if isinstance(input_mask, np.ndarray): + mask_np = input_mask + elif isinstance(input_mask, Image.Image): + mask_np = np.array(input_mask) + else: + raise TypeError(f"Unsupported mask input type: {type(input_mask)}") + + mask_binary = (mask_np > 0).astype(np.float32) + mask_pil = Image.fromarray((mask_binary * 255).astype(np.uint8)).convert("L") + + if mask_pil.size != (self.config.image_size, self.config.image_size): + mask_pil = mask_pil.resize((self.config.image_size, self.config.image_size), Image.NEAREST) + + mask_np = np.array(mask_pil).astype(np.float32) / 255.0 + mask_tensor = torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0) # [1,1,H,W] + return mask_tensor.to(device=self.device, dtype=self.dtype) + + @torch.no_grad() + def inpaint( + self, + image: Union[np.ndarray, Image.Image, str], + mask: Union[np.ndarray, Image.Image], + prompt: str, + defect_type: Optional[str] = None, + seed: Optional[int] = None, + guidance_scale: Optional[float] = None, + num_steps: Optional[int] = None + ) -> Image.Image: + if not self._models_loaded: + self._load_models() + + img_tensor = self._format_image(image) + mask_tensor = self._format_mask(mask) + + num_inference_steps = num_steps if num_steps is not None else self.config.num_steps + guidance = guidance_scale if guidance_scale is not None else self.config.guidance_scale + current_seed = seed if seed is not None else self.config.seed + + generator = torch.Generator(device=self.device).manual_seed(current_seed) + + print(f"[DefectFill] Generating: prompt='{prompt}', steps={num_inference_steps}, seed={current_seed}") + + # Use the custom generate() from DefectFillModel (9-ch input, CFG, bg preservation) + output_tensor = self.model.generate( + image=img_tensor, + mask=mask_tensor, + prompt=prompt, + num_inference_steps=num_inference_steps, + guidance_scale=guidance, + generator=generator + ) # Returns [0, 1] + + # Convert to PIL + output_np = output_tensor.squeeze(0).permute(1, 2, 0).cpu().float().numpy() + output_np = (output_np * 255).clip(0, 255).astype(np.uint8) + return Image.fromarray(output_np) + + @torch.no_grad() + def generate( + self, + image: Union[np.ndarray, Image.Image, str], + mask: Union[np.ndarray, Image.Image], + prompt: str, + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray]: + pil_mask, bool_mask = self._format_mask_for_return(mask) + inpainted_image = self.inpaint(image=image, mask=pil_mask if isinstance(mask, Image.Image) else mask, + prompt=prompt, seed=seed) + return inpainted_image, bool_mask + + def _format_mask_for_return(self, input_mask: Union[np.ndarray, Image.Image]) -> Tuple[Image.Image, np.ndarray]: + if isinstance(input_mask, np.ndarray): + mask_np = input_mask + elif isinstance(input_mask, Image.Image): + mask_np = np.array(input_mask) + else: + raise TypeError(f"Unsupported mask input type: {type(input_mask)}") + + mask_binary = (mask_np > 0).astype(np.uint8) * 255 + mask_pil = Image.fromarray(mask_binary).convert("L") + + if mask_pil.size != (self.config.image_size, self.config.image_size): + mask_pil = mask_pil.resize((self.config.image_size, self.config.image_size), Image.NEAREST) + mask_binary = np.array(mask_pil) > 0 + + return mask_pil, (mask_binary > 0) + + @torch.no_grad() + def generate_from_plan( + self, + product_description: str, + defect_description: str, + image_patch: Union[np.ndarray, Image.Image], + mask_patch: Union[np.ndarray, Image.Image], + seed: Optional[int] = None + ) -> Tuple[Image.Image, np.ndarray, Dict]: + prompt = f"a photo of {product_description} with {defect_description} defect" + inpainted_patch, bool_mask = self.generate( + image=image_patch, + mask=mask_patch, + prompt=prompt, + seed=seed + ) + meta = { + "prompt": prompt, + "product_description": product_description, + "defect_description": defect_description, + "seed": seed or self.config.seed + } + return inpainted_patch, bool_mask, meta + + def unload_models(self): + self.model = None + self._models_loaded = False + if self.device.type == "cuda": + torch.cuda.empty_cache() + print("[DefectFill] Models unloaded from memory.") + + +def get_defectfill_generator( + ckpt_path: str = "", + device: str = "cuda", + **kwargs +) -> DefectFillGenerator: + config = DefectFillConfig( + ckpt_path=ckpt_path, + device=device, + **kwargs + ) + return DefectFillGenerator(config) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/domain_router.py b/ArtiAgent - DefectFill/src/pipeline/domain_router.py new file mode 100644 index 0000000000000000000000000000000000000000..65a2c95169e997d27bbcb33207150bad3c049010 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/domain_router.py @@ -0,0 +1,97 @@ +# src/pipeline/domain_router.py +# ============================================ +# Routes queries to correct collection based on product description +# ============================================ + +from typing import Optional + + +class DomainRouter: + """ + Routes defect queries to domain-specific DB collections or filters. + + Usage: + router = DomainRouter() + domain = router.route("VCSEL laser diode in TO-can package...") + # Returns: "glass/vcsel" + """ + + # Keyword โ†’ domain mapping + DOMAIN_MAP = { + # VCSEL / Optoelectronics + "vcsel": "glass/vcsel", + "laser diode": "glass/vcsel", + "to-can": "glass/vcsel", + "to can": "glass/vcsel", + "glass lens": "glass/vcsel", + "emission aperture": "glass/vcsel", + "bond pad": "glass/vcsel", + "optoelectronic": "glass/vcsel", + + # PCB + "pcb": "pcb", + "printed circuit": "pcb", + "solder joint": "pcb", + "smt": "pcb", + "surface mount": "pcb", + "trace": "pcb", + "pad": "pcb", + + # Metal / Steel + "steel": "metal/steel", + "sheet metal": "metal/steel", + "rolled": "metal/steel", + "metal surface": "metal/steel", + + # Semiconductor + "wafer": "semiconductor", + "die": "semiconductor", + "chip": "semiconductor", + "silicon": "semiconductor", + } + + def __init__(self): + self._cache = {} + + def route(self, product_desc: str) -> str: + """Determine domain from product description.""" + if not product_desc: + return "general" + + desc_lower = product_desc.lower() + + # Check cache + if desc_lower in self._cache: + return self._cache[desc_lower] + + # Match keywords + for keyword, domain in sorted(self.DOMAIN_MAP.items(), key=lambda x: -len(x[0])): + if keyword in desc_lower: + self._cache[desc_lower] = domain + return domain + + self._cache[desc_lower] = "general" + return "general" + + def get_collection_name(self, domain: str) -> str: + """Map domain to DB collection name.""" + # All domains share one collection with domain metadata + # Or use separate collections if needed + return "defect_patches" + + def get_domain_filter(self, product_desc: str) -> Optional[str]: + """Get domain filter for RAG query.""" + domain = self.route(product_desc) + if domain == "general": + return None + return domain + + +# Singleton +_router_instance = None + +def get_router() -> DomainRouter: + global _router_instance + if _router_instance is None: + _router_instance = DomainRouter() + return _router_instance \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/flux_generator.py b/ArtiAgent - DefectFill/src/pipeline/flux_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..ef665818e072980d7bc9afe363e49d5c238f2487 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/flux_generator.py @@ -0,0 +1,438 @@ +import torch +import argparse +from typing import Dict, Optional, Union, List +import numpy as np +from dataclasses import dataclass +import os +import re +import json +import time +from glob import iglob +from einops import rearrange +from PIL import Image + +# FLUX imports +import flux +from flux.sampling import denoise, denoise_first_order, denoise_fireflow, get_schedule, prepare, unpack +from flux.util import (configs, load_ae, load_clip, + load_flow_model, load_t5) + +@dataclass +class FluxConfig: + """Configuration for FLUX model""" + name: str = 'flux-dev' + guidance: float = 5.0 + num_steps: int = 25 + inject_step: int = 15 + pe_step: Union[int, Dict[str, int]] = 25 # Can be int or dict with artifact type keys + attn_mask_step: int = 0 + seed: int = 42 + masks: list = None + alpha: float = 0.0 + feature_path: str = 'feature' + percentage_of_steps: float = 1.0 + offload: bool = False + use_rf_solver: bool = False # Use denoise (RF solver) instead of denoise_first_order + + def __post_init__(self): + if self.masks is None: + self.masks = ['none', 'none', 'none', 'none'] + + # Validate pe_step configuration + if isinstance(self.pe_step, dict): + required_keys = {'addition', 'removal', 'distortion', 'fusion'} + provided_keys = set(self.pe_step.keys()) + if not required_keys.issubset(provided_keys): + missing_keys = required_keys - provided_keys + raise ValueError(f"pe_step dict missing required artifact types: {missing_keys}") + + # Validate all values are integers + for artifact_type, value in self.pe_step.items(): + if not isinstance(value, int): + raise ValueError(f"pe_step value for '{artifact_type}' must be an integer, got {type(value)}") + + def get_pe_step(self, artifact_type: str) -> int: + """ + Get pe_step value for specific artifact type + + Args: + artifact_type: Type of artifact ('addition', 'removal', 'distortion') + + Returns: + pe_step value for the artifact type + """ + if isinstance(self.pe_step, dict): + if artifact_type not in self.pe_step: + raise ValueError(f"Unknown artifact type '{artifact_type}'. Available types: {list(self.pe_step.keys())}") + return self.pe_step[artifact_type] + else: + return self.pe_step + + +class FluxGenerator: + """Handler for FLUX model operations and image generation""" + + def __init__(self, device: str = 'cuda', config: Optional[FluxConfig] = None): + """ + Initialize FLUX generator + + Args: + device: Device to run models on ('cuda' or 'cpu') + config: FLUX configuration object + """ + self.device = device + self.config = config or FluxConfig() + + # Model components + self.t5 = None + self.clip = None + self.model = None + self.ae = None + + self._models_loaded = False + self.load_models() + + def load_models(self): + """Load all FLUX model components""" + print("Loading FLUX models...") + + # Determine max_length based on model name + max_length = 256 if self.config.name == "flux-schnell" else 512 + + # Load model components + self.t5 = load_t5(self.device, max_length=max_length).to(dtype=torch.bfloat16) + self.clip = load_clip(self.device).to(dtype=torch.bfloat16) + self.offload = self.config.offload # store it so you know it's defined + self.model = load_flow_model(self.config.name, device=self.device) + if self.offload: + self.model.enable_sequential_cpu_offload() + self.ae = load_ae(self.config.name, device=self.device) + + self._models_loaded = True + print("FLUX models loaded successfully.") + + def create_default_flux_args(self) -> argparse.Namespace: + """ + Create default FLUX arguments based on current configuration + + Returns: + Default argparse.Namespace object with config values + """ + # Create parser and args + parser = argparse.ArgumentParser() + flux_args = parser.parse_args(args=[]) + + # Set FLUX configuration defaults + flux_args.name = self.config.name + flux_args.feature_path = self.config.feature_path + flux_args.guidance = self.config.guidance + flux_args.num_steps = self.config.num_steps + flux_args.inject_step = self.config.inject_step + flux_args.attn_mask_step = self.config.attn_mask_step + flux_args.pe_step = self.config.pe_step + flux_args.pe_step_addition = self.config.pe_step['addition'] + flux_args.pe_step_removal = self.config.pe_step['removal'] + flux_args.pe_step_distortion = self.config.pe_step['distortion'] + flux_args.pe_step_fusion = self.config.pe_step['fusion'] + flux_args.seed = self.config.seed + flux_args.masks = self.config.masks.copy() + flux_args.alpha = self.config.alpha + flux_args.percentage_of_steps = self.config.percentage_of_steps + flux_args.offload = self.config.offload + + # Initialize task-specific arguments to None + flux_args.source_prompt = None + flux_args.target_prompt = None + flux_args.artifact_type = None + flux_args.output_dir = None + flux_args.source_img = None + + # Initialize optional patch mapping information + flux_args.patch_mapping = None + flux_args.reference_patch_indices = None + flux_args.target_patch_indices = None + + return flux_args + + @torch.inference_mode() + def inject_artifacts(self, + source_prompt: str, + target_prompt: str, + artifact_data: Dict, + source_img: Union[np.ndarray, str], + output_dir: str = None, + pe_step_addition: Optional[int] = None, + pe_step_removal: Optional[int] = None, + pe_step_distortion: Optional[int] = None, + pe_step_fusion: Optional[int] = None, + inject_step: Optional[int] = None, + num_steps: Optional[int] = None, + use_fireflow: bool = False, + reference_images: Optional[List] = None + ): + """ + Sample the flux model with artifact injection supporting arbitrary shapes. + NEW: reference_images โ€” list of example defect images from RAG retrieval + + Args: + source_prompt: Source image prompt/caption + target_prompt: Target prompt for generation + artifact_type: Type of artifact ('addition', 'removal', 'distortion') + source_img: Source image array or path + output_dir: Output directory for generated images + reference_patch_indices: List of reference patch indices + target_patch_indices: List of target patch indices + """ + torch.set_grad_enabled(False) + + # Create default flux args and update with passed parameters + flux_args = self.create_default_flux_args() + + # Update with required parameters + flux_args.source_prompt = source_prompt + flux_args.target_prompt = target_prompt + flux_args.artifact_data = artifact_data + flux_args.source_img = source_img + flux_args.output_dir = output_dir + flux_args.inject_step = inject_step if inject_step is not None else self.config.inject_step + torch_device = torch.device(self.device) + if num_steps is not None: + flux_args.num_steps = num_steps + + init_image = None + init_image = self.load_image(flux_args.source_img) + + shape = init_image.shape + + new_h = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16 + new_w = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16 + + init_image = init_image[:new_h, :new_w, :] + + width, height = init_image.shape[0], init_image.shape[1] + init_image = self.encode(init_image, torch_device, self.ae) + + rng = torch.Generator(device="cpu").manual_seed(flux_args.seed) + + if flux_args.seed is None: + flux_args.seed = rng.seed() + print(f"Generating with seed {flux_args.seed}:\n{flux_args.source_prompt}") + t0 = time.perf_counter() + + flux_args.seed = None + if flux_args.offload: + self.ae = self.ae.cpu() + torch.cuda.empty_cache() + self.t5, self.clip = self.t5.to(torch_device), self.clip.to(torch_device) + + info = {} + info['feature_path'] = flux_args.feature_path + info['feature'] = {} + info['inject_step'] = flux_args.inject_step + info['attn_mask_step'] = flux_args.attn_mask_step + info['alpha'] = flux_args.alpha + if pe_step_distortion is not None: + info['pe_step_distortion'] = pe_step_distortion + else: + info['pe_step_distortion'] = flux_args.pe_step_distortion + if pe_step_removal is not None: + info['pe_step_removal'] = pe_step_removal + else: + info['pe_step_removal'] = flux_args.pe_step_removal + if pe_step_addition is not None: + info['pe_step_addition'] = pe_step_addition + else: + info['pe_step_addition'] = flux_args.pe_step_addition + if pe_step_fusion is not None: + info['pe_step_fusion'] = pe_step_fusion + else: + info['pe_step_fusion'] = flux_args.pe_step_fusion + info['artifact_data'] = flux_args.artifact_data + info['guidance'] = flux_args.guidance + + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + # NEW: Encode RAG reference images and inject into info + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + if reference_images is not None and len(reference_images) > 0: + print(f"[FLUX] Encoding {len(reference_images)} reference image(s) from RAG...") + + # โ”€โ”€ FIX: Ensure AE is on GPU for encoding โ”€โ”€ + self.ae = self.ae.to(torch_device) + + ref_latents = [] + for i, ref_img in enumerate(reference_images): + # Load if path, convert if PIL + if isinstance(ref_img, str): + ref_img = self.load_image(ref_img) + elif isinstance(ref_img, Image.Image): + ref_img = np.array(ref_img.convert('RGB')) + + # Match dimensions to source image + # Resize to match source dimensions for consistent latent encoding + ref_pil = Image.fromarray(ref_img).resize((new_w, new_h), Image.LANCZOS) + ref_img = np.array(ref_pil) + # Encode through VAE (same as source image) + ref_latent = self.encode(ref_img, torch_device, self.ae) + ref_latents.append(ref_latent) + print(f"[FLUX] Reference {i+1} encoded: {ref_latent.shape}") + + info['reference_latents'] = ref_latents + info['num_reference_images'] = len(ref_latents) + + # โ”€โ”€ MOVE AE BACK TO CPU TO SAVE MEMORY DURING DENOISING โ”€โ”€ + if flux_args.offload: + self.ae = self.ae.cpu() + torch.cuda.empty_cache() + else: + info['reference_latents'] = [] + info['num_reference_images'] = 0 + # โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ• + + if not os.path.exists(flux_args.feature_path): + os.mkdir(flux_args.feature_path) + + # Prepare inputs with shape-aware approach + inp, (patch_h, patch_w) = prepare( + self.t5, self.clip, init_image, + prompt=flux_args.source_prompt, + info=info + ) + + inp_target, _ = prepare( + self.t5, self.clip, init_image, + prompt=flux_args.target_prompt, + info=info + ) + + timesteps = get_schedule(flux_args.num_steps, inp["img"].shape[1], shift=(flux_args.name != "flux-schnell")) + + info['patch_h'] = patch_h + info['patch_w'] = patch_w + + L = inp['img'].shape[1] + inp['txt'].shape[1] + + # Choose denoising function based on configuration + # RF solver (denoise) is more accurate but slower than first-order denoising + denoise_func = denoise_fireflow if use_fireflow else denoise_first_order + # denoise_func = denoise_fireflow + # denoise_func = denoise_fireflow + # inversion initial noise + + # 1. Convert any Float32 tensor inputs to bfloat16 to avoid bitsandbytes float32->float16 warnings + inp = { + k: v.to(dtype=torch.bfloat16) if isinstance(v, torch.Tensor) and v.dtype == torch.float32 else v + for k, v in inp.items() + } + + # BEFORE: + # z, info = denoise_func(self.model, **inp, timesteps=timesteps, guidance=1, inverse=True, info=info, percentage_of_steps=flux_args.percentage_of_steps) + # AFTER: + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + z, info = denoise_func( + self.model, + **inp, + timesteps=timesteps, + guidance=1, + inverse=True, + info=info, + percentage_of_steps=flux_args.percentage_of_steps + ) + + inp_target["img"] = z + + timesteps = get_schedule(flux_args.num_steps, inp_target["img"].shape[1], shift=(flux_args.name != "flux-schnell")) + + # denoise initial noise + x, _ = denoise_func(self.model, **inp_target, timesteps=timesteps, guidance=info['guidance'], inverse=False, info=info, percentage_of_steps=flux_args.percentage_of_steps) + + # Bring AE back to GPU for decoding (it was offloaded to CPU earlier) + if flux_args.offload: + self.t5, self.clip = self.t5.cpu(), self.clip.cpu() + torch.cuda.empty_cache() + self.ae = self.ae.to(torch_device) + + # decode latents to pixel space + + ####################################### + #### TODO: allow batch computation #### + ####################################### + + x = unpack(x.float(), width, height) + + if output_dir is not None: + output_name = os.path.join(output_dir, "img.jpg") + if not os.path.exists(output_dir): + os.makedirs(output_dir) + idx = 0 + else: + fns = [fn for fn in iglob(output_name.format(idx="*")) if re.search(r"img_[0-9]+\.jpg$", fn)] + if len(fns) > 0: + idx = max(int(fn.split("_")[-1].split(".")[0]) for fn in fns) + 1 + else: + idx = 0 + + with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16): + x = self.ae.decode(x.float()) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + t1 = time.perf_counter() + + print(f"Done in {t1 - t0:.1f}s.") + + # bring into PIL format and save + x = x.clamp(-1, 1) + x = rearrange(x[0], "c h w -> h w c") + img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy()) + + return img + + def load_image(self, source): + """Load image from various sources (numpy array, PIL Image, or file path)""" + if isinstance(source, np.ndarray): + # Already a NumPy array + return source + elif isinstance(source, Image.Image): + # Already a PIL Image + return np.array(source.convert('RGB')) + elif isinstance(source, str): + if os.path.isfile(source): + # It's a file path to an image + return np.array(Image.open(source).convert('RGB')) + else: + raise ValueError(f"Provided string is not a valid file: {source}") + else: + raise TypeError(f"Unsupported input type: {type(source)}") + + @torch.inference_mode() + def encode(self, init_image, torch_device, ae): + """Encode image to latent space""" + init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1 + init_image = init_image.unsqueeze(0) + init_image = init_image.to(torch_device) + init_image = ae.encode(init_image).to(torch.bfloat16) + return init_image + + def update_config(self, **kwargs): + """Update FLUX configuration parameters""" + for key, value in kwargs.items(): + if hasattr(self.config, key): + setattr(self.config, key, value) + else: + print(f"Warning: Unknown config parameter '{key}'") + + def unload_models(self): + """Unload models to free memory""" + self.t5 = None + self.clip = None + self.model = None + self.ae = None + self._models_loaded = False + + # Clear GPU cache if using CUDA + if self.device == 'cuda' and torch.cuda.is_available(): + torch.cuda.empty_cache() + + def __del__(self): + """Cleanup when object is destroyed""" + self.unload_models() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/gsam_detector - Copy.py b/ArtiAgent - DefectFill/src/pipeline/gsam_detector - Copy.py new file mode 100644 index 0000000000000000000000000000000000000000..16746db931c2c1c85f0355a5a4189952c611acc3 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/gsam_detector - Copy.py @@ -0,0 +1,379 @@ +import sys +import os +import multiprocessing as mp +import numpy as np +from typing import List, Dict, Tuple, Optional, Union +import torch +from PIL import Image +import torchvision +import supervision as sv +from flux.artifacts_util import mask_to_patch_coords +from logging import Logger as logger +from pydantic import BaseModel, RootModel +import json +import re + +# Add GroundingDINO and SAM to path +sys.path.append(os.path.join(os.getcwd(), 'GroundingDINO')) +sys.path.append(os.path.join(os.getcwd(), 'segment_anything')) +sys.path.append(os.path.join(os.getcwd(), 'pipeline')) + +# Grounding DINO +from groundingdino.util.inference import Model + +# S +from segment_anything import ( + sam_model_registry, + sam_hq_model_registry, + SamPredictor +) + +def robust_json_parse(raw_text: str): + """Try to parse VLM JSON output, with automatic repair.""" + # 1. Try direct parse + try: + return json.loads(raw_text) + except json.JSONDecodeError: + pass + + # 2. Extract JSON block from markdown fences + match = re.search(r'```json\s*(.*?)\s*```', raw_text, re.DOTALL) + if match: + try: + return json.loads(match.group(1)) + except json.JSONDecodeError: + pass + + # 3. Find the outermost {...} or [...] + match = re.search(r'(\{.*\}|\[.*\])', raw_text, re.DOTALL) + if match: + candidate = match.group(1) + # Fix common VLM JSON mistakes: + # - trailing commas before ] or } + candidate = re.sub(r',\s*([}\]])', r'\1', candidate) + # - extra closing brackets + while candidate.count('[') < candidate.count(']'): + candidate = candidate[:-1] # remove trailing ] + while candidate.count('{') < candidate.count('}'): + candidate = candidate[:-1] + # - missing closing brace + if candidate.count('{') > candidate.count('}'): + candidate += '}' + try: + return robust_json_parse(candidate) + except json.JSONDecodeError: + pass + + # 4. Fallback: regex extract bboxes directly + bboxes = re.findall(r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]', raw_text) + if bboxes: + return {"bboxes": [[int(x) for x in box] for box in bboxes]} + + raise ValueError(f"Could not parse JSON from: {raw_text[:200]}") + +class GSAMDetector: + """Handler for Grounded SAM part detection model""" + + # Constants for better code maintainability + DEFAULT_CONTAINMENT_THRESHOLD = 0.9 + DEFAULT_MIN_AREA_RATIO = 0.005 + DEFAULT_MAX_AREA_RATIO = 0.5 + + def __init__(self, + grounding_config_file: Optional[str] = None, + grounding_checkpoint: Optional[str] = None, + sam_version: str = "vit_h", + sam_checkpoint: Optional[str] = None, + sam_hq_checkpoint: Optional[str] = None, + use_sam_hq: bool = False, + box_threshold: float = 0.3, + text_threshold: float = 0.25, + nms_threshold: float = 0.5, + bert_base_uncased_path: Optional[str] = None, + device: str = "cuda", + openai_client: Optional[any] = None + ): + """ + Initialize GSAM detector + + Args: + grounding_config_file: Path to GroundingDINO config file + grounding_checkpoint: Path to GroundingDINO checkpoint + sam_version: SAM model version (vit_b, vit_l, vit_h) + sam_checkpoint: Path to SAM checkpoint + sam_hq_checkpoint: Path to SAM-HQ checkpoint + use_sam_hq: Whether to use SAM-HQ + box_threshold: Box threshold for detection + text_threshold: Text threshold for detection + nms_threshold: NMS threshold for detection + bert_base_uncased_path: Path to BERT model + device: Device to use (cuda/cpu) + openai_client: OpenAI client for vocabulary generation + """ + self.gsam_path = os.getcwd() + + # Set default paths if not provided + if grounding_config_file is None: + grounding_config_file = os.path.join(self.gsam_path, "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py") + if grounding_checkpoint is None: + grounding_checkpoint = os.path.join(self.gsam_path, "weight/groundingdino_swint_ogc.pth") + if sam_checkpoint is None and not use_sam_hq: + sam_checkpoint = os.path.join(self.gsam_path, "weight/sam_vit_h_4b8939.pth") + if sam_hq_checkpoint is None and use_sam_hq: + sam_hq_checkpoint = os.path.join(self.gsam_path, "weight/sam_hq_vit_h.pth") + + self.grounding_config_file = grounding_config_file + self.grounding_checkpoint = grounding_checkpoint + self.sam_version = sam_version + self.sam_checkpoint = sam_checkpoint + self.sam_hq_checkpoint = sam_hq_checkpoint + self.nms_threshold = nms_threshold + self.box_threshold = box_threshold + self.text_threshold = text_threshold + self.device = device + self.openai_client = openai_client + + # Model components + self.grounding_model = Model(model_config_path=self.grounding_config_file, model_checkpoint_path=self.grounding_checkpoint) + if use_sam_hq: + self.sam_predictor = SamPredictor(sam_hq_model_registry[self.sam_version](checkpoint=self.sam_hq_checkpoint).to(self.device)) + else: + self.sam_predictor = SamPredictor(sam_model_registry[self.sam_version](checkpoint=self.sam_checkpoint).to(self.device)) + + # Set multiprocessing start method + mp.set_start_method('spawn', force=True) + + + # Prompting SAM with detected boxes (same as original) + def segment(self, image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: + self.sam_predictor.set_image(image) + result_masks = [] + for box in xyxy: + masks, scores, logits = self.sam_predictor.predict( + box=box, + multimask_output=True + ) + index = np.argmax(scores) + result_masks.append(masks[index]) + return np.array(result_masks) + + def detect_parts(self, image, entities, subentities, entity_subentity_mapping, + min_area_ratio=0.005, max_area_ratio=0.5, openai_client=None): + """ + Detect parts using VLM for bboxes + SAM for masks. + Replaces GroundingDINO with VLM-guided detection. + """ + import cv2 + h, w = image.shape[:2] + total_area = h * w + + predictions = [] + entity_predictions = [] + + # Use VLM to get bboxes for each entity + from prompts import get_entity_bboxes + + all_bboxes = {} # entity -> list of bboxes + + for entity in entities: + bboxes = get_entity_bboxes(openai_client, image, entity) + if bboxes: + all_bboxes[entity] = bboxes + print(f"VLM found {len(bboxes)} instances of '{entity}'") + + # Use SAM to get masks from bboxes + self.sam_predictor.set_image(image) # <-- ADD THIS LINE + for entity, bboxes in all_bboxes.items(): + for bbox in bboxes: + x1, y1, x2, y2 = map(int, bbox) + + # Validate bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(w, x2), min(h, y2) + if x2 <= x1 or y2 <= y1: + continue + + bbox_area = (x2 - x1) * (y2 - y1) + area_ratio = bbox_area / total_area + + if not (min_area_ratio <= area_ratio <= max_area_ratio): + print(f"Discarded '{entity}' bbox - area ratio {area_ratio:.4f} out of range") + continue + + # SAM mask from bbox + input_box = np.array([x1, y1, x2, y2]) + masks, scores, _ = self.sam_predictor.predict( + point_coords=None, + point_labels=None, + box=input_box[None, :], + multimask_output=False + ) + + if masks is None or len(masks) == 0: + continue + + mask = masks[0] if len(masks.shape) == 3 else masks + mask_binary = (mask > 0).astype(np.uint8) + + # Find contours for precise bbox + contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + continue + + cnt = max(contours, key=cv2.contourArea) + ex, ey, ew, eh = cv2.boundingRect(cnt) + + pred = { + 'entity': entity, + 'subentity': entity, # For now, entity = subentity + 'pred_box': torch.tensor([ex, ey, ex + ew, ey + eh]).float(), + 'bbox': [ex, ey, ex + ew, ey + eh], # alias for downstream + 'pred_class': torch.tensor(0).long(), # dummy class + 'pred_mask': torch.from_numpy(mask_binary).bool(), # numpy โ†’ torch tensor # <-- change 'mask' to 'pred_mask' + 'mask': mask_binary, # alias for downstream + 'score': torch.tensor(float(scores[0]) if len(scores) > 0 else 1.0).float(), + 'area_ratio': area_ratio + } + predictions.append(pred) + entity_predictions.append(pred) + + print(f"VLM+SAM: Found {len(predictions)} valid detections") + + # Visualization (optional) + visualized_output = None + try: + import supervision as sv + if predictions: + img_viz = image.copy() if hasattr(image, 'copy') else np.array(image) + # Simple bbox visualization + for p in predictions: + x1, y1, x2, y2 = map(int, p['pred_box'].tolist()) + cv2.rectangle(img_viz, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(img_viz, p['entity'], (x1, y1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + visualized_output = img_viz + except Exception as e: + print(f"Visualization error: {e}") + + return predictions, entity_predictions, visualized_output + + def detect_entities(self, image: np.ndarray, entities: List[str], + min_area_ratio: float = 0.01, max_area_ratio: float = 1.0) -> Tuple[List[Dict], any]: + """ + Run entity detection on image (entities only, no subentities) + + Args: + image: Input image as numpy array (RGB format) + entities: List of entity names + min_area_ratio: Minimum area ratio for filtering + max_area_ratio: Maximum area ratio for filtering + + Returns: + Tuple of (entity_predictions, visualized_output): + - entity_predictions: List of dictionaries, each containing entity detection with keys: + 'pred_box', 'pred_class', 'score', 'pred_mask', 'entity_name' + - visualized_output: PIL Image with annotations + """ + # Store current image size for area calculations + self.current_image_size = image.shape[:2] + + # Get grounding output + detections, phrases = self.grounding_model.predict_with_caption( + image=image, + caption=", ".join(entities), + box_threshold=self.box_threshold, + text_threshold=self.text_threshold + ) + + # Generate class_id from phrases since predict_with_caption doesn't include it + detections.class_id = Model.phrases2classes(phrases=phrases, classes=entities) + + # NMS post process + print(f"Before NMS: {len(detections.xyxy)} boxes") + nms_idx = torchvision.ops.nms( + torch.from_numpy(detections.xyxy), + torch.from_numpy(detections.confidence), + self.nms_threshold + ).numpy().tolist() + + detections.xyxy = detections.xyxy[nms_idx] + detections.confidence = detections.confidence[nms_idx] + detections.class_id = detections.class_id[nms_idx] + # Also filter phrases to match the filtered detections + phrases = [phrases[i] for i in nms_idx] + + detections.mask = self.segment( + image=image, + xyxy=detections.xyxy, + ) + + print(f"Found {len(detections.class_id)} entity detections") + + # Filter entities by area ratio + filtered_entities = [] + image_area = self.current_image_size[0] * self.current_image_size[1] + + for i in range(len(detections.class_id)): + entity_mask = torch.from_numpy(detections.mask[i]) + entity_class = detections.class_id[i] + entity_name = entities[entity_class] + area_ratio = torch.sum(entity_mask > 0) / image_area + + if min_area_ratio <= area_ratio <= max_area_ratio: + filtered_entities.append(i) + print(f"Kept entity '{entity_name}' (class {entity_class}) with area ratio {area_ratio:.4f}") + else: + print(f"Discarded entity '{entity_name}' (class {entity_class}) - area ratio {area_ratio:.4f} outside range [{min_area_ratio}, {max_area_ratio}]") + + if len(filtered_entities) == 0: + raise ValueError("No entities detected after filtering") + + # Filter detections to keep only valid entities + filtered_xyxy = detections.xyxy[filtered_entities] + filtered_confidence = detections.confidence[filtered_entities] + filtered_class_id = detections.class_id[filtered_entities] + filtered_mask = detections.mask[filtered_entities] + + # Create filtered detections object for annotation + filtered_detections = sv.Detections( + xyxy=filtered_xyxy, + confidence=filtered_confidence, + class_id=filtered_class_id, + mask=filtered_mask + ) + + # Annotate image with filtered detections + box_annotator = sv.BoundingBoxAnnotator() + mask_annotator = sv.MaskAnnotator() + label_annotator = sv.LabelAnnotator() + + labels = [ + f"{entities[class_id]} {confidence:0.2f}" + for class_id, confidence in zip(filtered_class_id, filtered_confidence)] + + annotated_image = mask_annotator.annotate(scene=image.copy(), detections=filtered_detections) + annotated_image = box_annotator.annotate(scene=annotated_image, detections=filtered_detections) + annotated_image = label_annotator.annotate(scene=annotated_image, detections=filtered_detections, labels=labels) + # Convert annotated image to PIL Image + annotated_image = Image.fromarray(annotated_image) + + # Create entity predictions as list of dictionaries + entity_predictions = [] + for i, entity_idx in enumerate(filtered_entities): + entity_pred_instance = { + 'pred_box': torch.from_numpy(detections.xyxy[entity_idx]).float(), + 'pred_class': torch.tensor(detections.class_id[entity_idx]).long(), + 'score': torch.tensor(detections.confidence[entity_idx]).float(), + 'pred_mask': torch.from_numpy(detections.mask[entity_idx]).bool(), + 'entity': entities[detections.class_id[entity_idx]], + } + entity_predictions.append(entity_pred_instance) + + print(f"Returning {len(entity_predictions)} entity detections") + return entity_predictions, annotated_image + + def cleanup(self): + """Clean up model resources""" + self.grounding_model = None + self.sam_predictor = None + self.current_vocabulary = [] \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/gsam_detector.py b/ArtiAgent - DefectFill/src/pipeline/gsam_detector.py new file mode 100644 index 0000000000000000000000000000000000000000..0e9248a5c87afc72c7f2490d90c09bef19df6dc6 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/gsam_detector.py @@ -0,0 +1,716 @@ +import sys +import os +import multiprocessing as mp +import numpy as np +from typing import List, Dict, Tuple, Optional, Union +import torch +from PIL import Image +import torchvision +import supervision as sv +from flux.artifacts_util import mask_to_patch_coords +from logging import Logger as logger +from pydantic import BaseModel, RootModel +import json +import re + +# Add GroundingDINO and SAM to path +sys.path.append(os.path.join(os.getcwd(), 'GroundingDINO')) +sys.path.append(os.path.join(os.getcwd(), 'segment_anything')) +sys.path.append(os.path.join(os.getcwd(), 'pipeline')) + +# Grounding DINO +from groundingdino.util.inference import Model + +# S +from segment_anything import ( + sam_model_registry, + sam_hq_model_registry, + SamPredictor +) + +def robust_json_parse(raw_text: str): + """Try to parse VLM JSON output, with automatic repair.""" + # 1. Try direct parse + try: + return json.loads(raw_text) + except json.JSONDecodeError: + pass + + # 2. Extract JSON block from markdown fences + match = re.search(r'```json\s*(.*?)\s*```', raw_text, re.DOTALL) + if match: + try: + return json.loads(match.group(1)) + except json.JSONDecodeError: + pass + + # 3. Find the outermost {...} or [...] + match = re.search(r'(\{.*\}|\[.*\])', raw_text, re.DOTALL) + if match: + candidate = match.group(1) + # Fix common VLM JSON mistakes: + # - trailing commas before ] or } + candidate = re.sub(r',\s*([}\]])', r'\1', candidate) + # - extra closing brackets + while candidate.count('[') < candidate.count(']'): + candidate = candidate[:-1] # remove trailing ] + while candidate.count('{') < candidate.count('}'): + candidate = candidate[:-1] + # - missing closing brace + if candidate.count('{') > candidate.count('}'): + candidate += '}' + try: + return robust_json_parse(candidate) + except json.JSONDecodeError: + pass + + # 4. Fallback: regex extract bboxes directly + bboxes = re.findall(r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]', raw_text) + if bboxes: + return {"bboxes": [[int(x) for x in box] for box in bboxes]} + + raise ValueError(f"Could not parse JSON from: {raw_text[:200]}") + +class GSAMDetector: + """Handler for Grounded SAM part detection model""" + + # Constants for better code maintainability + DEFAULT_CONTAINMENT_THRESHOLD = 0.9 + DEFAULT_MIN_AREA_RATIO = 0.005 + DEFAULT_MAX_AREA_RATIO = 0.5 + + def __init__(self, + grounding_config_file: Optional[str] = None, + grounding_checkpoint: Optional[str] = None, + sam_version: str = "vit_h", + sam_checkpoint: Optional[str] = None, + sam_hq_checkpoint: Optional[str] = None, + use_sam_hq: bool = False, + box_threshold: float = 0.3, + text_threshold: float = 0.25, + nms_threshold: float = 0.5, + bert_base_uncased_path: Optional[str] = None, + device: str = "cuda", + openai_client: Optional[any] = None + ): + """ + Initialize GSAM detector + + Args: + grounding_config_file: Path to GroundingDINO config file + grounding_checkpoint: Path to GroundingDINO checkpoint + sam_version: SAM model version (vit_b, vit_l, vit_h) + sam_checkpoint: Path to SAM checkpoint + sam_hq_checkpoint: Path to SAM-HQ checkpoint + use_sam_hq: Whether to use SAM-HQ + box_threshold: Box threshold for detection + text_threshold: Text threshold for detection + nms_threshold: NMS threshold for detection + bert_base_uncased_path: Path to BERT model + device: Device to use (cuda/cpu) + openai_client: OpenAI client for vocabulary generation + """ + # OLD โ€” fragile, depends on CWD + # self.gsam_path = os.getcwd() + # NEW โ€” relative to gsam_detector.py location (src/pipeline/) + # GroundingDINO is in src/ (one level up) + # weight/ is at project root (two levels up) + self.gsam_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + # Set default paths if not provided + if grounding_config_file is None: + grounding_config_file = os.path.join(self.gsam_path, "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py") + if grounding_checkpoint is None: + grounding_checkpoint = os.path.join(self.gsam_path, "weight/groundingdino_swint_ogc.pth") + if sam_checkpoint is None and not use_sam_hq: + sam_checkpoint = os.path.join(self.gsam_path, "weight/sam_vit_h_4b8939.pth") + if sam_hq_checkpoint is None and use_sam_hq: + sam_hq_checkpoint = os.path.join(self.gsam_path, "weight/sam_hq_vit_h.pth") + + self.grounding_config_file = grounding_config_file + self.grounding_checkpoint = grounding_checkpoint + self.sam_version = sam_version + self.sam_checkpoint = sam_checkpoint + self.sam_hq_checkpoint = sam_hq_checkpoint + self.nms_threshold = nms_threshold + self.box_threshold = box_threshold + self.text_threshold = text_threshold + self.device = device + self.openai_client = openai_client + + # Model components + self.grounding_model = Model(model_config_path=self.grounding_config_file, model_checkpoint_path=self.grounding_checkpoint) + if use_sam_hq: + self.sam_predictor = SamPredictor(sam_hq_model_registry[self.sam_version](checkpoint=self.sam_hq_checkpoint).to(self.device)) + else: + self.sam_predictor = SamPredictor(sam_model_registry[self.sam_version](checkpoint=self.sam_checkpoint).to(self.device)) + + # Set multiprocessing start method + mp.set_start_method('spawn', force=True) + + + # Prompting SAM with detected boxes (same as original) + def segment(self, image: np.ndarray, xyxy: np.ndarray) -> np.ndarray: + self.sam_predictor.set_image(image) + result_masks = [] + for box in xyxy: + masks, scores, logits = self.sam_predictor.predict( + box=box, + multimask_output=True + ) + index = np.argmax(scores) + result_masks.append(masks[index]) + return np.array(result_masks) + + def detect_parts(self, image, entities, subentities, entity_subentity_mapping, location_hint=None, + min_area_ratio=0.005, max_area_ratio=0.5, openai_client=None): + """ + Detect parts using VLM for bboxes + SAM for masks. + Replaces GroundingDINO with VLM-guided detection. + """ + import cv2 + h, w = image.shape[:2] + total_area = h * w + + predictions = [] + entity_predictions = [] + + # Use VLM to get bboxes for each entity + from prompts import get_entity_bboxes + + all_bboxes = {} # entity -> list of bboxes + + for entity in entities: + bboxes = get_entity_bboxes(openai_client, image, entity, subentity=subentities[0] if subentities else None, + location_hint=location_hint) + + if bboxes: + all_bboxes[entity] = bboxes + print(f"VLM found {len(bboxes)} instances of '{entity}'") + + # Use SAM to get masks from bboxes + self.sam_predictor.set_image(image) # <-- ADD THIS LINE + for entity, bboxes in all_bboxes.items(): + for bbox in bboxes: + x1, y1, x2, y2 = map(int, bbox) + + # Validate bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(w, x2), min(h, y2) + if x2 <= x1 or y2 <= y1: + continue + + bbox_area = (x2 - x1) * (y2 - y1) + area_ratio = bbox_area / total_area + + if not (min_area_ratio <= area_ratio <= max_area_ratio): + print(f"Discarded '{entity}' bbox - area ratio {area_ratio:.4f} out of range") + continue + + # SAM mask from bbox + input_box = np.array([x1, y1, x2, y2]) + masks, scores, _ = self.sam_predictor.predict( + point_coords=None, + point_labels=None, + box=input_box[None, :], + multimask_output=False + ) + + if masks is None or len(masks) == 0: + continue + + mask = masks[0] if len(masks.shape) == 3 else masks + mask_binary = (mask > 0).astype(np.uint8) + + # Find contours for precise bbox + contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + continue + + cnt = max(contours, key=cv2.contourArea) + ex, ey, ew, eh = cv2.boundingRect(cnt) + + pred = { + 'entity': entity, + 'subentity': entity, # For now, entity = subentity + 'pred_box': torch.tensor([ex, ey, ex + ew, ey + eh]).float(), + 'bbox': [ex, ey, ex + ew, ey + eh], # alias for downstream + 'pred_class': torch.tensor(0).long(), # dummy class + 'pred_mask': torch.from_numpy(mask_binary).bool(), # numpy โ†’ torch tensor # <-- change 'mask' to 'pred_mask' + 'mask': mask_binary, # alias for downstream + 'score': torch.tensor(float(scores[0]) if len(scores) > 0 else 1.0).float(), + 'area_ratio': area_ratio + } + predictions.append(pred) + entity_predictions.append(pred) + + print(f"VLM+SAM: Found {len(predictions)} valid detections") + + # Visualization (optional) + visualized_output = None + try: + import supervision as sv + if predictions: + img_viz = image.copy() if hasattr(image, 'copy') else np.array(image) + # Simple bbox visualization + for p in predictions: + x1, y1, x2, y2 = map(int, p['pred_box'].tolist()) + cv2.rectangle(img_viz, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(img_viz, p['entity'], (x1, y1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + visualized_output = img_viz + except Exception as e: + print(f"Visualization error: {e}") + + return predictions, entity_predictions, visualized_output + + def detect_entities(self, image: np.ndarray, entities: List[str], + min_area_ratio: float = 0.01, max_area_ratio: float = 1.0) -> Tuple[List[Dict], any]: + """ + Run entity detection on image (entities only, no subentities) + + Args: + image: Input image as numpy array (RGB format) + entities: List of entity names + min_area_ratio: Minimum area ratio for filtering + max_area_ratio: Maximum area ratio for filtering + + Returns: + Tuple of (entity_predictions, visualized_output): + - entity_predictions: List of dictionaries, each containing entity detection with keys: + 'pred_box', 'pred_class', 'score', 'pred_mask', 'entity_name' + - visualized_output: PIL Image with annotations + """ + # Store current image size for area calculations + self.current_image_size = image.shape[:2] + + # Get grounding output + detections, phrases = self.grounding_model.predict_with_caption( + image=image, + caption=", ".join(entities), + box_threshold=self.box_threshold, + text_threshold=self.text_threshold + ) + + # Generate class_id from phrases since predict_with_caption doesn't include it + detections.class_id = Model.phrases2classes(phrases=phrases, classes=entities) + + # NMS post process + print(f"Before NMS: {len(detections.xyxy)} boxes") + nms_idx = torchvision.ops.nms( + torch.from_numpy(detections.xyxy), + torch.from_numpy(detections.confidence), + self.nms_threshold + ).numpy().tolist() + + detections.xyxy = detections.xyxy[nms_idx] + detections.confidence = detections.confidence[nms_idx] + detections.class_id = detections.class_id[nms_idx] + # Also filter phrases to match the filtered detections + phrases = [phrases[i] for i in nms_idx] + + detections.mask = self.segment( + image=image, + xyxy=detections.xyxy, + ) + + print(f"Found {len(detections.class_id)} entity detections") + + # Filter entities by area ratio + filtered_entities = [] + image_area = self.current_image_size[0] * self.current_image_size[1] + + for i in range(len(detections.class_id)): + entity_mask = torch.from_numpy(detections.mask[i]) + entity_class = detections.class_id[i] + entity_name = entities[entity_class] + area_ratio = torch.sum(entity_mask > 0) / image_area + + if min_area_ratio <= area_ratio <= max_area_ratio: + filtered_entities.append(i) + print(f"Kept entity '{entity_name}' (class {entity_class}) with area ratio {area_ratio:.4f}") + else: + print(f"Discarded entity '{entity_name}' (class {entity_class}) - area ratio {area_ratio:.4f} outside range [{min_area_ratio}, {max_area_ratio}]") + + if len(filtered_entities) == 0: + raise ValueError("No entities detected after filtering") + + # Filter detections to keep only valid entities + filtered_xyxy = detections.xyxy[filtered_entities] + filtered_confidence = detections.confidence[filtered_entities] + filtered_class_id = detections.class_id[filtered_entities] + filtered_mask = detections.mask[filtered_entities] + + # Create filtered detections object for annotation + filtered_detections = sv.Detections( + xyxy=filtered_xyxy, + confidence=filtered_confidence, + class_id=filtered_class_id, + mask=filtered_mask + ) + + # Annotate image with filtered detections + box_annotator = sv.BoundingBoxAnnotator() + mask_annotator = sv.MaskAnnotator() + label_annotator = sv.LabelAnnotator() + + labels = [ + f"{entities[class_id]} {confidence:0.2f}" + for class_id, confidence in zip(filtered_class_id, filtered_confidence)] + + annotated_image = mask_annotator.annotate(scene=image.copy(), detections=filtered_detections) + annotated_image = box_annotator.annotate(scene=annotated_image, detections=filtered_detections) + annotated_image = label_annotator.annotate(scene=annotated_image, detections=filtered_detections, labels=labels) + # Convert annotated image to PIL Image + annotated_image = Image.fromarray(annotated_image) + + # Create entity predictions as list of dictionaries + entity_predictions = [] + for i, entity_idx in enumerate(filtered_entities): + entity_pred_instance = { + 'pred_box': torch.from_numpy(detections.xyxy[entity_idx]).float(), + 'pred_class': torch.tensor(detections.class_id[entity_idx]).long(), + 'score': torch.tensor(detections.confidence[entity_idx]).float(), + 'pred_mask': torch.from_numpy(detections.mask[entity_idx]).bool(), + 'entity': entities[detections.class_id[entity_idx]], + } + entity_predictions.append(entity_pred_instance) + + print(f"Returning {len(entity_predictions)} entity detections") + return entity_predictions, annotated_image + + def cleanup(self): + """Clean up model resources""" + self.grounding_model = None + self.sam_predictor = None + self.current_vocabulary = [] + + def detect_feature_array( + self, + image: np.ndarray, + feature_type: str = "dots", # "dots" | "lines" | "rectangles" | "single dot" + blob_color: int = 0, # 0=dark features, 255=bright + min_area: Optional[float] = None, # keep for backward compatibility + max_area: Optional[float] = None, + min_area_ratio: float = 0.2, # new + max_area_ratio: float = 0.5, # new + min_circularity: float = 0.5, + min_convexity: float = 0.6, + min_inertia_ratio: float = 0.1, + min_dist_between_blobs: float = 4, + pad_x: int = 60, + pad_y: int = 40, + max_y_span: int = 90, + min_cluster_size: int = 4, + aspect_ratio_range: Optional[Tuple[float, float]] = None, # (h/w min, h/w max) + vertical_align_threshold: float = 0.3, + entity_name: str = "feature_array", + use_adaptive_threshold: bool = False, + morph_kernel_size: int =3, + location_hint: Optional[str] = None # <--- NEW: pass location to pick the right single blob + ): + """ + Generic feature detector for industrial components. + + Modes: + - "dots" : Rounded arrays (solder balls, BGA, die pads) + - "lines" : Elongated features (pins, legs, leads, traces) + - "rectangles": Blocky regions (die body, package, pads) + """ + import cv2 + h, w = image.shape[:2] + total_pixels = h * w + gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) if len(image.shape) == 3 else image.copy() + + # --- Resolve area constraints --- + if min_area is None: + min_area = int(total_pixels * min_area_ratio) + if max_area is None: + max_area = int(total_pixels * max_area_ratio) + + # ===================================================================== + # MODE 1: Blob detection (good for dots / circles) + # ===================================================================== + if feature_type == "dots": + params = cv2.SimpleBlobDetector_Params() + params.filterByColor = True + params.blobColor = blob_color + params.minThreshold = 5 + params.maxThreshold = 150 + params.thresholdStep = 10 + params.minRepeatability = 2 + params.minDistBetweenBlobs = min_dist_between_blobs + + params.filterByArea = True + params.minArea = min_area + params.maxArea = max_area + + params.filterByCircularity = True + params.minCircularity = min_circularity + + params.filterByConvexity = True + params.minConvexity = min_convexity + + params.filterByInertia = True + params.minInertiaRatio = min_inertia_ratio + + detector = cv2.SimpleBlobDetector_create(params) + keypoints = detector.detect(gray) + + if len(keypoints) < min_cluster_size: + return [], [], None + + pts = np.array([[k.pt[0], k.pt[1]] for k in keypoints]) + + # Cluster by vertical proximity (densest group) + y_sorted = pts[np.argsort(pts[:, 1])] + y_vals = y_sorted[:, 1] + + best_cluster = [] + best_count = 0 + for i in range(len(y_vals)): + cluster = [y_sorted[i]] + for j in range(i + 1, len(y_vals)): + if y_vals[j] - y_vals[i] < max_y_span: + cluster.append(y_sorted[j]) + else: + break + if len(cluster) > best_count: + best_count = len(cluster) + best_cluster = cluster + + if len(best_cluster) < min_cluster_size: + return [], [], None + + valid = np.array(best_cluster) + + # ===================================================================== + # MODE 2: Contour detection (good for lines / legs / rectangles) + # ===================================================================== + elif feature_type in ("lines", "rectangles"): + # Threshold + # --- NEW: Adaptive Threshold for thin lines --- + gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + # Edge detection before Hough is usually better than thresholding + edges = cv2.Canny(gray, 50, 150, apertureSize=3) + + # Probabilistic Hough Lines + lines = cv2.HoughLinesP( + edges, + rho=1, + theta=np.pi/180, + threshold=50, # Minimum votes to consider a line + minLineLength=30, # Minimum length of a line (legs are tall) + maxLineGap=10 # Max gap to connect broken leg segments + ) + + if lines is not None and len(lines) >= min_cluster_size: + # Collect all the endpoints of the detected line segments + all_points = [] + for line in lines: + # --- SAFE CHANGE: Use flatten() --- + x1, y1, x2, y2 = line.flatten() + all_points.append([x1, y1]) + all_points.append([x2, y2]) + pts = np.array(all_points) + + # Cluster by X-position to group the legs + x_sorted = pts[np.argsort(pts[:, 0])] + x_vals = x_sorted[:, 0] + x_gap = w * vertical_align_threshold + + best_cluster = [] + best_count = 0 + for i in range(len(x_vals)): + cluster = [x_sorted[i]] + for j in range(i + 1, len(x_vals)): + if x_vals[j] - x_vals[i] < x_gap: + cluster.append(x_sorted[j]) + else: + break + if len(cluster) > best_count: + best_count = len(cluster) + best_cluster = cluster + + if len(best_cluster) >= min_cluster_size: + valid = np.array(best_cluster) # Successful assignment + + # ===================================================================== + # NEW MODE: Single dot / isolated circle (Perfect for one screw/hole) + # ===================================================================== + elif feature_type == "single_dot": + params = cv2.SimpleBlobDetector_Params() + params.filterByColor = True + params.blobColor = blob_color # 0 for dark empty holes, 255 for shiny screw heads + params.minThreshold = 10 + params.maxThreshold = 200 + params.thresholdStep = 10 + params.minRepeatability = 2 + params.minDistBetweenBlobs = 10 # Spacing, but we'll take best one + + params.filterByArea = True + params.minArea = min_area + params.maxArea = max_area + + params.filterByCircularity = True + params.minCircularity = 0.4 # Screws/holes are quite round + + params.filterByConvexity = True + params.minConvexity = 0.5 + + params.filterByInertia = True + params.minInertiaRatio = 0.1 + + detector = cv2.SimpleBlobDetector_create(params) + keypoints = detector.detect(gray) + + if len(keypoints) == 0: + return [], [], None + + # If we have a location hint, filter blobs by spatial proximity + # to avoid picking the top-left screw instead of the white-block screw. + best_kp = None + best_score = -1.0 + + # Define center points based on image quadrants + center_x, center_y = w // 2, h // 2 + if location_hint: + loc = location_hint.lower() + + for kp in keypoints: + kx, ky = kp.pt[0], kp.pt[1] + + # Score = Area + Proximity to hint + score = kp.size # bigger is better + + # Bonus logic based on location_hint + if location_hint: + if "block" in loc or "ceramic" in loc or "bottom" in loc: + if ky > center_y: # Lower half of image + score += 100 + if "right" in loc and kx > center_x: + score += 50 + if "center" in loc: + score += 50 - abs(kx - center_x) / 5 + elif "top" in loc or "left" in loc: + if ky < center_y: # Upper half + score += 100 + + if score > best_score: + best_score = score + best_kp = kp + + if best_kp is None: + return [], [], None + + # Convert to single point array to reuse downstream logic + pts = np.array([[best_kp.pt[0], best_kp.pt[1]]]) + valid = pts # Skip array clustering + + else: + raise ValueError(f"Unknown feature_type: {feature_type}") + + # ===================================================================== + # Common: Build bbox from valid points + # ===================================================================== + xs = valid[:, 0].astype(int) + ys = valid[:, 1].astype(int) + + x1 = max(0, int(xs.min()) - pad_x) + y1 = max(0, int(ys.min()) - pad_y) + x2 = min(w, int(xs.max()) + pad_x) + y2 = min(h, int(ys.max()) + pad_y) + + bbox = [x1, y1, x2, y2] + + # --- CONDITIONAL CROP: Only if entity is a 'leg' --- + if 'leg' in entity_name.lower(): + # Use the ACTUAL detected Hough line points (ignore the heavy padding) + # feat_x_min = int(valid[:, 0].min()) + # feat_x_max = int(valid[:, 0].max()) + # feat_y_min = int(valid[:, 1].min()) + # feat_y_max = int(valid[:, 1].max()) + + # detected_h = feat_y_max - feat_y_min + + # # Heuristic: if the vertical span is very small, we only caught the body bottom edge. + # # Extend downward by an estimated leg length (~2.5ร— the detected edge thickness, + # # minimum 60 px for a TO-220 leg), but NEVER force it to the image border. + # if detected_h < (h * 0.10): + # leg_ext = max(int(detected_h * 2.5), 60) + # feat_y_max = min(h, feat_y_max + leg_ext) + # print(f"[FeatureDetector] Short detection ({detected_h}px), extending legs by {leg_ext}px") + # else: + # print(f"[FeatureDetector] Good leg detection span ({detected_h}px), keeping as-is") + + # # Tight padding around the real features (10 % of feature size, min 4 px) + # pad_x = max(4, int((feat_x_max - feat_x_min) * 0.10)) + # pad_y = max(4, int((feat_y_max - feat_y_min) * 0.10)) + + # new_x1 = max(0, feat_x_min - pad_x) + # new_x2 = min(w, feat_x_max + pad_x) + # new_y1 = max(0, feat_y_min - pad_y) + # new_y2 = min(h, feat_y_max + pad_y) + + # bbox = [new_x1, new_y1, new_x2, new_y2] + # print(f"[FeatureDetector] Final leg bbox: {bbox}") + + # If HoughLinesP only caught the plastic body (raw_y2 < 85% of image height), + # we MUST override it to cover the actual leg area. + + if y2 < (h * 0.85): + # Force detection to extend to the bottom edge of the image + # effective_bottom = h + # # Use the full height to calculate the bottom 35% + # effective_height = effective_bottom - y1 + # new_y1 = effective_bottom - int(effective_height * 0.35) + # new_y2 = effective_bottom + # new_y1 = y2 + # new_y2 = h + # Calculate the height of the detected box + box_height = y2 - y1 + + # Target the bottom 35% of the DETECTED feature box, anchored at y2 + new_y1 = y2 - int(box_height * 0.35) + new_y2 = y2 + print("[FeatureDetector] HoughLinesP might caught the plastic body (raw_y2 < 85% of image height)") + else: + # Standard case: detection caught the whole component. + # Use the bottom 35% of the detected bbox. + bbox_height = y2 - y1 + new_y1 = y2 - int(bbox_height * 0.35) + new_y2 = y2 + print("[FeatureDetector] using standard case") + + new_y1 = max(y1, new_y1) # Safety clamp to prevent negative height + bbox = [x1, new_y1, x2, new_y2] # Overwrite bbox to only point to the legs! + # --------------------------------------------------- + + bbox_area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) + area_ratio = bbox_area / (h * w) + + if area_ratio > max_area_ratio: + print(f"[FeatureDetector] Bbox too large ({area_ratio:.2%}), rejecting.") + return [], [], None + + mask = np.zeros((h, w), dtype=np.uint8) + cv2.rectangle(mask, (bbox[0], bbox[1]), (bbox[2], bbox[3]), 255, -1) + + pred = { + 'entity': entity_name, + 'subentity': entity_name, + 'pred_box': torch.tensor([x1, y1, x2, y2]).float(), + 'bbox': bbox, + 'pred_class': torch.tensor(0).long(), + 'pred_mask': torch.from_numpy(mask).bool(), + 'mask': mask, + 'score': torch.tensor(1.0).float(), + 'area_ratio': area_ratio + } + + viz = image.copy() + cv2.rectangle(viz, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(viz, entity_name, (x1, max(10, y1 - 5)), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + + return [pred], [pred], viz \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/instance_processor.py b/ArtiAgent - DefectFill/src/pipeline/instance_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..e24d90df3f1f8a68308e4dc2835e2c6a8b263ff0 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/instance_processor.py @@ -0,0 +1,101 @@ +""" +Instance Processor โ€” DefectDiffu Edition + +Stripped-down version: DefectDiffu does NOT use 16x16 patch mappings, +so all patch-based artifact logic has been removed. + +Retained utilities: + - bbox โ†” mask helpers (for verification cropping) + - IoU calculation + - Visualization helpers +""" + +import random +import torch +import numpy as np +from PIL import Image +from typing import List, Dict, Tuple, Optional, Union +import matplotlib.pyplot as plt +import matplotlib.patches as patches + + +class InstanceProcessor: + """Utility class for detection post-processing and mask operations.""" + + @staticmethod + def calculate_iou(box1: Union[List, np.ndarray], box2: Union[List, np.ndarray]) -> float: + x1 = max(box1[0], box2[0]) + y1 = max(box1[1], box2[1]) + x2 = min(box1[2], box2[2]) + y2 = min(box1[3], box2[3]) + if x2 <= x1 or y2 <= y1: + return 0.0 + intersection = (x2 - x1) * (y2 - y1) + area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) + area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) + union = area1 + area2 - intersection + return intersection / union if union > 0 else 0.0 + + @staticmethod + def mask_from_bbox(bbox: Tuple[int, int, int, int], img_shape: Tuple[int, ...]) -> np.ndarray: + """Create a binary mask from a bounding box.""" + h, w = img_shape[:2] + mask = np.zeros((h, w), dtype=np.uint8) + x1, y1, x2, y2 = bbox + x1, y1 = max(0, x1), max(0, y1) + x2, y2 = min(w, x2), min(h, y2) + if x2 > x1 and y2 > y1: + mask[y1:y2, x1:x2] = 1 + return mask + + @staticmethod + def get_bbox_from_mask(mask: np.ndarray, margin: int = 0) -> Tuple[int, int, int, int]: + """Compute tight bounding box from binary mask, with optional margin.""" + ys, xs = np.where(mask > 0) + if len(ys) == 0: + return (0, 0, 0, 0) + y1, y2 = ys.min(), ys.max() + x1, x2 = xs.min(), xs.max() + h, w = mask.shape + x1 = max(0, x1 - margin) + y1 = max(0, y1 - margin) + x2 = min(w, x2 + margin) + y2 = min(h, y2 + margin) + return (x1, y1, x2, y2) + + @staticmethod + def visualize_generation_result( + original_image: np.ndarray, + generated_image: np.ndarray, + defect_mask: np.ndarray, + output_path: str, + title: str = "DefectDiffu Generation Result" + ): + """Create a 3-panel visualization: original, generated, mask overlay.""" + fig, axes = plt.subplots(1, 3, figsize=(18, 6)) + + axes[0].imshow(original_image) + axes[0].set_title("Original (Planning Reference)") + axes[0].axis("off") + + axes[1].imshow(generated_image) + axes[1].set_title("Generated Defect Image") + axes[1].axis("off") + + axes[2].imshow(generated_image) + axes[2].imshow(defect_mask, alpha=0.5, cmap="Reds") + axes[2].set_title("Defect Mask Overlay") + axes[2].axis("off") + + fig.suptitle(title, fontsize=14) + plt.tight_layout() + plt.savefig(output_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"[Viz] Saved result visualization to {output_path}") + + @staticmethod + def resize_to_square(image: np.ndarray, size: int = 512) -> np.ndarray: + """Resize image to square (DefectDiffu expects 512x512).""" + pil_img = Image.fromarray(image) if isinstance(image, np.ndarray) else image + pil_img = pil_img.resize((size, size), Image.LANCZOS) + return np.array(pil_img) diff --git a/ArtiAgent - DefectFill/src/pipeline/local_vlm_client.py b/ArtiAgent - DefectFill/src/pipeline/local_vlm_client.py new file mode 100644 index 0000000000000000000000000000000000000000..e20c4aecc9a6432b98b31c84e142e60d420ed079 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/local_vlm_client.py @@ -0,0 +1,192 @@ +import base64 +import io +import json +import requests +from typing import List, Optional, Union +from PIL import Image +import numpy as np +import re + + +class Usage: + def __init__(self, input_tokens=0, output_tokens=0): + self.input_tokens = input_tokens + self.output_tokens = output_tokens + + +class ChatCompletion: + def __init__(self, output_parsed, usage=None): + self.output_parsed = output_parsed + self.usage = usage or Usage() + + +class LocalVLMClient: + def __init__(self, base_url="http://localhost:11434", model="gemma3:12b"): + self.base_url = base_url + self.model = model + self.api_url = f"{base_url}/api/chat" + # Magic alias: Allows client.responses.parse(...) calls in prompts.py to work directly + self.responses = self + + def _strip_markdown_fences(self, text: str) -> str: + """Strip ```json ... ``` markdown code fences from VLM output.""" + cleaned = text.strip() + cleaned = re.sub(r'^```(?:json)?\s*', '', cleaned, flags=re.IGNORECASE) + cleaned = re.sub(r'\s*```\s*$', '', cleaned) + return cleaned.strip() + + def parse(self, model: str = None, input: list = None, temperature: float = 0.2, text_format=None): + """Alias for responses_parse to support client.responses.parse(...)""" + return self.responses_parse(model=model, input=input, temperature=temperature, text_format=text_format) + + def _pil_to_base64(self, img: Image.Image) -> str: + """Helper to convert a PIL Image into a base64 string for Ollama.""" + buffered = io.BytesIO() + img.convert("RGB").save(buffered, format="JPEG") + return base64.b64encode(buffered.getvalue()).decode("utf-8") + + def _prepare_image_base64(self, image_input) -> str: + """Converts any image input type into a pure base64 string.""" + if isinstance(image_input, np.ndarray): + img = Image.fromarray(image_input) + return self._pil_to_base64(img) + + elif isinstance(image_input, Image.Image): + return self._pil_to_base64(image_input) + + elif isinstance(image_input, str): + if image_input.startswith("data:image"): + return image_input.split(",")[1] + elif image_input.startswith("http"): + import urllib.request + with urllib.request.urlopen(image_input) as response: + img = Image.open(io.BytesIO(response.read())) + return self._pil_to_base64(img) + else: + # Local file path + with Image.open(image_input) as img: + return self._pil_to_base64(img) + else: + raise ValueError(f"Unsupported image input type: {type(image_input)}") + + def _call_ollama_chat(self, messages: list, temperature: float = 0.2, format_json: bool = False): + """Calls Ollama Chat API with system/user messages and images.""" + payload = { + "model": self.model, + "messages": messages, + "stream": False, + "options": { + "temperature": temperature + } + } + + if format_json: + payload["format"] = "json" + + try: + response = requests.post(self.api_url, json=payload, timeout=600) + response.raise_for_status() + result = response.json() + return result["message"]["content"] + except Exception as e: + print(f"Ollama API Error: {e}") + raise + + def responses_parse(self, model: str = None, input: list = None, temperature: float = 0.2, text_format=None): + # Override gpt-4o / external model strings with local Gemma model + active_model = self.model + formatted_messages = [] + + for msg in input: + role = msg["role"] + content = msg["content"] + + if isinstance(content, str): + formatted_messages.append({"role": role, "content": content}) + elif isinstance(content, list): + text_parts = [] + b64_images = [] + + for item in content: + if item.get("type") == "input_text" or "text" in item: + text_parts.append(item.get("text", "")) + elif item.get("type") == "input_image" or "image_url" in item: + img_src = item.get("image_url", item.get("image")) + b64_images.append(self._prepare_image_base64(img_src)) + + msg_obj = { + "role": role, + "content": "\n".join(text_parts) + } + if b64_images: + msg_obj["images"] = b64_images + + formatted_messages.append(msg_obj) + + # FIX 1: Provide explicit JSON structure examples instead of schema definitions + if text_format: + instruction = ( + "\n\nIMPORTANT: Do NOT output the schema structure or field types. " + "Output ONLY a populated JSON object like this:\n" + "{\n" + ' "has_artifact": true,\n' + ' "explanation": "Visual description of what was detected",\n' + ' "label": "artifact name"\n' + "}" + ) + if formatted_messages and formatted_messages[-1]["role"] == "user": + formatted_messages[-1]["content"] += instruction + else: + formatted_messages.append({"role": "user", "content": instruction}) + + payload = { + "model": active_model, + "messages": formatted_messages, + "stream": False, + "options": {"temperature": temperature} + } + + # Make the request to Ollama + try: + response = requests.post(self.api_url, json=payload, timeout=600) + response.raise_for_status() + response_text = response.json()["message"]["content"] + except Exception as e: + print(f"Ollama API error: {e}") + raise + + # Determine the parsed output based on whether a Pydantic text_format was passed + if text_format: + try: + parsed_json = json.loads(self._strip_markdown_fences(response_text)) + parsed_output = text_format(**parsed_json) + except Exception as e: + print(f"Failed to parse JSON response: {e}\nRaw output: {response_text}") + raise + else: + # Fallback for plain text standard responses + parsed_output = type("Output", (), {"explanation": response_text})() + + # Return a single standardized response object + return type("Response", (), { + "output_parsed": parsed_output, + "usage": Usage() + })() + + +# Quick Test Example +if __name__ == "__main__": + client = LocalVLMClient(model="gemma3:12b") + + test_input = [ + {"role": "system", "content": "You are a visual AI inspector."}, + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Analyze if this image is distorted."} + ] + } + ] + + res = client.responses_parse(input=test_input) + print("Result:", res.output_parsed.explanation) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/prompts.py b/ArtiAgent - DefectFill/src/pipeline/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..c79d8f0ffb7f852725089fc5876bc8206b46f823 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/prompts.py @@ -0,0 +1,653 @@ +""" +Prompts and VLM interaction layer โ€” DefectFill Edition (with VLM list selection) + +Key changes: + - VLM selects object_class from a given list based on product description. + - If defect_type is given, VLM selects the closest match from the valid list. + - If defect_type is NOT given, VLM picks the most plausible one from the list. + - All selections are constrained to the provided valid lists. +""" + +import json +from PIL import Image +import base64 +import io +import numpy as np +import re +from pipeline.local_vlm_client import LocalVLMClient, ChatCompletion +from typing import Union, List, Optional, Dict +from pydantic import BaseModel, Field +import os + +DEFAULT_MODEL = "gemma3:12b" +default_client = LocalVLMClient(model=DEFAULT_MODEL) +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) + + +# ============================================================================= +# Pydantic schemas +# ============================================================================= + +class DefectPlan(BaseModel): + """Single defect plan for DefectFill.""" + defect_type: str = Field( + description="Defect type / checkpoint folder name (e.g., xray_die, bubble, scratch)" + ) + object_class: str = Field( + description="Object class for checkpoint routing (e.g., xray_PCB, vcsel)" + ) + target_entity: str = Field(description="Main product entity/component for placement") + target_subentity: Optional[str] = Field(None, description="Specific sub-part for placement") + location_hint: str = Field(description="Spatial location for defect placement") + defect_coverage_ratio: float = Field( + default=0.5, + ge=0.10, + le=1.0, + description="Coverage ratio of defect relative to target ROI" + ) + mask_shape: str = Field( + default="free", + description="Boundary geometry: 'circle', 'square', 'rectangle', or 'free'" + ) + description: str = Field(description="Detailed visual description for verification") + + +class ProductDefectPlan(BaseModel): + product_type: str = Field(description="Identified product category / object_class") + analysis: str = Field(description="Visual inspection analysis") + possible_defects: List[DefectPlan] = Field(description="List of proposed defects") + + +class BboxResponse(BaseModel): + bboxes: List[List[int]] + + +class ArtifactDescriptionResponse(BaseModel): + has_artifact: bool + explanation: str + label: str + + +class ArtifactExplanationResponse(BaseModel): + explanation: str + + +class ArtifactSuccessResponse(BaseModel): + reasoning: str + success: bool + + +class VocabResponse(BaseModel): + peripheral: Optional[Dict[str, List[str]]] = None + intermediate: Optional[Dict[str, List[str]]] = None + + +# ============================================================================= +# Image encoding utility +# ============================================================================= + +def encode_image_to_base64(image): + try: + if isinstance(image, str): + if not os.path.exists(image): + print(f"Warning: File path '{image}' not found. Skipping...") + return "" + pil_image = Image.open(image) + elif isinstance(image, np.ndarray): + pil_image = Image.fromarray(image) + else: + pil_image = image + + if pil_image.mode != 'RGB': + pil_image = pil_image.convert('RGB') + + buffer = io.BytesIO() + pil_image.save(buffer, format='JPEG') + buffer.seek(0) + return base64.b64encode(buffer.getvalue()).decode('utf-8') + except Exception as e: + print(f"Warning: Could not encode image ({e}). Skipping...") + return "" + + +def clean_json_string(raw_string: str) -> str: + if not isinstance(raw_string, str): + return raw_string + cleaned = raw_string.strip() + cleaned = re.sub(r'^```(?:json)?\s*', '', cleaned, flags=re.IGNORECASE) + cleaned = re.sub(r'\s*```$', '', cleaned) + return cleaned.strip() + + +# ============================================================================= +# Step 1: Defect Planning (DefectFill with list-constrained selection) +# ============================================================================= + +def plan_defects_for_product( + client, + product_description: str, + image, + money_manager=None, + target_defect_type: Optional[str] = None, + object_class: Optional[str] = None, + valid_object_classes: Optional[List[str]] = None, + valid_defect_types: Optional[Dict[str, List[str]]] = None, + num_defects: int = 3, + domain_hint: Optional[str] = None, +): + """ + Agent reasons about possible defects and outputs DefectFill-native plans. + The VLM selects object_class from valid_object_classes and defect_type from + the corresponding valid_defect_types list. + """ + if client is None: + client = default_client + base64_image = encode_image_to_base64(image) + + # ------------------------------------------------------------------ + # Build list-selection instructions for the VLM + # ------------------------------------------------------------------ + list_instructions = "" + + if valid_object_classes: + list_instructions += f"\nAVAILABLE OBJECT CLASSES: {json.dumps(valid_object_classes)}\n" + if not object_class: + list_instructions += ( + "Since NO object_class was provided by the user, you MUST select the " + "best-matching object_class from the AVAILABLE OBJECT CLASSES list above, " + "based on the product description and the reference image.\n" + ) + else: + list_instructions += ( + f"The user provided object_class: '{object_class}'. " + f"Use this exact value if it appears in the list; otherwise pick the closest match.\n" + ) + + if valid_defect_types: + if object_class and object_class in valid_defect_types: + list_instructions += ( + f"\nAVAILABLE DEFECT TYPES for '{object_class}': " + f"{json.dumps(valid_defect_types[object_class])}\n" + ) + elif not object_class and valid_object_classes: + list_instructions += "\nAVAILABLE DEFECT TYPES per object class:\n" + for oc, dtypes in valid_defect_types.items(): + list_instructions += f" {oc}: {json.dumps(dtypes)}\n" + + if not target_defect_type: + list_instructions += ( + "Since NO defect_type was provided by the user, you MUST select the " + "most plausible defect_type from the AVAILABLE DEFECT TYPES list for the " + "chosen object_class.\n" + ) + else: + list_instructions += ( + f"The user requested defect_type: '{target_defect_type}'. " + f"Select the closest match from the AVAILABLE DEFECT TYPES list.\n" + ) + + if valid_object_classes or valid_defect_types: + list_instructions += ( + "\nCRITICAL CONSTRAINT: You MUST ONLY use object_class and defect_type " + "values that exist in the provided lists above. Do NOT invent new object classes " + "or defect types. If the user's request does not match any available option, " + "pick the closest valid alternative and note it in the description.\n" + ) + + # ------------------------------------------------------------------ + # System prompt + # ------------------------------------------------------------------ + system_prompt = f""" +You are an expert manufacturing quality control engineer and defect analysis agent. + + +You are given: +1. A product description: "{product_description}" +2. A clean reference image of the product +{list_instructions} + + +DOMAIN-SPECIFIC INSTRUCTIONS: +If the image is an X-ray or CT scan of a PCB/IC package: +- The die appears as a large dark or light rectangular region, usually in the center. +- Solder balls/BGA appear as an array of circular dots. +- Defects in X-ray appear as grayscale anomalies (darker or brighter regions), NOT as circuit traces, lines, or wire patterns. +- A "xray_die" defect should look like: a natural void, crack, or delamination INSIDE the die material โ€” blobby, irregular, or circular grayscale spots. +- A "xray_solder" defect should look like: missing solder, solder bridges, or voids in solder balls. +- NEVER describe the defect as "circuit traces", "wiring", "grid lines", or "PCB patterns". The defect must look like a manufacturing flaw, not an overlaid schematic. +- The `defect_coverage_ratio` should be 1.0. + +If the image is a standard TO-220 package (like a Triac): +- The component consists of three metallic pins/leads, a black plastic mold body, a metal mounting tab with a mounting hole, and printed silkscreen text. +- **CRITICAL CONSTRAINT**: You MUST exclusively propose defects located on the **metallic pins/leads**. Do NOT propose any defects on the black plastic body, the metal mounting tab, the mounting hole, or the silkscreen text. +- Acceptable pin/lead defects include: pin bending (crooked leads), lead deformation, surface scratches on the metallic pins, or contamination/solder dross on the lead surfaces. +- Double means 2 pin will bend out +- Unsorted means 1 pin will bend out +- The target entity and target_subentity should be 3 pin rather than specific one pin only. +- The `description` must clearly detail exactly which pin is affected and how it is deformed or damaged (e.g., "The right pin is bent sharply outwards near the mid-section"). +- The `defect_coverage_ratio` should be 0.4 - 0.60. + + +{domain_hint} + +Your task is to analyze the product and propose {num_defects} realistic manufacturing or handling defects. +Each defect_type must be selected at least ONCE but MUST less than {num_defects}. + +For EACH defect, you must produce the following fields: + - object_class: The object category for checkpoint routing (MUST be from the available list). + - defect_type: The checkpoint folder name (MUST be from the available list for the chosen object_class). + - target_entity: The main component where the defect should be placed. + - target_subentity: A specific sub-part if applicable.bbox_w = x2 - x1 + - location_hint: Spatial location on the image. + - defect_coverage_ratio: Compulsory is 1.0 unless specify by DOMAIN-SPECIFIC INSTRUCTIONS + - mask_shape: Compulsory is "rectangle" + - description: Detailed visual description of what the defect looks like. + +Think step by step, then output EXACTLY one JSON object with this structure: +{{ + "product_type": "short product name / object_class", + "analysis": "one sentence summary of plausible defects", + "possible_defects": [ + {{ + "object_class": "xray_PCB", + "defect_type": "xray_die", + "target_entity": "die surface", + "target_subentity": "bond pad area", + "location_hint": "center of die", + "defect_coverage_ratio": 1.0, + "mask_shape": "rectangle", + "description": "a small dark void under the die surface near the bond pad" + }} + ] +}} + + +IMPORTANT RULES: +1. object_class MUST be selected from the available list. Do not invent new classes. +2. defect_type MUST be selected from the available list for the chosen object_class. Do not invent new types. +3. target_entity should be a concrete, visually distinctive component name. +4. Do NOT target thin wire-like structures as primary target_entity. +5. Return ONLY the keys shown above. Do NOT return has_artifact, explanation, label, c_p, c_d, c_f, w_d, or w_p. +""" + + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{base64_image}"}, + {"type": "input_text", "text": f"Product: {product_description}\nAnalyze and propose realistic defects with placement plans."} + ] + } + ], + temperature=0.3, + text_format=ProductDefectPlan + ) + + if money_manager: + money_manager(response) + + if hasattr(response, "output_parsed") and response.output_parsed: + plan = response.output_parsed + # Validate selections against valid lists + if valid_object_classes and plan.possible_defects: + for d in plan.possible_defects: + if d.object_class not in valid_object_classes: + print(f"[Plan] Warning: VLM selected object_class '{d.object_class}' not in valid list. " + f"Falling back to first valid: {valid_object_classes[0]}") + d.object_class = valid_object_classes[0] + if valid_defect_types and d.object_class in valid_defect_types: + if d.defect_type not in valid_defect_types[d.object_class]: + print(f"[Plan] Warning: VLM selected defect_type '{d.defect_type}' not valid for '{d.object_class}'. " + f"Falling back to first valid: {valid_defect_types[d.object_class][0]}") + d.defect_type = valid_defect_types[d.object_class][0] + return plan + + # Fallback parsing for raw string output + raw_content = getattr(response, "content", str(response)) + cleaned_text = clean_json_string(raw_content) + parsed = json.loads(cleaned_text) + + if isinstance(parsed, list): + defects = [] + for item in parsed: + oc = item.get("object_class", object_class or (valid_object_classes[0] if valid_object_classes else "")) + dt = item.get("defect_type", target_defect_type or "") + # Validate + if valid_object_classes and oc not in valid_object_classes: + oc = valid_object_classes[0] + if valid_defect_types and oc in valid_defect_types and dt not in valid_defect_types[oc]: + dt = valid_defect_types[oc][0] if valid_defect_types[oc] else dt + defects.append({ + "object_class": oc, + "defect_type": dt, + "description": item.get("description", item.get("explanation", "")), + "target_entity": item.get("target_entity", item.get("entity", "")), + "target_subentity": item.get("target_subentity", item.get("subentity")), + "location_hint": item.get("location_hint", ""), + "defect_coverage_ratio": item.get("defect_coverage_ratio", 0.25), + "mask_shape": item.get("mask_shape", "free"), + }) + plan_dict = { + "product_type": object_class or (valid_object_classes[0] if valid_object_classes else "component"), + "analysis": f"Planned {len(defects)} defects", + "possible_defects": defects + } + return ProductDefectPlan(**plan_dict) + + elif isinstance(parsed, dict): + if "possible_defects" not in parsed: + for key in ["defects", "defect_plan", "results"]: + if key in parsed: + parsed["possible_defects"] = parsed.pop(key) + break + if "possible_defects" in parsed: + for d in parsed["possible_defects"]: + if "explanation" in d and "description" not in d: + d["description"] = d.pop("explanation") + if "entity" in d and "target_entity" not in d: + d["target_entity"] = d.pop("entity") + d.setdefault("object_class", object_class or (valid_object_classes[0] if valid_object_classes else "")) + d.setdefault("defect_type", target_defect_type or "") + d.setdefault("defect_coverage_ratio", 0.25) + d.setdefault("mask_shape", "free") + # Validate + oc = d.get("object_class", "") + dt = d.get("defect_type", "") + if valid_object_classes and oc not in valid_object_classes: + d["object_class"] = valid_object_classes[0] + if valid_defect_types and d["object_class"] in valid_defect_types and dt not in valid_defect_types[d["object_class"]]: + d["defect_type"] = valid_defect_types[d["object_class"]][0] if valid_defect_types[d["object_class"]] else dt + return ProductDefectPlan(**parsed) + + except Exception as e: + print(f"Error in defect planning: {e}") + return None + + +# ============================================================================= +# Entity vocabulary & bbox detection (unchanged) +# ============================================================================= + +def get_entity_subentities(client, image, money_manager=None): + if client is None: + client = default_client + base64_image = encode_image_to_base64(image) + system_prompt = """ +You are given a microscopic or optical inspection image of a precision component. +Identify visible entities and subentities, split into peripheral and intermediate layers. + +Output exactly one JSON object with two keys: "peripheral" and "intermediate". +Each value is a dict mapping entity names to lists of subentity names. + +Hard rules: +1) Each entity MUST have at least one subentity. +2) Subentities must be clearly visible and segmentable. +3) Use concise, lowercase nouns. +4) Do not invent occluded parts. +""" + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": [{"type": "input_image", "image_url": f"data:image/jpeg;base64,{base64_image}"}]} + ], + temperature=0.2, + text_format=VocabResponse + ) + if money_manager: + money_manager(response) + if hasattr(response, "output_parsed") and response.output_parsed: + return response.output_parsed + raw_content = getattr(response, "content", str(response)) + cleaned_text = clean_json_string(raw_content) + return VocabResponse.model_validate_json(cleaned_text) + except Exception as e: + print(f"Error in entity analysis: {e}") + return None + + +def get_entity_bboxes(client, image, entity_name, subentity=None, location_hint=None, money_manager=None): + if client is None: + client = default_client + base64_image = encode_image_to_base64(image) + + # --- DYNAMIC GROUPING INSTRUCTION --- + group_instruction = "" + # If the entity implies a group/array (triac pins, IC leads, solder ball arrays) + if any(k in entity_name.lower() for k in ['pin', 'leg', 'lead', 'array', 'grid', 'group', 'bank']): + group_instruction = ( + f"CRITICAL RULE: The entity '{entity_name}' represents a SET or GROUP of elements. " + "Do NOT return individual boxes for each element. Return exactly ONE bounding box " + "that tightly encloses the ENTIRE group. " + # --- NEW INSTRUCTION --- + "**CRITICAL CONSTRAINT**: If the entity is '3 pin' or 'leg', your bounding box MUST EXCLUDE the black plastic body and metal tab. " + "The bounding box MUST ONLY enclose the vertical metal leads/pins at the BOTTOM of the component. " + "Do NOT include the plastic body or the top half of the component." + ) + # ------------------------------------- + + specific_instruction = "" + if subentity: + specific_instruction += f"\nSpecifically find the '{subentity}' instance of this entity." + if location_hint: + specific_instruction += f"\nThe target is located at: {location_hint}. Look specifically in this area." + + system_prompt = f""" +You are an industrial inspection bounding box detector. +Find ALL instances of: {entity_name} +{specific_instruction} +{group_instruction} +Return JSON with "bboxes" key: [[x1, y1, x2, y2], ...] in ABSOLUTE PIXEL coordinates. +If none found, return {{"bboxes": []}}. +""" + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": [{"type": "input_image", "image_url": f"data:image/jpeg;base64,{base64_image}"}]} + ], + temperature=0.2, + text_format=BboxResponse + ) + if money_manager: + money_manager(response) + if hasattr(response, "output_parsed") and response.output_parsed: + return response.output_parsed.bboxes + raw_content = getattr(response, "content", str(response)) + cleaned = clean_json_string(raw_content) + parsed = json.loads(cleaned) + return parsed.get("bboxes", []) if isinstance(parsed, dict) else [] + except Exception as e: + print(f"Error getting bboxes for {entity_name}: {e}") + return [] + + +# ============================================================================= +# Verification prompts (unchanged) +# ============================================================================= + +def artifact_description(client, masked_original_image, target_original_image, + target_artifact_image, object_name, artifact_type, money_manager=None): + if client is None: + client = default_client + + original_masked = encode_image_to_base64(masked_original_image) + original_target = encode_image_to_base64(target_original_image) + artifact_target = encode_image_to_base64(target_artifact_image) + + combined_instruction = { + "addition": ( + "You are an expert industrial quality control inspector.\n\n" + "ANOMALY DEFINITION (ADDITION): extra material, dust, debris, contamination, " + "solder blob, particulate, smudge, residue, bubble, scratch.\n\n" + "Checklist: 1) Is there extra material in the target region? " + "2) Does it have a distinct boundary? 3) Is it consistent with the defect class?" + ), + "removal": ( + "You are an expert industrial quality control inspector.\n\n" + "ANOMALY DEFINITION (REMOVAL): missing part, lifted pad, absent component, " + "incomplete geometric path, broken lead.\n\n" + "Checklist: 1) Is there a clear gap or missing structure? " + "2) Is the absence NOT explained by occlusion or viewpoint?" + ), + "distortion": ( + "You are an expert industrial quality control inspector.\n\n" + "ANOMALY DEFINITION (DISTORTION): bent, warped, misaligned, cracked, " + "deformed, wavy, stress pattern, bowing.\n\n" + "Checklist: 1) Are straight lines or edges warped? " + "2) Is geometric symmetry broken?" + ), + "fusion": ( + "You are an expert industrial quality control inspector.\n\n" + "ANOMALY DEFINITION (FUSION/BRIDGING): solder bridge, shorted pins, " + "merged adjacent structures.\n\n" + "Checklist: 1) Are separate conductive paths unnaturally merged? " + "2) Is the boundary between components degraded?" + ) + } + + prompt = f""" +{combined_instruction.get(artifact_type, combined_instruction["addition"])} + +Return JSON: + "has_artifact": true/false + "explanation": "Detailed description of what looks wrong" + "label": "Brief description (empty if no artifact)" + +Rules: Focus on visible evidence. Do not refer to images by number. + +CRITICAL VERIFICATION RULE: +Before answering, compare the second image (original target) and third image (artifact target) pixel-by-pixel. +If they are visually IDENTICAL with no discernible difference in texture, color, or structure, you MUST return: + "has_artifact": false + "explanation": "No visible anomaly detected in the target region." + "label": "" +Only return has_artifact=true if you can point to a specific, concrete visual difference. +""" + + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": prompt}, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_masked}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_target}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{artifact_target}"}, + {"type": "input_text", "text": f"{object_name}"} + ] + } + ], + temperature=0.2, + text_format=ArtifactDescriptionResponse + ) + if money_manager: + money_manager(response) + return response.output_parsed + except Exception as e: + print(f"Error in artifact description ({artifact_type}): {e}") + return ArtifactDescriptionResponse(has_artifact=False, explanation="", label="") + + +# ============================================================================= +# Explanation generation (unchanged) +# ============================================================================= + +def artifact_explanation_from_triplet(client, masked_original_image, target_original_image, + target_artifact_image, object_name, artifact_type, money_manager=None): + if client is None: + client = default_client + + original_masked = encode_image_to_base64(masked_original_image) + original_target = encode_image_to_base64(target_original_image) + artifact_target = encode_image_to_base64(target_artifact_image) + + type_guidance = { + "addition": "Describe an ADDED instance: duplicated parts, extra elements, foreign material.", + "removal": "Describe a MISSING part: gaps, smoothed-over areas, discontinuity.", + "distortion": "Describe WARPING: bent shapes, irregular textures, malformed geometry.", + "fusion": "Describe MERGING: boundary loss, texture bleed, interpenetration." + } + + prompt = f""" +You receive three images and an object name: + 1) Original WITHOUT target region + 2) Original showing ONLY target region + 3) Artifact showing ONLY target region (describe THIS one) + 4) Object name: {object_name} + +TASK: One-sentence description of what looks wrong in Image 3, consistent with {artifact_type}. +{type_guidance.get(artifact_type, "")} + +Rules: visible evidence only, no JSON, no image numbers, simple language. +""" + + try: + response = client.responses.parse( + model=DEFAULT_MODEL, + input=[ + {"role": "system", "content": prompt}, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_masked}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{original_target}"}, + {"type": "input_image", "image_url": f"data:image/jpeg;base64,{artifact_target}"}, + {"type": "input_text", "text": f"{object_name}"} + ] + } + ], + temperature=0.2, + text_format=ArtifactExplanationResponse + ) + if money_manager: + money_manager(response) + return response.output_parsed + except Exception as e: + print(f"Error in explanation ({artifact_type}): {e}") + return ArtifactExplanationResponse(explanation="") + + +# ============================================================================= +# Cost tracking (unchanged) +# ============================================================================= + +class MoneyManager: + def __init__(self, model: str = DEFAULT_MODEL): + self.total_cost = 0.0 + self.model = model + cost_table = { + "gpt-4o": (2.5/1000, 10/1000), + "gpt-4o-mini": (0.15/1000, 0.6/1000), + "gemini-2.5-flash": (0.3/1000, 2.5/1000), + "gemini-2.5-pro": (1.25/1000, 10/1000), + } + self.input_cost, self.output_cost = cost_table.get(model, (0.0, 0.0)) + + def __call__(self, response=None): + if response is None or not hasattr(response, "usage"): + return + try: + if hasattr(response.usage, "input_tokens"): + inp = response.usage.input_tokens + out = response.usage.output_tokens + elif hasattr(response, "usage_metadata"): + inp = response.usage_metadata.prompt_token_count + out = (response.usage_metadata.candidates_token_count + + getattr(response.usage_metadata, "thoughts_token_count", 0)) + else: + return + self.total_cost += (inp / 1000 * self.input_cost + out / 1000 * self.output_cost) + except Exception: + pass + + def refresh(self): + self.total_cost = 0.0 \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/pipeline/sample_manifest.csv b/ArtiAgent - DefectFill/src/pipeline/sample_manifest.csv new file mode 100644 index 0000000000000000000000000000000000000000..0aceaad08e7ac89e602f28759013e245288e2443 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/sample_manifest.csv @@ -0,0 +1,8 @@ +image_path,product_description +vcsel/clean_01.png,VCSEL laser diode with emission aperture and surrounding mesa structure +vcsel/clean_02.png,VCSEL laser diode with emission aperture and surrounding mesa structure +lens/clean_01.png,Optical lens with anti-reflective coating and mounting frame +lens/clean_02.png,Optical lens with anti-reflective coating and mounting frame +die/clean_01.png,Semiconductor die with bond pads and scribe lines +photodiode/clean_01.png,Photodiode sensor with photosensitive area and electrode contacts +optical_sensor/clean_01.png,Optical sensor with active area and reflective cavity diff --git a/ArtiAgent - DefectFill/src/pipeline/test_prompts.py b/ArtiAgent - DefectFill/src/pipeline/test_prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..24ce945da3961923b844377905b62f718987cd71 --- /dev/null +++ b/ArtiAgent - DefectFill/src/pipeline/test_prompts.py @@ -0,0 +1,13 @@ +from prompts import clean_json_string, VocabResponse + +# Test 1: clean_json_string +test = '```json\n{"peripheral": {"pin": ["lead"]}}\n```' +result = clean_json_string(test) +print("Test 1 - clean_json_string:") +print(repr(result)) +print() + +# Test 2: VocabResponse validation +v = VocabResponse.model_validate_json(result) +print("Test 2 - VocabResponse validation:") +print(v.peripheral) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/scripts/clean_up_stale_records.py b/ArtiAgent - DefectFill/src/scripts/clean_up_stale_records.py new file mode 100644 index 0000000000000000000000000000000000000000..f35843c426a84f808ffdc8483303785f00b62db8 --- /dev/null +++ b/ArtiAgent - DefectFill/src/scripts/clean_up_stale_records.py @@ -0,0 +1,33 @@ +import json +from pathlib import Path + +import sys +import os +# Add parent directory (src/) to path so 'pipeline' is findable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from pipeline.defect_rag import get_rag + +rag = get_rag() +data = rag.collection.get() + +stale_ids = [] + +for doc_id, metadata in zip(data['ids'], data['metadatas']): + # Parse stored patch paths + paths = json.loads(metadata['paths']) + + # Check key with .png extension + sample_path_str = paths.get('original_masked.png') or list(paths.values())[0] + sample_file = Path(sample_path_str) + + # If the file/folder no longer exists on disk, mark for deletion + if not sample_file.exists(): + print(f" โŒ Missing on disk: {doc_id} -> {sample_file}") + stale_ids.append(doc_id) + +if stale_ids: + rag.collection.delete(ids=stale_ids) + print(f"\n๐Ÿงน Cleaned up {len(stale_ids)} stale record(s) from database.") +else: + print("\nโœจ Database is clean! All entries match folders on disk.") \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/scripts/convert_roboflow_to_triplets.py b/ArtiAgent - DefectFill/src/scripts/convert_roboflow_to_triplets.py new file mode 100644 index 0000000000000000000000000000000000000000..f9cd12f154c4db48771f7e0a4f1121e9b9bbb0b0 --- /dev/null +++ b/ArtiAgent - DefectFill/src/scripts/convert_roboflow_to_triplets.py @@ -0,0 +1,102 @@ +# scripts/convert_roboflow_to_triplets.py +import json +import shutil +from pathlib import Path +from PIL import Image +import numpy as np + +def convert_roboflow_to_triplets(roboflow_dir: str, output_dir: str): + """ + Convert Roboflow COCO format to patch triplet format. + + Roboflow COCO structure: + train/ + _annotations.coco.json + image1.jpg + image2.jpg + """ + roboflow_path = Path(roboflow_dir) + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Load COCO annotations + anno_file = roboflow_path / "train" / "_annotations.coco.json" + with open(anno_file) as f: + coco = json.load(f) + + # Build image ID โ†’ filename map + images = {img['id']: img['file_name'] for img in coco['images']} + + # Group annotations by image + from collections import defaultdict + image_annotations = defaultdict(list) + for ann in coco['annotations']: + image_annotations[ann['image_id']].append(ann) + + triplet_id = 0 + for img_id, filename in images.items(): + img_path = roboflow_path / "train" / filename + if not img_path.exists(): + continue + + img = Image.open(img_path).convert('RGB') + img_array = np.array(img) + H, W = img_array.shape[:2] + + anns = image_annotations.get(img_id, []) + if not anns: + continue + + for ann in anns: + # 1. Skip non-defect category annotations (category_id == 2) + if ann.get('category_id') == 2: + continue + + # 2. Skip full-image bounding boxes + x, y, w, h = ann['bbox'] + if w >= W and h >= H: + continue + + triplet_id += 1 + triplet_dir = output_path / f"defect_{triplet_id:04d}" + triplet_dir.mkdir(exist_ok=True) + + # Get bbox + x, y, w, h = ann['bbox'] + x1, y1, x2, y2 = int(x), int(y), int(x+w), int(y+h) + + # Create mask from segmentation if available, else bbox + if 'segmentation' in ann and ann['segmentation']: + # COCO polygon segmentation + from pycocotools import mask as maskUtils + rles = maskUtils.frPyObjects(ann['segmentation'], H, W) + mask = maskUtils.decode(rles) + if len(mask.shape) == 3: + mask = np.any(mask, axis=2).astype(np.uint8) * 255 + else: + # Fallback: bbox mask + mask = np.zeros((H, W), dtype=np.uint8) + mask[y1:y2, x1:x2] = 255 + + # Save clean image (original_target) + img.save(triplet_dir / "original_target.png") + + # Save mask (original_masked) + Image.fromarray(mask).save(triplet_dir / "original_masked.png") + + # For artifact_target, we need the defective version. + # Since this is a real defect dataset, the original image IS the defect. + # For synthetic training, you may want to inpaint the defect out to create "clean", + # but for RAG retrieval, we can use the same image as artifact_target. + img.save(triplet_dir / "artifact_target.png") + + print(f"Created triplet {triplet_id}: {filename} โ†’ {triplet_dir}") + + print(f"\nTotal triplets created: {triplet_id}") + print(f"Output: {output_path}") + +if __name__ == "__main__": + convert_roboflow_to_triplets( + roboflow_dir="./data/external/Manufacturing_Defect_Detection", + output_dir="data/external/roboflow_manufacturing" + ) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/scripts/ingest_public_data.py b/ArtiAgent - DefectFill/src/scripts/ingest_public_data.py new file mode 100644 index 0000000000000000000000000000000000000000..a8a4afe4def500aa93833c20208dfba3d8735eab --- /dev/null +++ b/ArtiAgent - DefectFill/src/scripts/ingest_public_data.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# scripts/ingest_public_data.py +# ============================================ +# RUN ONCE: python scripts/ingest_public_data.py +# Populates ChromaDB with defect patch triplets +# ============================================ + +import chromadb +from sentence_transformers import SentenceTransformer +from PIL import Image +import json +from pathlib import Path +import argparse + + +# ============================================ +# DATASET REGISTRY โ€” Add new sources here +# ============================================ + +DATASETS = { + # โ”€โ”€ Commercial-friendly datasets โ”€โ”€ + "roboflow_manufacturing": { + "path": "data/external/roboflow_manufacturing", + "domain": "pcb", + "license": "cc-by-4.0", + "commercial_ok": True, + "attribution": "Roboflow Universe โ€” Manufacturing Defect Detection", + "description": "PCB manufacturing defects, CC BY 4.0 licensed" + }, + + # โ”€โ”€ Your proprietary data โ”€โ”€ + "my_vcsel_captures": { + "path": "data/my_captures/vcsel", + "domain": "glass/vcsel", + "license": "proprietary", + "commercial_ok": True, + "attribution": "Internal", + "description": "VCSEL laser diode defect captures" + }, + + # โ”€โ”€ Partner data (NDA) โ”€โ”€ + # "partner_pcb_nda": { + # "path": "data/partners/pcb_client_a", + # "domain": "pcb", + # "license": "partner_nda", + # "commercial_ok": True, + # "attribution": "Partner NDA", + # "description": "PCB client defect images under NDA" + # }, +} + + +def validate_triplet(folder: Path) -> dict | None: + """Check if folder contains required patch triplet files.""" + required = ["original_masked.png", "original_target.png", "artifact_target.png"] + paths = {f: folder / f for f in required} + + for f, p in paths.items(): + if not p.exists(): + print(f" โš ๏ธ Missing {f} in {folder}, skipping") + return None + + return {k: str(v) for k, v in paths.items()} + + +def generate_caption(folder_name: str, domain: str, dataset_desc: str) -> str: + """Generate a text caption for embedding.""" + # You can replace this with a VLM call for better captions + defect_type = folder_name.replace("_", " ") + return f"{defect_type} defect on {domain}. {dataset_desc}" + + +def ingest_dataset(collection, encoder, name: str, cfg: dict): + """Ingest one dataset into ChromaDB.""" + print(f"\n{'='*50}") + print(f"Dataset: {name}") + print(f"License: {cfg['license']} | Commercial: {'โœ…' if cfg['commercial_ok'] else 'โŒ'}") + print(f"Path: {cfg['path']}") + print(f"{'='*50}") + + base_path = Path(cfg["path"]) + if not base_path.exists(): + print(f" โš ๏ธ Path not found: {base_path}") + print(f" Create it and add patch triplet folders:") + print(f" {base_path}/defect_name_001/original_masked.png") + print(f" {base_path}/defect_name_001/original_target.png") + print(f" {base_path}/defect_name_001/artifact_target.png") + return 0 + + # NEW: Get set of IDs already in the collection to avoid re-processing + existing_ids = set(collection.get()["ids"]) + + count = 0 + for triplet_folder in sorted(base_path.iterdir()): + if not triplet_folder.is_dir(): + continue + + # NEW: Skip if already processed + doc_id = f"{name}_{triplet_folder.name}" + if doc_id in existing_ids: + print(f" โญ๏ธ Skipping {triplet_folder.name} (already in DB)") + continue + + paths = validate_triplet(triplet_folder) + if paths is None: + continue + + caption = generate_caption( + triplet_folder.name, + cfg["domain"], + cfg.get("description", "") + ) + + embedding = encoder.encode(caption) + + collection.upsert( + documents=[caption], + embeddings=[embedding.tolist()], + metadatas=[{ + "paths": json.dumps(paths), + "domain": cfg["domain"], + "license": cfg["license"], + "commercial_ok": cfg["commercial_ok"], + "source": name, + "attribution": cfg["attribution"], + "defect_name": triplet_folder.name + }], + ids=[f"{name}_{triplet_folder.name}"] + ) + + count += 1 + print(f" โœ… {triplet_folder.name}: {caption[:60]}...") + + return count + + +def main(): + parser = argparse.ArgumentParser(description="Ingest defect patch triplets into RAG DB") + parser.add_argument("--db-path", default="data/defect_db", help="ChromaDB persistent path") + parser.add_argument("--collection", default="defect_patches", help="Collection name") + parser.add_argument("--model", default="all-MiniLM-L6-v2", help="SentenceTransformer model") + args = parser.parse_args() + + # Initialize DB + Path(args.db_path).mkdir(parents=True, exist_ok=True) + client = chromadb.PersistentClient(path=args.db_path) + + # Delete existing collection if you want fresh start + # client.delete_collection(args.collection) + + collection = client.get_or_create_collection( + name=args.collection, + metadata={"hnsw:space": "cosine"} + ) + + print(f"DB path: {args.db_path}") + print(f"Collection: {args.collection}") + print(f"Existing entries: {collection.count()}") + + # Initialize encoder + print(f"\nLoading encoder: {args.model}") + encoder = SentenceTransformer(args.model) + + # Ingest all datasets + total = 0 + for name, cfg in DATASETS.items(): + # Skip non-commercial datasets in commercial builds + if not cfg.get("commercial_ok", False): + print(f"\nโญ๏ธ Skipping {name} โ€” not commercial-friendly") + continue + + count = ingest_dataset(collection, encoder, name, cfg) + total += count + + print(f"\n{'='*50}") + print(f"TOTAL INGESTED: {total} patch triplets") + print(f"TOTAL IN DB: {collection.count()}") + print(f"{'='*50}") + + # Print commercial summary + print("\n๐Ÿ“‹ Commercial License Summary:") + results = collection.get() + licenses = {} + for meta in results["metadatas"]: + lic = meta["license"] + licenses[lic] = licenses.get(lic, 0) + 1 + + for lic, count in licenses.items(): + icon = "โœ…" if any( + cfg["license"] == lic and cfg.get("commercial_ok") + for cfg in DATASETS.values() + ) else "โŒ" + print(f" {icon} {lic}: {count} entries") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/scripts/resize-file-for-training.py b/ArtiAgent - DefectFill/src/scripts/resize-file-for-training.py new file mode 100644 index 0000000000000000000000000000000000000000..ea277bc5e36e4e95de512d537bf28fbaa3127c79 --- /dev/null +++ b/ArtiAgent - DefectFill/src/scripts/resize-file-for-training.py @@ -0,0 +1,54 @@ +import os +from PIL import Image + +def preprocess_dataset(input_dir, output_dir, target_size=(512, 512)): + """ + Center crops images to a square, resizes them to 512x512, + and converts them to PNG format. + """ + os.makedirs(output_dir, exist_ok=True) + + valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff') + + for root, _, files in os.walk(input_dir): + for file in files: + if file.lower().endswith(valid_extensions): + # Setup paths + rel_path = os.path.relpath(root, input_dir) + out_folder = os.path.join(output_dir, rel_path) + os.makedirs(out_folder, exist_ok=True) + + img_path = os.path.join(root, file) + filename_without_ext = os.path.splitext(file)[0] + save_path = os.path.join(out_folder, f"{filename_without_ext}.png") + + with Image.open(img_path) as img: + w, h = img.size + + # 1. Calculate center crop box + min_dim = min(w, h) + left = (w - min_dim) // 2 + top = (h - min_dim) // 2 + right = left + min_dim + bottom = top + min_dim + + # 2. Crop to square + img_cropped = img.crop((left, top, right, bottom)) + + # 3. Resize to target resolution (512x512) + # For masks (binary), use NEAREST; for images, use LANCZOS + if "ground_truth" in root.lower() or "mask" in root.lower(): + img_resized = img_cropped.resize(target_size, Image.Resampling.NEAREST) + else: + img_resized = img_cropped.resize(target_size, Image.Resampling.LANCZOS) + + # 4. Save as PNG + img_resized.save(save_path, "PNG") + print(f"Processed: {file} -> {save_path}") + +# Example Usage: +preprocess_dataset( + input_dir="./engine/DefectFill/data/xray_PCB", # + # "./engine/DefectFill/data/xray_PCB/train/defective_masks/xray_die" + output_dir="./engine/DefectFill/data/xray_PCB_dataset_512" +) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/scripts/test-chromadb.py b/ArtiAgent - DefectFill/src/scripts/test-chromadb.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e82e13258a9ae77627e0f430eaf4b0edfc9dd4 --- /dev/null +++ b/ArtiAgent - DefectFill/src/scripts/test-chromadb.py @@ -0,0 +1,19 @@ +import sys +import os +import json +# Add parent directory (src/) to path so 'pipeline' is findable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from pipeline.defect_rag import get_rag + +rag = get_rag() +all_data = rag.collection.get() + +print("Found IDs in ChromaDB:") +for doc_id in all_data['ids']: + print(" -", doc_id) + +results = rag.collection.get(ids=["my_vcsel_captures_bubble_001"]) # or any ID you know exists +print(results['metadatas'][0]) +print("---") +print(json.loads(results['metadatas'][0]['paths'])) \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/scripts/test-safetensor.py b/ArtiAgent - DefectFill/src/scripts/test-safetensor.py new file mode 100644 index 0000000000000000000000000000000000000000..af5c848160d04e2c5b7f6b7ce88cfad83f43c4ff --- /dev/null +++ b/ArtiAgent - DefectFill/src/scripts/test-safetensor.py @@ -0,0 +1,12 @@ +import os +from safetensors.torch import load_file + +ckpt_path = r"C:\path\to\your\flux1-dev.safetensors" # or wherever your checkpoint is +print(f"File size: {os.path.getsize(ckpt_path) / 1e9:.2f} GB") + +state_dict = load_file(ckpt_path) +print(f"Total keys in checkpoint: {len(state_dict)}") + +# Check which double_blocks are present +blocks = sorted(set([k.split('.')[1] for k in state_dict.keys() if 'double_blocks' in k])) +print(f"Blocks present: {blocks}") \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/segment_anything/.flake8 b/ArtiAgent - DefectFill/src/segment_anything/.flake8 new file mode 100644 index 0000000000000000000000000000000000000000..6b0759587aa5756e66a13ef034c6bcdd76a885f5 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/.flake8 @@ -0,0 +1,7 @@ +[flake8] +ignore = W503, E203, E221, C901, C408, E741, C407, B017, F811, C101, EXE001, EXE002 +max-line-length = 100 +max-complexity = 18 +select = B,C,E,F,W,T4,B9 +per-file-ignores = + **/__init__.py:F401,F403,E402 diff --git a/ArtiAgent - DefectFill/src/segment_anything/CODE_OF_CONDUCT.md b/ArtiAgent - DefectFill/src/segment_anything/CODE_OF_CONDUCT.md new file mode 100644 index 0000000000000000000000000000000000000000..08b500a221857ec3f451338e80b4a9ab1173a1af --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/CODE_OF_CONDUCT.md @@ -0,0 +1,80 @@ +# Code of Conduct + +## Our Pledge + +In the interest of fostering an open and welcoming environment, we as +contributors and maintainers pledge to make participation in our project and +our community a harassment-free experience for everyone, regardless of age, body +size, disability, ethnicity, sex characteristics, gender identity and expression, +level of experience, education, socio-economic status, nationality, personal +appearance, race, religion, or sexual identity and orientation. + +## Our Standards + +Examples of behavior that contributes to creating a positive environment +include: + +* Using welcoming and inclusive language +* Being respectful of differing viewpoints and experiences +* Gracefully accepting constructive criticism +* Focusing on what is best for the community +* Showing empathy towards other community members + +Examples of unacceptable behavior by participants include: + +* The use of sexualized language or imagery and unwelcome sexual attention or + advances +* Trolling, insulting/derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or electronic + address, without explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Our Responsibilities + +Project maintainers are responsible for clarifying the standards of acceptable +behavior and are expected to take appropriate and fair corrective action in +response to any instances of unacceptable behavior. + +Project maintainers have the right and responsibility to remove, edit, or +reject comments, commits, code, wiki edits, issues, and other contributions +that are not aligned to this Code of Conduct, or to ban temporarily or +permanently any contributor for other behaviors that they deem inappropriate, +threatening, offensive, or harmful. + +## Scope + +This Code of Conduct applies within all project spaces, and it also applies when +an individual is representing the project or its community in public spaces. +Examples of representing a project or community include using an official +project e-mail address, posting via an official social media account, or acting +as an appointed representative at an online or offline event. Representation of +a project may be further defined and clarified by project maintainers. + +This Code of Conduct also applies outside the project spaces when there is a +reasonable belief that an individual's behavior may have a negative impact on +the project or its community. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported by contacting the project team at . All +complaints will be reviewed and investigated and will result in a response that +is deemed necessary and appropriate to the circumstances. The project team is +obligated to maintain confidentiality with regard to the reporter of an incident. +Further details of specific enforcement policies may be posted separately. + +Project maintainers who do not follow or enforce the Code of Conduct in good +faith may face temporary or permanent repercussions as determined by other +members of the project's leadership. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, +available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see +https://www.contributor-covenant.org/faq diff --git a/ArtiAgent - DefectFill/src/segment_anything/CONTRIBUTING.md b/ArtiAgent - DefectFill/src/segment_anything/CONTRIBUTING.md new file mode 100644 index 0000000000000000000000000000000000000000..263991c9496cf29ed4b99e03a9fb9a38e6bfaf86 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/CONTRIBUTING.md @@ -0,0 +1,31 @@ +# Contributing to segment-anything +We want to make contributing to this project as easy and transparent as +possible. + +## Pull Requests +We actively welcome your pull requests. + +1. Fork the repo and create your branch from `main`. +2. If you've added code that should be tested, add tests. +3. If you've changed APIs, update the documentation. +4. Ensure the test suite passes. +5. Make sure your code lints, using the `linter.sh` script in the project's root directory. Linting requires `black==23.*`, `isort==5.12.0`, `flake8`, and `mypy`. +6. If you haven't already, complete the Contributor License Agreement ("CLA"). + +## Contributor License Agreement ("CLA") +In order to accept your pull request, we need you to submit a CLA. You only need +to do this once to work on any of Facebook's open source projects. + +Complete your CLA here: + +## Issues +We use GitHub issues to track public bugs. Please ensure your description is +clear and has sufficient instructions to be able to reproduce the issue. + +Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe +disclosure of security bugs. In those cases, please go through the process +outlined on that page and do not file a public issue. + +## License +By contributing to segment-anything, you agree that your contributions will be licensed +under the LICENSE file in the root directory of this source tree. diff --git a/ArtiAgent - DefectFill/src/segment_anything/LICENSE b/ArtiAgent - DefectFill/src/segment_anything/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/ArtiAgent - DefectFill/src/segment_anything/README.md b/ArtiAgent - DefectFill/src/segment_anything/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6256d2b7f5a387988338d538df4e699eb17ba702 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/README.md @@ -0,0 +1,107 @@ +# Segment Anything + +**[Meta AI Research, FAIR](https://ai.facebook.com/research/)** + +[Alexander Kirillov](https://alexander-kirillov.github.io/), [Eric Mintun](https://ericmintun.github.io/), [Nikhila Ravi](https://nikhilaravi.com/), [Hanzi Mao](https://hanzimao.me/), Chloe Rolland, Laura Gustafson, [Tete Xiao](https://tetexiao.com), [Spencer Whitehead](https://www.spencerwhitehead.com/), Alex Berg, Wan-Yen Lo, [Piotr Dollar](https://pdollar.github.io/), [Ross Girshick](https://www.rossgirshick.info/) + +[[`Paper`](https://ai.facebook.com/research/publications/segment-anything/)] [[`Project`](https://segment-anything.com/)] [[`Demo`](https://segment-anything.com/demo)] [[`Dataset`](https://segment-anything.com/dataset/index.html)] [[`Blog`](https://ai.facebook.com/blog/segment-anything-foundation-model-image-segmentation/)] + +![SAM design](assets/model_diagram.png?raw=true) + +The **Segment Anything Model (SAM)** produces high quality object masks from input prompts such as points or boxes, and it can be used to generate masks for all objects in an image. It has been trained on a [dataset](https://segment-anything.com/dataset/index.html) of 11 million images and 1.1 billion masks, and has strong zero-shot performance on a variety of segmentation tasks. + +

+ + +

+ +## Installation + +The code requires `python>=3.8`, as well as `pytorch>=1.7` and `torchvision>=0.8`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended. + +Install Segment Anything: + +``` +pip install git+https://github.com/facebookresearch/segment-anything.git +``` + +or clone the repository locally and install with + +``` +git clone git@github.com:facebookresearch/segment-anything.git +cd segment-anything; pip install -e . +``` + +The following optional dependencies are necessary for mask post-processing, saving masks in COCO format, the example notebooks, and exporting the model in ONNX format. `jupyter` is also required to run the example notebooks. +``` +pip install opencv-python pycocotools matplotlib onnxruntime onnx +``` + + +## Getting Started + +First download a [model checkpoint](#model-checkpoints). Then the model can be used in just a few lines to get masks from a given prompt: + +``` +from segment_anything import build_sam, SamPredictor +predictor = SamPredictor(build_sam(checkpoint="")) +predictor.set_image() +masks, _, _ = predictor.predict() +``` + +or generate masks for an entire image: + +``` +from segment_anything import build_sam, SamAutomaticMaskGenerator +mask_generator = SamAutomaticMaskGenerator(build_sam(checkpoint="")) +masks = mask_generator_generate() +``` + +Additionally, masks can be generated for images from the command line: + +``` +python scripts/amg.py --checkpoint --input --output +``` + +See the examples notebooks on [using SAM with prompts](/notebooks/predictor_example.ipynb) and [automatically generating masks](/notebooks/automatic_mask_generator_example.ipynb) for more details. + +

+ + +

+ +## ONNX Export + +SAM's lightweight mask decoder can be exported to ONNX format so that it can be run in any environment that supports ONNX runtime, such as in-browser as showcased in the [demo](https://segment-anything.com/demo). Export the model with + +``` +python scripts/export_onnx_model.py --checkpoint --output +``` + +See the [example notebook](https://github.com/facebookresearch/segment-anything/blob/main/notebooks/onnx_model_example.ipynb) for details on how to combine image preprocessing via SAM's backbone with mask prediction using the ONNX model. It is recommended to use the latest stable version of PyTorch for ONNX export. + +## Model Checkpoints + +Three model versions of the model are available with different backbone sizes. These models can be instantiated by running +``` +from segment_anything import sam_model_registry +sam = sam_model_registry[""](checkpoint="") +``` +Click the links below to download the checkpoint for the corresponding model name. The default model in bold can also be instantiated with `build_sam`, as in the examples in [Getting Started](#getting-started). + +* **`default` or `vit_h`: [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)** +* `vit_l`: [ViT-L SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth) +* `vit_b`: [ViT-B SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth) + +## License +The model is licensed under the [Apache 2.0 license](LICENSE). + +## Contributing + +See [contributing](CONTRIBUTING.md) and the [code of conduct](CODE_OF_CONDUCT.md). + +## Contributors + +The Segment Anything project was made possible with the help of many contributors (alphabetical): + +Aaron Adcock, Vaibhav Aggarwal, Morteza Behrooz, Cheng-Yang Fu, Ashley Gabriel, Ahuva Goldstand, Allen Goodman, Sumanth Gurram, Jiabo Hu, Somya Jain, Devansh Kukreja, Robert Kuo, Joshua Lane, Yanghao Li, Lilian Luong, Jitendra Malik, Mallika Malhotra, William Ngan, Omkar Parkhi, Nikhil Raina, Dirk Rowe, Neil Sejoor, Vanessa Stark, Bala Varadarajan, Bram Wasti, Zachary Winstrom diff --git a/ArtiAgent - DefectFill/src/segment_anything/linter.sh b/ArtiAgent - DefectFill/src/segment_anything/linter.sh new file mode 100644 index 0000000000000000000000000000000000000000..df2e17436d30e89ff1728109301599f425f1ad6b --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/linter.sh @@ -0,0 +1,32 @@ +#!/bin/bash -e +# Copyright (c) Facebook, Inc. and its affiliates. + +{ + black --version | grep -E "23\." > /dev/null +} || { + echo "Linter requires 'black==23.*' !" + exit 1 +} + +ISORT_VERSION=$(isort --version-number) +if [[ "$ISORT_VERSION" != 5.12* ]]; then + echo "Linter requires isort==5.12.0 !" + exit 1 +fi + +echo "Running isort ..." +isort . --atomic + +echo "Running black ..." +black -l 100 . + +echo "Running flake8 ..." +if [ -x "$(command -v flake8)" ]; then + flake8 . +else + python3 -m flake8 . +fi + +echo "Running mypy..." + +mypy --exclude 'setup.py|notebooks' . diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__init__.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b3668adc7817cb24a54cfe4405184a8409c6cb44 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from .build_sam import ( + build_sam, + build_sam_vit_h, + build_sam_vit_l, + build_sam_vit_b, + sam_model_registry, +) +from .build_sam_hq import ( + build_sam_hq, + build_sam_hq_vit_h, + build_sam_hq_vit_l, + build_sam_hq_vit_b, + sam_hq_model_registry, +) +from .predictor import SamPredictor +from .automatic_mask_generator import SamAutomaticMaskGenerator diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92fb0fdb0d5f178dc331050f0aac3aa1d3a387d7 Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eed9e46a76b34836e4ef869f5ec1da73639fcf2f Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/automatic_mask_generator.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/build_sam.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/build_sam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de4e075c325215d0aacd2fea2a3a4c3d70a45a0d Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/build_sam.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/build_sam_hq.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/build_sam_hq.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebe56fbc467c52e21016263e1e0995d5cd2c7a39 Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/build_sam_hq.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/predictor.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/predictor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db31f5b2be04f0852b497825356ee3917c11d424 Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/__pycache__/predictor.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/automatic_mask_generator.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/automatic_mask_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..23264971b7ff5aa0b4f499ade7773b68dce984b6 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/automatic_mask_generator.py @@ -0,0 +1,372 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torchvision.ops.boxes import batched_nms, box_area # type: ignore + +from typing import Any, Dict, List, Optional, Tuple + +from .modeling import Sam +from .predictor import SamPredictor +from .utils.amg import ( + MaskData, + area_from_rle, + batch_iterator, + batched_mask_to_box, + box_xyxy_to_xywh, + build_all_layer_point_grids, + calculate_stability_score, + coco_encode_rle, + generate_crop_boxes, + is_box_near_crop_edge, + mask_to_rle_pytorch, + remove_small_regions, + rle_to_mask, + uncrop_boxes_xyxy, + uncrop_masks, + uncrop_points, +) + + +class SamAutomaticMaskGenerator: + def __init__( + self, + model: Sam, + points_per_side: Optional[int] = 32, + points_per_batch: int = 64, + pred_iou_thresh: float = 0.88, + stability_score_thresh: float = 0.95, + stability_score_offset: float = 1.0, + box_nms_thresh: float = 0.7, + crop_n_layers: int = 0, + crop_nms_thresh: float = 0.7, + crop_overlap_ratio: float = 512 / 1500, + crop_n_points_downscale_factor: int = 1, + point_grids: Optional[List[np.ndarray]] = None, + min_mask_region_area: int = 0, + output_mode: str = "binary_mask", + ) -> None: + """ + Using a SAM model, generates masks for the entire image. + Generates a grid of point prompts over the image, then filters + low quality and duplicate masks. The default settings are chosen + for SAM with a ViT-H backbone. + + Arguments: + model (Sam): The SAM model to use for mask prediction. + points_per_side (int or None): The number of points to be sampled + along one side of the image. The total number of points is + points_per_side**2. If None, 'point_grids' must provide explicit + point sampling. + points_per_batch (int): Sets the number of points run simultaneously + by the model. Higher numbers may be faster but use more GPU memory. + pred_iou_thresh (float): A filtering threshold in [0,1], using the + model's predicted mask quality. + stability_score_thresh (float): A filtering threshold in [0,1], using + the stability of the mask under changes to the cutoff used to binarize + the model's mask predictions. + stability_score_offset (float): The amount to shift the cutoff when + calculated the stability score. + box_nms_thresh (float): The box IoU cutoff used by non-maximal + suppression to filter duplicate masks. + crops_n_layers (int): If >0, mask prediction will be run again on + crops of the image. Sets the number of layers to run, where each + layer has 2**i_layer number of image crops. + crops_nms_thresh (float): The box IoU cutoff used by non-maximal + suppression to filter duplicate masks between different crops. + crop_overlap_ratio (float): Sets the degree to which crops overlap. + In the first crop layer, crops will overlap by this fraction of + the image length. Later layers with more crops scale down this overlap. + crop_n_points_downscale_factor (int): The number of points-per-side + sampled in layer n is scaled down by crop_n_points_downscale_factor**n. + point_grids (list(np.ndarray) or None): A list over explicit grids + of points used for sampling, normalized to [0,1]. The nth grid in the + list is used in the nth crop layer. Exclusive with points_per_side. + min_mask_region_area (int): If >0, postprocessing will be applied + to remove disconnected regions and holes in masks with area smaller + than min_mask_region_area. Requires opencv. + output_mode (str): The form masks are returned in. Can be 'binary_mask', + 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools. + For large resolutions, 'binary_mask' may consume large amounts of + memory. + """ + + assert (points_per_side is None) != ( + point_grids is None + ), "Exactly one of points_per_side or point_grid must be provided." + if points_per_side is not None: + self.point_grids = build_all_layer_point_grids( + points_per_side, + crop_n_layers, + crop_n_points_downscale_factor, + ) + elif point_grids is not None: + self.point_grids = point_grids + else: + raise ValueError("Can't have both points_per_side and point_grid be None.") + + assert output_mode in [ + "binary_mask", + "uncompressed_rle", + "coco_rle", + ], f"Unknown output_mode {output_mode}." + if output_mode == "coco_rle": + from pycocotools import mask as mask_utils # type: ignore # noqa: F401 + + if min_mask_region_area > 0: + import cv2 # type: ignore # noqa: F401 + + self.predictor = SamPredictor(model) + self.points_per_batch = points_per_batch + self.pred_iou_thresh = pred_iou_thresh + self.stability_score_thresh = stability_score_thresh + self.stability_score_offset = stability_score_offset + self.box_nms_thresh = box_nms_thresh + self.crop_n_layers = crop_n_layers + self.crop_nms_thresh = crop_nms_thresh + self.crop_overlap_ratio = crop_overlap_ratio + self.crop_n_points_downscale_factor = crop_n_points_downscale_factor + self.min_mask_region_area = min_mask_region_area + self.output_mode = output_mode + + @torch.no_grad() + def generate(self, image: np.ndarray) -> List[Dict[str, Any]]: + """ + Generates masks for the given image. + + Arguments: + image (np.ndarray): The image to generate masks for, in HWC uint8 format. + + Returns: + list(dict(str, any)): A list over records for masks. Each record is + a dict containing the following keys: + segmentation (dict(str, any) or np.ndarray): The mask. If + output_mode='binary_mask', is an array of shape HW. Otherwise, + is a dictionary containing the RLE. + bbox (list(float)): The box around the mask, in XYWH format. + area (int): The area in pixels of the mask. + predicted_iou (float): The model's own prediction of the mask's + quality. This is filtered by the pred_iou_thresh parameter. + point_coords (list(list(float))): The point coordinates input + to the model to generate this mask. + stability_score (float): A measure of the mask's quality. This + is filtered on using the stability_score_thresh parameter. + crop_box (list(float)): The crop of the image used to generate + the mask, given in XYWH format. + """ + + # Generate masks + mask_data = self._generate_masks(image) + + # Filter small disconnected regions and holes in masks + if self.min_mask_region_area > 0: + mask_data = self.postprocess_small_regions( + mask_data, + self.min_mask_region_area, + max(self.box_nms_thresh, self.crop_nms_thresh), + ) + + # Encode masks + if self.output_mode == "coco_rle": + mask_data["segmentations"] = [coco_encode_rle(rle) for rle in mask_data["rles"]] + elif self.output_mode == "binary_mask": + mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]] + else: + mask_data["segmentations"] = mask_data["rles"] + + # Write mask records + curr_anns = [] + for idx in range(len(mask_data["segmentations"])): + ann = { + "segmentation": mask_data["segmentations"][idx], + "area": area_from_rle(mask_data["rles"][idx]), + "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(), + "predicted_iou": mask_data["iou_preds"][idx].item(), + "point_coords": [mask_data["points"][idx].tolist()], + "stability_score": mask_data["stability_score"][idx].item(), + "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(), + } + curr_anns.append(ann) + + return curr_anns + + def _generate_masks(self, image: np.ndarray) -> MaskData: + orig_size = image.shape[:2] + crop_boxes, layer_idxs = generate_crop_boxes( + orig_size, self.crop_n_layers, self.crop_overlap_ratio + ) + + # Iterate over image crops + data = MaskData() + for crop_box, layer_idx in zip(crop_boxes, layer_idxs): + crop_data = self._process_crop(image, crop_box, layer_idx, orig_size) + data.cat(crop_data) + + # Remove duplicate masks between crops + if len(crop_boxes) > 1: + # Prefer masks from smaller crops + scores = 1 / box_area(data["crop_boxes"]) + scores = scores.to(data["boxes"].device) + keep_by_nms = batched_nms( + data["boxes"].float(), + scores, + torch.zeros(len(data["boxes"])), # categories + iou_threshold=self.crop_nms_thresh, + ) + data.filter(keep_by_nms) + + data.to_numpy() + return data + + def _process_crop( + self, + image: np.ndarray, + crop_box: List[int], + crop_layer_idx: int, + orig_size: Tuple[int, ...], + ) -> MaskData: + # Crop the image and calculate embeddings + x0, y0, x1, y1 = crop_box + cropped_im = image[y0:y1, x0:x1, :] + cropped_im_size = cropped_im.shape[:2] + self.predictor.set_image(cropped_im) + + # Get points for this crop + points_scale = np.array(cropped_im_size)[None, ::-1] + points_for_image = self.point_grids[crop_layer_idx] * points_scale + + # Generate masks for this crop in batches + data = MaskData() + for (points,) in batch_iterator(self.points_per_batch, points_for_image): + batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size) + data.cat(batch_data) + del batch_data + self.predictor.reset_image() + + # Remove duplicates within this crop. + keep_by_nms = batched_nms( + data["boxes"].float(), + data["iou_preds"], + torch.zeros(len(data["boxes"])), # categories + iou_threshold=self.box_nms_thresh, + ) + data.filter(keep_by_nms) + + # Return to the original image frame + data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box) + data["points"] = uncrop_points(data["points"], crop_box) + data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))]) + + return data + + def _process_batch( + self, + points: np.ndarray, + im_size: Tuple[int, ...], + crop_box: List[int], + orig_size: Tuple[int, ...], + ) -> MaskData: + orig_h, orig_w = orig_size + + # Run model on this batch + transformed_points = self.predictor.transform.apply_coords(points, im_size) + in_points = torch.as_tensor(transformed_points, device=self.predictor.device) + in_labels = torch.ones(in_points.shape[0], dtype=torch.int, device=in_points.device) + masks, iou_preds, _ = self.predictor.predict_torch( + in_points[:, None, :], + in_labels[:, None], + multimask_output=True, + return_logits=True, + ) + + # Serialize predictions and store in MaskData + data = MaskData( + masks=masks.flatten(0, 1), + iou_preds=iou_preds.flatten(0, 1), + points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)), + ) + del masks + + # Filter by predicted IoU + if self.pred_iou_thresh > 0.0: + keep_mask = data["iou_preds"] > self.pred_iou_thresh + data.filter(keep_mask) + + # Calculate stability score + data["stability_score"] = calculate_stability_score( + data["masks"], self.predictor.model.mask_threshold, self.stability_score_offset + ) + if self.stability_score_thresh > 0.0: + keep_mask = data["stability_score"] >= self.stability_score_thresh + data.filter(keep_mask) + + # Threshold masks and calculate boxes + data["masks"] = data["masks"] > self.predictor.model.mask_threshold + data["boxes"] = batched_mask_to_box(data["masks"]) + + # Filter boxes that touch crop boundaries + keep_mask = ~is_box_near_crop_edge(data["boxes"], crop_box, [0, 0, orig_w, orig_h]) + if not torch.all(keep_mask): + data.filter(keep_mask) + + # Compress to RLE + data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w) + data["rles"] = mask_to_rle_pytorch(data["masks"]) + del data["masks"] + + return data + + @staticmethod + def postprocess_small_regions( + mask_data: MaskData, min_area: int, nms_thresh: float + ) -> MaskData: + """ + Removes small disconnected regions and holes in masks, then reruns + box NMS to remove any new duplicates. + + Edits mask_data in place. + + Requires open-cv as a dependency. + """ + if len(mask_data["rles"]) == 0: + return mask_data + + # Filter small disconnected regions and holes + new_masks = [] + scores = [] + for rle in mask_data["rles"]: + mask = rle_to_mask(rle) + + mask, changed = remove_small_regions(mask, min_area, mode="holes") + unchanged = not changed + mask, changed = remove_small_regions(mask, min_area, mode="islands") + unchanged = unchanged and not changed + + new_masks.append(torch.as_tensor(mask).unsqueeze(0)) + # Give score=0 to changed masks and score=1 to unchanged masks + # so NMS will prefer ones that didn't need postprocessing + scores.append(float(unchanged)) + + # Recalculate boxes and remove any new duplicates + masks = torch.cat(new_masks, dim=0) + boxes = batched_mask_to_box(masks) + keep_by_nms = batched_nms( + boxes.float(), + torch.as_tensor(scores), + torch.zeros(len(boxes)), # categories + iou_threshold=nms_thresh, + ) + + # Only recalculate RLEs for masks that have changed + for i_mask in keep_by_nms: + if scores[i_mask] == 0.0: + mask_torch = masks[i_mask].unsqueeze(0) + mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0] + mask_data["boxes"][i_mask] = boxes[i_mask] # update res directly + mask_data.filter(keep_by_nms) + + return mask_data diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/build_sam.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/build_sam.py new file mode 100644 index 0000000000000000000000000000000000000000..07abfca24e96eced7f13bdefd3212ce1b77b8999 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/build_sam.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from functools import partial + +from .modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer + + +def build_sam_vit_h(checkpoint=None): + return _build_sam( + encoder_embed_dim=1280, + encoder_depth=32, + encoder_num_heads=16, + encoder_global_attn_indexes=[7, 15, 23, 31], + checkpoint=checkpoint, + ) + + +build_sam = build_sam_vit_h + + +def build_sam_vit_l(checkpoint=None): + return _build_sam( + encoder_embed_dim=1024, + encoder_depth=24, + encoder_num_heads=16, + encoder_global_attn_indexes=[5, 11, 17, 23], + checkpoint=checkpoint, + ) + + +def build_sam_vit_b(checkpoint=None): + return _build_sam( + encoder_embed_dim=768, + encoder_depth=12, + encoder_num_heads=12, + encoder_global_attn_indexes=[2, 5, 8, 11], + checkpoint=checkpoint, + ) + + +sam_model_registry = { + "default": build_sam, + "vit_h": build_sam, + "vit_l": build_sam_vit_l, + "vit_b": build_sam_vit_b, +} + + +def _build_sam( + encoder_embed_dim, + encoder_depth, + encoder_num_heads, + encoder_global_attn_indexes, + checkpoint=None, +): + prompt_embed_dim = 256 + image_size = 1024 + vit_patch_size = 16 + image_embedding_size = image_size // vit_patch_size + sam = Sam( + image_encoder=ImageEncoderViT( + depth=encoder_depth, + embed_dim=encoder_embed_dim, + img_size=image_size, + mlp_ratio=4, + norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), + num_heads=encoder_num_heads, + patch_size=vit_patch_size, + qkv_bias=True, + use_rel_pos=True, + global_attn_indexes=encoder_global_attn_indexes, + window_size=14, + out_chans=prompt_embed_dim, + ), + prompt_encoder=PromptEncoder( + embed_dim=prompt_embed_dim, + image_embedding_size=(image_embedding_size, image_embedding_size), + input_image_size=(image_size, image_size), + mask_in_chans=16, + ), + mask_decoder=MaskDecoder( + num_multimask_outputs=3, + transformer=TwoWayTransformer( + depth=2, + embedding_dim=prompt_embed_dim, + mlp_dim=2048, + num_heads=8, + ), + transformer_dim=prompt_embed_dim, + iou_head_depth=3, + iou_head_hidden_dim=256, + ), + pixel_mean=[123.675, 116.28, 103.53], + pixel_std=[58.395, 57.12, 57.375], + ) + sam.eval() + if checkpoint is not None: + with open(checkpoint, "rb") as f: + state_dict = torch.load(f) + sam.load_state_dict(state_dict) + return sam diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/build_sam_hq.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/build_sam_hq.py new file mode 100644 index 0000000000000000000000000000000000000000..a113b745c9772cb2f5a34a81c0626e9161699796 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/build_sam_hq.py @@ -0,0 +1,114 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from functools import partial + +from .modeling import ImageEncoderViT, MaskDecoderHQ, PromptEncoder, Sam, TwoWayTransformer + + +def build_sam_hq_vit_h(checkpoint=None): + return _build_sam( + encoder_embed_dim=1280, + encoder_depth=32, + encoder_num_heads=16, + encoder_global_attn_indexes=[7, 15, 23, 31], + checkpoint=checkpoint, + ) + + +build_sam_hq = build_sam_hq_vit_h + + +def build_sam_hq_vit_l(checkpoint=None): + return _build_sam( + encoder_embed_dim=1024, + encoder_depth=24, + encoder_num_heads=16, + encoder_global_attn_indexes=[5, 11, 17, 23], + checkpoint=checkpoint, + ) + + +def build_sam_hq_vit_b(checkpoint=None): + return _build_sam( + encoder_embed_dim=768, + encoder_depth=12, + encoder_num_heads=12, + encoder_global_attn_indexes=[2, 5, 8, 11], + checkpoint=checkpoint, + ) + + +sam_hq_model_registry = { + "default": build_sam_hq_vit_h, + "vit_h": build_sam_hq_vit_h, + "vit_l": build_sam_hq_vit_l, + "vit_b": build_sam_hq_vit_b, +} + + +def _build_sam( + encoder_embed_dim, + encoder_depth, + encoder_num_heads, + encoder_global_attn_indexes, + checkpoint=None, +): + prompt_embed_dim = 256 + image_size = 1024 + vit_patch_size = 16 + image_embedding_size = image_size // vit_patch_size + sam = Sam( + image_encoder=ImageEncoderViT( + depth=encoder_depth, + embed_dim=encoder_embed_dim, + img_size=image_size, + mlp_ratio=4, + norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), + num_heads=encoder_num_heads, + patch_size=vit_patch_size, + qkv_bias=True, + use_rel_pos=True, + global_attn_indexes=encoder_global_attn_indexes, + window_size=14, + out_chans=prompt_embed_dim, + ), + prompt_encoder=PromptEncoder( + embed_dim=prompt_embed_dim, + image_embedding_size=(image_embedding_size, image_embedding_size), + input_image_size=(image_size, image_size), + mask_in_chans=16, + ), + mask_decoder=MaskDecoderHQ( + num_multimask_outputs=3, + transformer=TwoWayTransformer( + depth=2, + embedding_dim=prompt_embed_dim, + mlp_dim=2048, + num_heads=8, + ), + transformer_dim=prompt_embed_dim, + iou_head_depth=3, + iou_head_hidden_dim=256, + vit_dim=encoder_embed_dim, + ), + pixel_mean=[123.675, 116.28, 103.53], + pixel_std=[58.395, 57.12, 57.375], + ) + # sam.eval() + if checkpoint is not None: + with open(checkpoint, "rb") as f: + device = "cuda" if torch.cuda.is_available() else "cpu" + state_dict = torch.load(f, map_location=device) + info = sam.load_state_dict(state_dict, strict=False) + print(info) + for n, p in sam.named_parameters(): + if 'hf_token' not in n and 'hf_mlp' not in n and 'compress_vit_feat' not in n and 'embedding_encoder' not in n and 'embedding_maskfeature' not in n: + p.requires_grad = False + + return sam diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__init__.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..71172d22345eff1f3729c6326299feee17717ccc --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from .sam import Sam +from .image_encoder import ImageEncoderViT +from .mask_decoder_hq import MaskDecoderHQ +from .mask_decoder import MaskDecoder +from .prompt_encoder import PromptEncoder +from .transformer import TwoWayTransformer diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87f1d8c2ac93e4982fabb8c8e6920f5cbc1359d8 Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/common.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/common.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf2be6697355c877d54c9d107c6659efcfc4a84e Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/common.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..692059b3bca1c5c92ecc3f0c5d1eef99c9730abd Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/image_encoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6f379c2d75c665dc0576d4d8aecbf9a0cea1e3d Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder_hq.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder_hq.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..870de23dc2f41fa1c47fa40aed9ae92925620bfb Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/mask_decoder_hq.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22a6147114c235c720c5d5030c6ac80f7a1ed4f2 Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/prompt_encoder.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/sam.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/sam.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a224c8aa2b4fb41f9b049e4b4064762d6caccb2 Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/sam.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cba0adf9585bf375a4f187bd232e7e8727ea5fff Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/__pycache__/transformer.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/common.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/common.py new file mode 100644 index 0000000000000000000000000000000000000000..2bf15236a3eb24d8526073bc4fa2b274cccb3f96 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/common.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn + +from typing import Type + + +class MLPBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + mlp_dim: int, + act: Type[nn.Module] = nn.GELU, + ) -> None: + super().__init__() + self.lin1 = nn.Linear(embedding_dim, mlp_dim) + self.lin2 = nn.Linear(mlp_dim, embedding_dim) + self.act = act() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lin2(self.act(self.lin1(x))) + + +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa +# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa +class LayerNorm2d(nn.Module): + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/image_encoder.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/image_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..9c01f71c906e90d9b9b7d3252cdd3e5c555a2734 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/image_encoder.py @@ -0,0 +1,398 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from typing import Optional, Tuple, Type + +from .common import LayerNorm2d, MLPBlock + + +# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa +class ImageEncoderViT(nn.Module): + def __init__( + self, + img_size: int = 1024, + patch_size: int = 16, + in_chans: int = 3, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + out_chans: int = 256, + qkv_bias: bool = True, + norm_layer: Type[nn.Module] = nn.LayerNorm, + act_layer: Type[nn.Module] = nn.GELU, + use_abs_pos: bool = True, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + window_size: int = 0, + global_attn_indexes: Tuple[int, ...] = (), + ) -> None: + """ + Args: + img_size (int): Input image size. + patch_size (int): Patch size. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + depth (int): Depth of ViT. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_abs_pos (bool): If True, use absolute positional embeddings. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. + global_attn_indexes (list): Indexes for blocks using global attention. + """ + super().__init__() + self.img_size = img_size + + self.patch_embed = PatchEmbed( + kernel_size=(patch_size, patch_size), + stride=(patch_size, patch_size), + in_chans=in_chans, + embed_dim=embed_dim, + ) + + self.pos_embed: Optional[nn.Parameter] = None + if use_abs_pos: + # Initialize absolute positional embedding with pretrain image size. + self.pos_embed = nn.Parameter( + torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim) + ) + + self.blocks = nn.ModuleList() + for i in range(depth): + block = Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + norm_layer=norm_layer, + act_layer=act_layer, + use_rel_pos=use_rel_pos, + rel_pos_zero_init=rel_pos_zero_init, + window_size=window_size if i not in global_attn_indexes else 0, + input_size=(img_size // patch_size, img_size // patch_size), + ) + self.blocks.append(block) + + self.neck = nn.Sequential( + nn.Conv2d( + embed_dim, + out_chans, + kernel_size=1, + bias=False, + ), + LayerNorm2d(out_chans), + nn.Conv2d( + out_chans, + out_chans, + kernel_size=3, + padding=1, + bias=False, + ), + LayerNorm2d(out_chans), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.patch_embed(x) + if self.pos_embed is not None: + x = x + self.pos_embed + + interm_embeddings=[] + for blk in self.blocks: + x = blk(x) + if blk.window_size == 0: + interm_embeddings.append(x) + + x = self.neck(x.permute(0, 3, 1, 2)) + + return x, interm_embeddings + + +class Block(nn.Module): + """Transformer blocks with support of window attention and residual propagation blocks""" + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + norm_layer: Type[nn.Module] = nn.LayerNorm, + act_layer: Type[nn.Module] = nn.GELU, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + window_size: int = 0, + input_size: Optional[Tuple[int, int]] = None, + ) -> None: + """ + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads in each ViT block. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + norm_layer (nn.Module): Normalization layer. + act_layer (nn.Module): Activation layer. + use_rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + window_size (int): Window size for window attention blocks. If it equals 0, then + use global attention. + input_size (tuple(int, int) or None): Input resolution for calculating the relative + positional parameter size. + """ + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + use_rel_pos=use_rel_pos, + rel_pos_zero_init=rel_pos_zero_init, + input_size=input_size if window_size == 0 else (window_size, window_size), + ) + + self.norm2 = norm_layer(dim) + self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer) + + self.window_size = window_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = x + x = self.norm1(x) + # Window partition + if self.window_size > 0: + H, W = x.shape[1], x.shape[2] + x, pad_hw = window_partition(x, self.window_size) + + x = self.attn(x) + # Reverse window partition + if self.window_size > 0: + x = window_unpartition(x, self.window_size, pad_hw, (H, W)) + + x = shortcut + x + x = x + self.mlp(self.norm2(x)) + + return x + + +class Attention(nn.Module): + """Multi-head Attention block with relative position embeddings.""" + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = True, + use_rel_pos: bool = False, + rel_pos_zero_init: bool = True, + input_size: Optional[Tuple[int, int]] = None, + ) -> None: + """ + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + qkv_bias (bool): If True, add a learnable bias to query, key, value. + rel_pos (bool): If True, add relative positional embeddings to the attention map. + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. + input_size (tuple(int, int) or None): Input resolution for calculating the relative + positional parameter size. + """ + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.proj = nn.Linear(dim, dim) + + self.use_rel_pos = use_rel_pos + if self.use_rel_pos: + assert ( + input_size is not None + ), "Input size must be provided if using relative positional encoding." + # initialize relative positional embeddings + self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) + self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, H, W, _ = x.shape + # qkv with shape (3, B, nHead, H * W, C) + qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) + # q, k, v with shape (B * nHead, H * W, C) + q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0) + + attn = (q * self.scale) @ k.transpose(-2, -1) + + if self.use_rel_pos: + attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)) + + attn = attn.softmax(dim=-1) + x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1) + x = self.proj(x) + + return x + + +def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]: + """ + Partition into non-overlapping windows with padding if needed. + Args: + x (tensor): input tokens with [B, H, W, C]. + window_size (int): window size. + + Returns: + windows: windows after partition with [B * num_windows, window_size, window_size, C]. + (Hp, Wp): padded height and width before partition + """ + B, H, W, C = x.shape + + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + if pad_h > 0 or pad_w > 0: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + Hp, Wp = H + pad_h, W + pad_w + + x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows, (Hp, Wp) + + +def window_unpartition( + windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int] +) -> torch.Tensor: + """ + Window unpartition into original sequences and removing padding. + Args: + windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. + window_size (int): window size. + pad_hw (Tuple): padded height and width (Hp, Wp). + hw (Tuple): original height and width (H, W) before padding. + + Returns: + x: unpartitioned sequences with [B, H, W, C]. + """ + Hp, Wp = pad_hw + H, W = hw + B = windows.shape[0] // (Hp * Wp // window_size // window_size) + x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) + + if Hp > H or Wp > W: + x = x[:, :H, :W, :].contiguous() + return x + + +def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: + """ + Get relative positional embeddings according to the relative positions of + query and key sizes. + Args: + q_size (int): size of query q. + k_size (int): size of key k. + rel_pos (Tensor): relative position embeddings (L, C). + + Returns: + Extracted positional embeddings according to relative positions. + """ + max_rel_dist = int(2 * max(q_size, k_size) - 1) + # Interpolate rel pos if needed. + if rel_pos.shape[0] != max_rel_dist: + # Interpolate rel pos. + rel_pos_resized = F.interpolate( + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), + size=max_rel_dist, + mode="linear", + ) + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) + else: + rel_pos_resized = rel_pos + + # Scale the coords with short length if shapes for q and k are different. + q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) + k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) + + return rel_pos_resized[relative_coords.long()] + + +def add_decomposed_rel_pos( + attn: torch.Tensor, + q: torch.Tensor, + rel_pos_h: torch.Tensor, + rel_pos_w: torch.Tensor, + q_size: Tuple[int, int], + k_size: Tuple[int, int], +) -> torch.Tensor: + """ + Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. + https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950 + Args: + attn (Tensor): attention map. + q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C). + rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis. + rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis. + q_size (Tuple): spatial sequence size of query q with (q_h, q_w). + k_size (Tuple): spatial sequence size of key k with (k_h, k_w). + + Returns: + attn (Tensor): attention map with added relative positional embeddings. + """ + q_h, q_w = q_size + k_h, k_w = k_size + Rh = get_rel_pos(q_h, k_h, rel_pos_h) + Rw = get_rel_pos(q_w, k_w, rel_pos_w) + + B, _, dim = q.shape + r_q = q.reshape(B, q_h, q_w, dim) + rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) + rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) + + attn = ( + attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] + ).view(B, q_h * q_w, k_h * k_w) + + return attn + + +class PatchEmbed(nn.Module): + """ + Image to Patch Embedding. + """ + + def __init__( + self, + kernel_size: Tuple[int, int] = (16, 16), + stride: Tuple[int, int] = (16, 16), + padding: Tuple[int, int] = (0, 0), + in_chans: int = 3, + embed_dim: int = 768, + ) -> None: + """ + Args: + kernel_size (Tuple): kernel size of the projection layer. + stride (Tuple): stride of the projection layer. + padding (Tuple): padding size of the projection layer. + in_chans (int): Number of input image channels. + embed_dim (int): Patch embedding dimension. + """ + super().__init__() + + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + # B C H W -> B H W C + x = x.permute(0, 2, 3, 1) + return x \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/mask_decoder.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/mask_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c36c7b553c9df986dab91474de06d171feb7f93d --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/mask_decoder.py @@ -0,0 +1,178 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import nn +from torch.nn import functional as F + +from typing import List, Tuple, Type + +from .common import LayerNorm2d + + +class MaskDecoder(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), + activation(), + ) + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth + ) + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + hq_token_only: bool, + interm_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + """ + masks, iou_pred = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + ) + + # Select the correct mask or masks for output + if multimask_output: + mask_slice = slice(1, None) + else: + mask_slice = slice(0, 1) + masks = masks[:, mask_slice, :, :] + iou_pred = iou_pred[:, mask_slice] + + # Prepare output + return masks, iou_pred + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0) + output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + src = src + dense_prompt_embeddings + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, 0, :] + mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + upscaled_embedding = self.output_upscaling(src) + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])) + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding.shape + masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) + + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + + return masks, iou_pred + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/mask_decoder_hq.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/mask_decoder_hq.py new file mode 100644 index 0000000000000000000000000000000000000000..c4576f3495ae72d639b2278c4c252e3e02e5d424 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/mask_decoder_hq.py @@ -0,0 +1,232 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Modified by HQ-SAM team +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import nn +from torch.nn import functional as F + +from typing import List, Tuple, Type + +from .common import LayerNorm2d + + +class MaskDecoderHQ(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + vit_dim: int = 1024, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), + activation(), + ) + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth + ) + + # HQ-SAM parameters + self.hf_token = nn.Embedding(1, transformer_dim) # HQ-Ouptput-Token + self.hf_mlp = MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) # corresponding new MLP layer for HQ-Ouptput-Token + self.num_mask_tokens = self.num_mask_tokens + 1 + + # three conv fusion layers for obtaining HQ-Feature + self.compress_vit_feat = nn.Sequential( + nn.ConvTranspose2d(vit_dim, transformer_dim, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim), + nn.GELU(), + nn.ConvTranspose2d(transformer_dim, transformer_dim // 8, kernel_size=2, stride=2)) + + self.embedding_encoder = nn.Sequential( + nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2), + LayerNorm2d(transformer_dim // 4), + nn.GELU(), + nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2), + ) + self.embedding_maskfeature = nn.Sequential( + nn.Conv2d(transformer_dim // 8, transformer_dim // 4, 3, 1, 1), + LayerNorm2d(transformer_dim // 4), + nn.GELU(), + nn.Conv2d(transformer_dim // 4, transformer_dim // 8, 3, 1, 1)) + + + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + hq_token_only: bool, + interm_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the ViT image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + """ + vit_features = interm_embeddings[0].permute(0, 3, 1, 2) # early-layer ViT feature, after 1st global attention block in ViT + hq_features = self.embedding_encoder(image_embeddings) + self.compress_vit_feat(vit_features) + + masks, iou_pred = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + hq_features=hq_features, + ) + + # Select the correct mask or masks for output + if multimask_output: + # mask with highest score + mask_slice = slice(1,self.num_mask_tokens-1) + iou_pred = iou_pred[:, mask_slice] + iou_pred, max_iou_idx = torch.max(iou_pred,dim=1) + iou_pred = iou_pred.unsqueeze(1) + masks_multi = masks[:, mask_slice, :, :] + masks_sam = masks_multi[torch.arange(masks_multi.size(0)),max_iou_idx].unsqueeze(1) + else: + # singale mask output, default + mask_slice = slice(0, 1) + iou_pred = iou_pred[:,mask_slice] + masks_sam = masks[:,mask_slice] + + masks_hq = masks[:,slice(self.num_mask_tokens-1, self.num_mask_tokens)] + if hq_token_only: + masks = masks_hq + else: + masks = masks_sam + masks_hq + # Prepare output + return masks, iou_pred + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + hq_features: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight, self.hf_token.weight], dim=0) + output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + src = src + dense_prompt_embeddings + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, 0, :] + mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + + upscaled_embedding_sam = self.output_upscaling(src) + upscaled_embedding_hq = self.embedding_maskfeature(upscaled_embedding_sam) + hq_features.repeat(b,1,1,1) + + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + if i < self.num_mask_tokens - 1: + hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])) + else: + hyper_in_list.append(self.hf_mlp(mask_tokens_out[:, i, :])) + + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding_sam.shape + + masks_sam = (hyper_in[:,:self.num_mask_tokens-1] @ upscaled_embedding_sam.view(b, c, h * w)).view(b, -1, h, w) + masks_sam_hq = (hyper_in[:,self.num_mask_tokens-1:] @ upscaled_embedding_hq.view(b, c, h * w)).view(b, -1, h, w) + masks = torch.cat([masks_sam,masks_sam_hq],dim=1) + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + + return masks, iou_pred + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/prompt_encoder.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/prompt_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c3143f4f8e02ddd7ca8587b40ff5d47c3a6b7ef3 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/prompt_encoder.py @@ -0,0 +1,214 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torch import nn + +from typing import Any, Optional, Tuple, Type + +from .common import LayerNorm2d + + +class PromptEncoder(nn.Module): + def __init__( + self, + embed_dim: int, + image_embedding_size: Tuple[int, int], + input_image_size: Tuple[int, int], + mask_in_chans: int, + activation: Type[nn.Module] = nn.GELU, + ) -> None: + """ + Encodes prompts for input to SAM's mask decoder. + + Arguments: + embed_dim (int): The prompts' embedding dimension + image_embedding_size (tuple(int, int)): The spatial size of the + image embedding, as (H, W). + input_image_size (int): The padded size of the image as input + to the image encoder, as (H, W). + mask_in_chans (int): The number of hidden channels used for + encoding input masks. + activation (nn.Module): The activation to use when encoding + input masks. + """ + super().__init__() + self.embed_dim = embed_dim + self.input_image_size = input_image_size + self.image_embedding_size = image_embedding_size + self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) + + self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners + point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)] + self.point_embeddings = nn.ModuleList(point_embeddings) + self.not_a_point_embed = nn.Embedding(1, embed_dim) + + self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1]) + self.mask_downscaling = nn.Sequential( + nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans // 4), + activation(), + nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans), + activation(), + nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), + ) + self.no_mask_embed = nn.Embedding(1, embed_dim) + + def get_dense_pe(self) -> torch.Tensor: + """ + Returns the positional encoding used to encode point prompts, + applied to a dense set of points the shape of the image encoding. + + Returns: + torch.Tensor: Positional encoding with shape + 1x(embed_dim)x(embedding_h)x(embedding_w) + """ + return self.pe_layer(self.image_embedding_size).unsqueeze(0) + + def _embed_points( + self, + points: torch.Tensor, + labels: torch.Tensor, + pad: bool, + ) -> torch.Tensor: + """Embeds point prompts.""" + points = points + 0.5 # Shift to center of pixel + if pad: + padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device) + padding_label = -torch.ones((labels.shape[0], 1), device=labels.device) + points = torch.cat([points, padding_point], dim=1) + labels = torch.cat([labels, padding_label], dim=1) + point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size) + point_embedding[labels == -1] = 0.0 + point_embedding[labels == -1] += self.not_a_point_embed.weight + point_embedding[labels == 0] += self.point_embeddings[0].weight + point_embedding[labels == 1] += self.point_embeddings[1].weight + return point_embedding + + def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: + """Embeds box prompts.""" + boxes = boxes + 0.5 # Shift to center of pixel + coords = boxes.reshape(-1, 2, 2) + corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size) + corner_embedding[:, 0, :] += self.point_embeddings[2].weight + corner_embedding[:, 1, :] += self.point_embeddings[3].weight + return corner_embedding + + def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor: + """Embeds mask inputs.""" + mask_embedding = self.mask_downscaling(masks) + return mask_embedding + + def _get_batch_size( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> int: + """ + Gets the batch size of the output given the batch size of the input prompts. + """ + if points is not None: + return points[0].shape[0] + elif boxes is not None: + return boxes.shape[0] + elif masks is not None: + return masks.shape[0] + else: + return 1 + + def _get_device(self) -> torch.device: + return self.point_embeddings[0].weight.device + + def forward( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Embeds different types of prompts, returning both sparse and dense + embeddings. + + Arguments: + points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates + and labels to embed. + boxes (torch.Tensor or none): boxes to embed + masks (torch.Tensor or none): masks to embed + + Returns: + torch.Tensor: sparse embeddings for the points and boxes, with shape + BxNx(embed_dim), where N is determined by the number of input points + and boxes. + torch.Tensor: dense embeddings for the masks, in the shape + Bx(embed_dim)x(embed_H)x(embed_W) + """ + bs = self._get_batch_size(points, boxes, masks) + sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device()) + if points is not None: + coords, labels = points + point_embeddings = self._embed_points(coords, labels, pad=(boxes is None)) + sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1) + if boxes is not None: + box_embeddings = self._embed_boxes(boxes) + sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1) + + if masks is not None: + dense_embeddings = self._embed_masks(masks) + else: + dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( + bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] + ) + + return sparse_embeddings, dense_embeddings + + +class PositionEmbeddingRandom(nn.Module): + """ + Positional encoding using random spatial frequencies. + """ + + def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: + super().__init__() + if scale is None or scale <= 0.0: + scale = 1.0 + self.register_buffer( + "positional_encoding_gaussian_matrix", + scale * torch.randn((2, num_pos_feats)), + ) + + def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: + """Positionally encode points that are normalized to [0,1].""" + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coords = 2 * coords - 1 + coords = coords @ self.positional_encoding_gaussian_matrix + coords = 2 * np.pi * coords + # outputs d_1 x ... x d_n x C shape + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + def forward(self, size: Tuple[int, int]) -> torch.Tensor: + """Generate positional encoding for a grid of the specified size.""" + h, w = size + device: Any = self.positional_encoding_gaussian_matrix.device + grid = torch.ones((h, w), device=device, dtype=torch.float32) + y_embed = grid.cumsum(dim=0) - 0.5 + x_embed = grid.cumsum(dim=1) - 0.5 + y_embed = y_embed / h + x_embed = x_embed / w + + pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) + return pe.permute(2, 0, 1) # C x H x W + + def forward_with_coords( + self, coords_input: torch.Tensor, image_size: Tuple[int, int] + ) -> torch.Tensor: + """Positionally encode points that are not normalized to [0,1].""" + coords = coords_input.clone() + coords[:, :, 0] = coords[:, :, 0] / image_size[1] + coords[:, :, 1] = coords[:, :, 1] / image_size[0] + return self._pe_encoding(coords.to(torch.float)) # B x N x C diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/sam.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/sam.py new file mode 100644 index 0000000000000000000000000000000000000000..303bc2f40c3dbc84f5d4286bb73336e075a86589 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/sam.py @@ -0,0 +1,174 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import nn +from torch.nn import functional as F + +from typing import Any, Dict, List, Tuple + +from .image_encoder import ImageEncoderViT +from .mask_decoder import MaskDecoder +from .prompt_encoder import PromptEncoder + + +class Sam(nn.Module): + mask_threshold: float = 0.0 + image_format: str = "RGB" + + def __init__( + self, + image_encoder: ImageEncoderViT, + prompt_encoder: PromptEncoder, + mask_decoder: MaskDecoder, + pixel_mean: List[float] = [123.675, 116.28, 103.53], + pixel_std: List[float] = [58.395, 57.12, 57.375], + ) -> None: + """ + SAM predicts object masks from an image and input prompts. + + Arguments: + image_encoder (ImageEncoderViT): The backbone used to encode the + image into image embeddings that allow for efficient mask prediction. + prompt_encoder (PromptEncoder): Encodes various types of input prompts. + mask_decoder (MaskDecoder): Predicts masks from the image embeddings + and encoded prompts. + pixel_mean (list(float)): Mean values for normalizing pixels in the input image. + pixel_std (list(float)): Std values for normalizing pixels in the input image. + """ + super().__init__() + self.image_encoder = image_encoder + self.prompt_encoder = prompt_encoder + self.mask_decoder = mask_decoder + self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False) + self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False) + + @property + def device(self) -> Any: + return self.pixel_mean.device + + @torch.no_grad() + def forward( + self, + batched_input: List[Dict[str, Any]], + multimask_output: bool, + ) -> List[Dict[str, torch.Tensor]]: + """ + Predicts masks end-to-end from provided images and prompts. + If prompts are not known in advance, using SamPredictor is + recommended over calling the model directly. + + Arguments: + batched_input (list(dict)): A list over input images, each a + dictionary with the following keys. A prompt key can be + excluded if it is not present. + 'image': The image as a torch tensor in 3xHxW format, + already transformed for input to the model. + 'original_size': (tuple(int, int)) The original size of + the image before transformation, as (H, W). + 'point_coords': (torch.Tensor) Batched point prompts for + this image, with shape BxNx2. Already transformed to the + input frame of the model. + 'point_labels': (torch.Tensor) Batched labels for point prompts, + with shape BxN. + 'boxes': (torch.Tensor) Batched box inputs, with shape Bx4. + Already transformed to the input frame of the model. + 'mask_inputs': (torch.Tensor) Batched mask inputs to the model, + in the form Bx1xHxW. + multimask_output (bool): Whether the model should predict multiple + disambiguating masks, or return a single mask. + + Returns: + (list(dict)): A list over input images, where each element is + as dictionary with the following keys. + 'masks': (torch.Tensor) Batched binary mask predictions, + with shape BxCxHxW, where B is the number of input promts, + C is determiend by multimask_output, and (H, W) is the + original size of the image. + 'iou_predictions': (torch.Tensor) The model's predictions + of mask quality, in shape BxC. + 'low_res_logits': (torch.Tensor) Low resolution logits with + shape BxCxHxW, where H=W=256. Can be passed as mask input + to subsequent iterations of prediction. + """ + input_images = torch.stack([self.preprocess(x["image"]) for x in batched_input], dim=0) + image_embeddings = self.image_encoder(input_images) + + outputs = [] + for image_record, curr_embedding in zip(batched_input, image_embeddings): + if "point_coords" in image_record: + points = (image_record["point_coords"], image_record["point_labels"]) + else: + points = None + sparse_embeddings, dense_embeddings = self.prompt_encoder( + points=points, + boxes=image_record.get("boxes", None), + masks=image_record.get("mask_inputs", None), + ) + low_res_masks, iou_predictions = self.mask_decoder( + image_embeddings=curr_embedding.unsqueeze(0), + image_pe=self.prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + ) + masks = self.postprocess_masks( + low_res_masks, + input_size=image_record["image"].shape[-2:], + original_size=image_record["original_size"], + ) + masks = masks > self.mask_threshold + outputs.append( + { + "masks": masks, + "iou_predictions": iou_predictions, + "low_res_logits": low_res_masks, + } + ) + return outputs + + def postprocess_masks( + self, + masks: torch.Tensor, + input_size: Tuple[int, ...], + original_size: Tuple[int, ...], + ) -> torch.Tensor: + """ + Remove padding and upscale masks to the original image size. + + Arguments: + masks (torch.Tensor): Batched masks from the mask_decoder, + in BxCxHxW format. + input_size (tuple(int, int)): The size of the image input to the + model, in (H, W) format. Used to remove padding. + original_size (tuple(int, int)): The original size of the image + before resizing for input to the model, in (H, W) format. + + Returns: + (torch.Tensor): Batched masks in BxCxHxW format, where (H, W) + is given by original_size. + """ + masks = F.interpolate( + masks, + (self.image_encoder.img_size, self.image_encoder.img_size), + mode="bilinear", + align_corners=False, + ) + masks = masks[..., : input_size[0], : input_size[1]] + masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False) + return masks + + def preprocess(self, x: torch.Tensor) -> torch.Tensor: + """Normalize pixel values and pad to a square input.""" + # Normalize colors + x = (x - self.pixel_mean) / self.pixel_std + + # Pad + h, w = x.shape[-2:] + padh = self.image_encoder.img_size - h + padw = self.image_encoder.img_size - w + x = F.pad(x, (0, padw, 0, padh)) + return x diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/transformer.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f1a2812f613cc55b1d0b3e3e1d0c84a760d1fb87 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/modeling/transformer.py @@ -0,0 +1,240 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor, nn + +import math +from typing import Tuple, Type + +from .common import MLPBlock + + +class TwoWayTransformer(nn.Module): + def __init__( + self, + depth: int, + embedding_dim: int, + num_heads: int, + mlp_dim: int, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + ) -> None: + """ + A transformer decoder that attends to an input image using + queries whose positional embedding is supplied. + + Args: + depth (int): number of layers in the transformer + embedding_dim (int): the channel dimension for the input embeddings + num_heads (int): the number of heads for multihead attention. Must + divide embedding_dim + mlp_dim (int): the channel dimension internal to the MLP block + activation (nn.Module): the activation to use in the MLP block + """ + super().__init__() + self.depth = depth + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.layers = nn.ModuleList() + + for i in range(depth): + self.layers.append( + TwoWayAttentionBlock( + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_dim=mlp_dim, + activation=activation, + attention_downsample_rate=attention_downsample_rate, + skip_first_layer_pe=(i == 0), + ) + ) + + self.final_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm_final_attn = nn.LayerNorm(embedding_dim) + + def forward( + self, + image_embedding: Tensor, + image_pe: Tensor, + point_embedding: Tensor, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + image_embedding (torch.Tensor): image to attend to. Should be shape + B x embedding_dim x h x w for any h and w. + image_pe (torch.Tensor): the positional encoding to add to the image. Must + have the same shape as image_embedding. + point_embedding (torch.Tensor): the embedding to add to the query points. + Must have shape B x N_points x embedding_dim for any N_points. + + Returns: + torch.Tensor: the processed point_embedding + torch.Tensor: the processed image_embedding + """ + # BxCxHxW -> BxHWxC == B x N_image_tokens x C + bs, c, h, w = image_embedding.shape + image_embedding = image_embedding.flatten(2).permute(0, 2, 1) + image_pe = image_pe.flatten(2).permute(0, 2, 1) + + # Prepare queries + queries = point_embedding + keys = image_embedding + + # Apply transformer blocks and final layernorm + for layer in self.layers: + queries, keys = layer( + queries=queries, + keys=keys, + query_pe=point_embedding, + key_pe=image_pe, + ) + + # Apply the final attenion layer from the points to the image + q = queries + point_embedding + k = keys + image_pe + attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm_final_attn(queries) + + return queries, keys + + +class TwoWayAttentionBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + num_heads: int, + mlp_dim: int = 2048, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + skip_first_layer_pe: bool = False, + ) -> None: + """ + A transformer block with four layers: (1) self-attention of sparse + inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp + block on sparse inputs, and (4) cross attention of dense inputs to sparse + inputs. + + Arguments: + embedding_dim (int): the channel dimension of the embeddings + num_heads (int): the number of heads in the attention layers + mlp_dim (int): the hidden dimension of the mlp block + activation (nn.Module): the activation of the mlp block + skip_first_layer_pe (bool): skip the PE on the first layer + """ + super().__init__() + self.self_attn = Attention(embedding_dim, num_heads) + self.norm1 = nn.LayerNorm(embedding_dim) + + self.cross_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm2 = nn.LayerNorm(embedding_dim) + + self.mlp = MLPBlock(embedding_dim, mlp_dim, activation) + self.norm3 = nn.LayerNorm(embedding_dim) + + self.norm4 = nn.LayerNorm(embedding_dim) + self.cross_attn_image_to_token = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + + self.skip_first_layer_pe = skip_first_layer_pe + + def forward( + self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor + ) -> Tuple[Tensor, Tensor]: + # Self attention block + if self.skip_first_layer_pe: + queries = self.self_attn(q=queries, k=queries, v=queries) + else: + q = queries + query_pe + attn_out = self.self_attn(q=q, k=q, v=queries) + queries = queries + attn_out + queries = self.norm1(queries) + + # Cross attention block, tokens attending to image embedding + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm2(queries) + + # MLP block + mlp_out = self.mlp(queries) + queries = queries + mlp_out + queries = self.norm3(queries) + + # Cross attention block, image embedding attending to tokens + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) + keys = keys + attn_out + keys = self.norm4(keys) + + return queries, keys + + +class Attention(nn.Module): + """ + An attention layer that allows for downscaling the size of the embedding + after projection to queries, keys, and values. + """ + + def __init__( + self, + embedding_dim: int, + num_heads: int, + downsample_rate: int = 1, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.internal_dim = embedding_dim // downsample_rate + self.num_heads = num_heads + assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim." + + self.q_proj = nn.Linear(embedding_dim, self.internal_dim) + self.k_proj = nn.Linear(embedding_dim, self.internal_dim) + self.v_proj = nn.Linear(embedding_dim, self.internal_dim) + self.out_proj = nn.Linear(self.internal_dim, embedding_dim) + + def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: + b, n, c = x.shape + x = x.reshape(b, n, num_heads, c // num_heads) + return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head + + def _recombine_heads(self, x: Tensor) -> Tensor: + b, n_heads, n_tokens, c_per_head = x.shape + x = x.transpose(1, 2) + return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + # Attention + _, _, _, c_per_head = q.shape + attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens + attn = attn / math.sqrt(c_per_head) + attn = torch.softmax(attn, dim=-1) + + # Get output + out = attn @ v + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/predictor.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..882063e14c9eb346018cf7e0a668fce251fa18f6 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/predictor.py @@ -0,0 +1,276 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch + +from .modeling import Sam + +from typing import Optional, Tuple + +from .utils.transforms import ResizeLongestSide + + +class SamPredictor: + def __init__( + self, + sam_model: Sam, + ) -> None: + """ + Uses SAM to calculate the image embedding for an image, and then + allow repeated, efficient mask prediction given prompts. + + Arguments: + sam_model (Sam): The model to use for mask prediction. + """ + super().__init__() + self.model = sam_model + self.transform = ResizeLongestSide(sam_model.image_encoder.img_size) + self.reset_image() + + def set_image( + self, + image: np.ndarray, + image_format: str = "RGB", + ) -> None: + """ + Calculates the image embeddings for the provided image, allowing + masks to be predicted with the 'predict' method. + + Arguments: + image (np.ndarray): The image for calculating masks. Expects an + image in HWC uint8 format, with pixel values in [0, 255]. + image_format (str): The color format of the image, in ['RGB', 'BGR']. + """ + assert image_format in [ + "RGB", + "BGR", + ], f"image_format must be in ['RGB', 'BGR'], is {image_format}." + # import pdb;pdb.set_trace() + if image_format != self.model.image_format: + image = image[..., ::-1] + + # Transform the image to the form expected by the model + # import pdb;pdb.set_trace() + input_image = self.transform.apply_image(image) + input_image_torch = torch.as_tensor(input_image, device=self.device) + input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :] + + self.set_torch_image(input_image_torch, image.shape[:2]) + + @torch.no_grad() + def set_torch_image( + self, + transformed_image: torch.Tensor, + original_image_size: Tuple[int, ...], + ) -> None: + """ + Calculates the image embeddings for the provided image, allowing + masks to be predicted with the 'predict' method. Expects the input + image to be already transformed to the format expected by the model. + + Arguments: + transformed_image (torch.Tensor): The input image, with shape + 1x3xHxW, which has been transformed with ResizeLongestSide. + original_image_size (tuple(int, int)): The size of the image + before transformation, in (H, W) format. + """ + assert ( + len(transformed_image.shape) == 4 + and transformed_image.shape[1] == 3 + and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size + ), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}." + self.reset_image() + + self.original_size = original_image_size + self.input_size = tuple(transformed_image.shape[-2:]) + input_image = self.model.preprocess(transformed_image) + self.features, self.interm_features = self.model.image_encoder(input_image) + self.is_image_set = True + + def predict( + self, + point_coords: Optional[np.ndarray] = None, + point_labels: Optional[np.ndarray] = None, + box: Optional[np.ndarray] = None, + mask_input: Optional[np.ndarray] = None, + multimask_output: bool = True, + return_logits: bool = False, + hq_token_only: bool =False, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Predict masks for the given input prompts, using the currently set image. + + Arguments: + point_coords (np.ndarray or None): A Nx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (np.ndarray or None): A length N array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + box (np.ndarray or None): A length 4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form 1xHxW, where + for SAM, H=W=256. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + + Returns: + (np.ndarray): The output masks in CxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (np.ndarray): An array of length C containing the model's + predictions for the quality of each mask. + (np.ndarray): An array of shape CxHxW, where C is the number + of masks and H=W=256. These low resolution logits can be passed to + a subsequent iteration as mask input. + """ + if not self.is_image_set: + raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") + + # Transform input prompts + coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None + if point_coords is not None: + assert ( + point_labels is not None + ), "point_labels must be supplied if point_coords is supplied." + point_coords = self.transform.apply_coords(point_coords, self.original_size) + coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device) + labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device) + coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :] + if box is not None: + box = self.transform.apply_boxes(box, self.original_size) + box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device) + box_torch = box_torch[None, :] + if mask_input is not None: + mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.device) + mask_input_torch = mask_input_torch[None, :, :, :] + + masks, iou_predictions, low_res_masks = self.predict_torch( + coords_torch, + labels_torch, + box_torch, + mask_input_torch, + multimask_output, + return_logits=return_logits, + hq_token_only=hq_token_only, + ) + + masks_np = masks[0].detach().cpu().numpy() + iou_predictions_np = iou_predictions[0].detach().cpu().numpy() + low_res_masks_np = low_res_masks[0].detach().cpu().numpy() + return masks_np, iou_predictions_np, low_res_masks_np + + @torch.no_grad() + def predict_torch( + self, + point_coords: Optional[torch.Tensor], + point_labels: Optional[torch.Tensor], + boxes: Optional[torch.Tensor] = None, + mask_input: Optional[torch.Tensor] = None, + multimask_output: bool = True, + return_logits: bool = False, + hq_token_only: bool =False, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Predict masks for the given input prompts, using the currently set image. + Input prompts are batched torch tensors and are expected to already be + transformed to the input frame using ResizeLongestSide. + + Arguments: + point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (torch.Tensor or None): A BxN array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + boxes (np.ndarray or None): A Bx4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form Bx1xHxW, where + for SAM, H=W=256. Masks returned by a previous iteration of the + predict method do not need further transformation. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + + Returns: + (torch.Tensor): The output masks in BxCxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (torch.Tensor): An array of shape BxC containing the model's + predictions for the quality of each mask. + (torch.Tensor): An array of shape BxCxHxW, where C is the number + of masks and H=W=256. These low res logits can be passed to + a subsequent iteration as mask input. + """ + if not self.is_image_set: + raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") + + if point_coords is not None: + points = (point_coords, point_labels) + else: + points = None + + # Embed prompts + sparse_embeddings, dense_embeddings = self.model.prompt_encoder( + points=points, + boxes=boxes, + masks=mask_input, + ) + + # Predict masks + low_res_masks, iou_predictions = self.model.mask_decoder( + image_embeddings=self.features, + image_pe=self.model.prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + hq_token_only=hq_token_only, + interm_embeddings=self.interm_features, + ) + + # Upscale the masks to the original image resolution + masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size) + + if not return_logits: + masks = masks > self.model.mask_threshold + + return masks, iou_predictions, low_res_masks + + def get_image_embedding(self) -> torch.Tensor: + """ + Returns the image embeddings for the currently set image, with + shape 1xCxHxW, where C is the embedding dimension and (H,W) are + the embedding spatial dimension of SAM (typically C=256, H=W=64). + """ + if not self.is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) to generate an embedding." + ) + assert self.features is not None, "Features must exist if an image has been set." + return self.features + + @property + def device(self) -> torch.device: + return self.model.device + + def reset_image(self) -> None: + """Resets the currently set image.""" + self.is_image_set = False + self.features = None + self.orig_h = None + self.orig_w = None + self.input_h = None + self.input_w = None \ No newline at end of file diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__init__.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5277f46157403e47fd830fc519144b97ef69d4ae --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/__init__.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0efd72923b79790d6fb0b7bcea89a4ca0b16f00a Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/__init__.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/amg.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/amg.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a78cb163484c2b015feb56b61fd8bb605170cb8e Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/amg.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/transforms.cpython-310.pyc b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/transforms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b698b4f636f09cafb4c30d46be407acca8270be Binary files /dev/null and b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/__pycache__/transforms.cpython-310.pyc differ diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/amg.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/amg.py new file mode 100644 index 0000000000000000000000000000000000000000..3a137778e45c464c079658ecb87ec53270e789f7 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/amg.py @@ -0,0 +1,346 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch + +import math +from copy import deepcopy +from itertools import product +from typing import Any, Dict, Generator, ItemsView, List, Tuple + + +class MaskData: + """ + A structure for storing masks and their related data in batched format. + Implements basic filtering and concatenation. + """ + + def __init__(self, **kwargs) -> None: + for v in kwargs.values(): + assert isinstance( + v, (list, np.ndarray, torch.Tensor) + ), "MaskData only supports list, numpy arrays, and torch tensors." + self._stats = dict(**kwargs) + + def __setitem__(self, key: str, item: Any) -> None: + assert isinstance( + item, (list, np.ndarray, torch.Tensor) + ), "MaskData only supports list, numpy arrays, and torch tensors." + self._stats[key] = item + + def __delitem__(self, key: str) -> None: + del self._stats[key] + + def __getitem__(self, key: str) -> Any: + return self._stats[key] + + def items(self) -> ItemsView[str, Any]: + return self._stats.items() + + def filter(self, keep: torch.Tensor) -> None: + for k, v in self._stats.items(): + if v is None: + self._stats[k] = None + elif isinstance(v, torch.Tensor): + self._stats[k] = v[torch.as_tensor(keep, device=v.device)] + elif isinstance(v, np.ndarray): + self._stats[k] = v[keep.detach().cpu().numpy()] + elif isinstance(v, list) and keep.dtype == torch.bool: + self._stats[k] = [a for i, a in enumerate(v) if keep[i]] + elif isinstance(v, list): + self._stats[k] = [v[i] for i in keep] + else: + raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") + + def cat(self, new_stats: "MaskData") -> None: + for k, v in new_stats.items(): + if k not in self._stats or self._stats[k] is None: + self._stats[k] = deepcopy(v) + elif isinstance(v, torch.Tensor): + self._stats[k] = torch.cat([self._stats[k], v], dim=0) + elif isinstance(v, np.ndarray): + self._stats[k] = np.concatenate([self._stats[k], v], axis=0) + elif isinstance(v, list): + self._stats[k] = self._stats[k] + deepcopy(v) + else: + raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") + + def to_numpy(self) -> None: + for k, v in self._stats.items(): + if isinstance(v, torch.Tensor): + self._stats[k] = v.detach().cpu().numpy() + + +def is_box_near_crop_edge( + boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0 +) -> torch.Tensor: + """Filter masks at the edge of a crop, but not at the edge of the original image.""" + crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device) + orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device) + boxes = uncrop_boxes_xyxy(boxes, crop_box).float() + near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0) + near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0) + near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge) + return torch.any(near_crop_edge, dim=1) + + +def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor: + box_xywh = deepcopy(box_xyxy) + box_xywh[2] = box_xywh[2] - box_xywh[0] + box_xywh[3] = box_xywh[3] - box_xywh[1] + return box_xywh + + +def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]: + assert len(args) > 0 and all( + len(a) == len(args[0]) for a in args + ), "Batched iteration must have inputs of all the same size." + n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0) + for b in range(n_batches): + yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args] + + +def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]: + """ + Encodes masks to an uncompressed RLE, in the format expected by + pycoco tools. + """ + # Put in fortran order and flatten h,w + b, h, w = tensor.shape + tensor = tensor.permute(0, 2, 1).flatten(1) + + # Compute change indices + diff = tensor[:, 1:] ^ tensor[:, :-1] + change_indices = diff.nonzero() + + # Encode run length + out = [] + for i in range(b): + cur_idxs = change_indices[change_indices[:, 0] == i, 1] + cur_idxs = torch.cat( + [ + torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device), + cur_idxs + 1, + torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device), + ] + ) + btw_idxs = cur_idxs[1:] - cur_idxs[:-1] + counts = [] if tensor[i, 0] == 0 else [0] + counts.extend(btw_idxs.detach().cpu().tolist()) + out.append({"size": [h, w], "counts": counts}) + return out + + +def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray: + """Compute a binary mask from an uncompressed RLE.""" + h, w = rle["size"] + mask = np.empty(h * w, dtype=bool) + idx = 0 + parity = False + for count in rle["counts"]: + mask[idx : idx + count] = parity + idx += count + parity ^= True + mask = mask.reshape(w, h) + return mask.transpose() # Put in C order + + +def area_from_rle(rle: Dict[str, Any]) -> int: + return sum(rle["counts"][1::2]) + + +def calculate_stability_score( + masks: torch.Tensor, mask_threshold: float, threshold_offset: float +) -> torch.Tensor: + """ + Computes the stability score for a batch of masks. The stability + score is the IoU between the binary masks obtained by thresholding + the predicted mask logits at high and low values. + """ + # One mask is always contained inside the other. + # Save memory by preventing unnecesary cast to torch.int64 + intersections = ( + (masks > (mask_threshold + threshold_offset)) + .sum(-1, dtype=torch.int16) + .sum(-1, dtype=torch.int32) + ) + unions = ( + (masks > (mask_threshold - threshold_offset)) + .sum(-1, dtype=torch.int16) + .sum(-1, dtype=torch.int32) + ) + return intersections / unions + + +def build_point_grid(n_per_side: int) -> np.ndarray: + """Generates a 2D grid of points evenly spaced in [0,1]x[0,1].""" + offset = 1 / (2 * n_per_side) + points_one_side = np.linspace(offset, 1 - offset, n_per_side) + points_x = np.tile(points_one_side[None, :], (n_per_side, 1)) + points_y = np.tile(points_one_side[:, None], (1, n_per_side)) + points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2) + return points + + +def build_all_layer_point_grids( + n_per_side: int, n_layers: int, scale_per_layer: int +) -> List[np.ndarray]: + """Generates point grids for all crop layers.""" + points_by_layer = [] + for i in range(n_layers + 1): + n_points = int(n_per_side / (scale_per_layer**i)) + points_by_layer.append(build_point_grid(n_points)) + return points_by_layer + + +def generate_crop_boxes( + im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float +) -> Tuple[List[List[int]], List[int]]: + """ + Generates a list of crop boxes of different sizes. Each layer + has (2**i)**2 boxes for the ith layer. + """ + crop_boxes, layer_idxs = [], [] + im_h, im_w = im_size + short_side = min(im_h, im_w) + + # Original image + crop_boxes.append([0, 0, im_w, im_h]) + layer_idxs.append(0) + + def crop_len(orig_len, n_crops, overlap): + return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops)) + + for i_layer in range(n_layers): + n_crops_per_side = 2 ** (i_layer + 1) + overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side)) + + crop_w = crop_len(im_w, n_crops_per_side, overlap) + crop_h = crop_len(im_h, n_crops_per_side, overlap) + + crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)] + crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)] + + # Crops in XYWH format + for x0, y0 in product(crop_box_x0, crop_box_y0): + box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)] + crop_boxes.append(box) + layer_idxs.append(i_layer + 1) + + return crop_boxes, layer_idxs + + +def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor: + x0, y0, _, _ = crop_box + offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device) + # Check if boxes has a channel dimension + if len(boxes.shape) == 3: + offset = offset.unsqueeze(1) + return boxes + offset + + +def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor: + x0, y0, _, _ = crop_box + offset = torch.tensor([[x0, y0]], device=points.device) + # Check if points has a channel dimension + if len(points.shape) == 3: + offset = offset.unsqueeze(1) + return points + offset + + +def uncrop_masks( + masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int +) -> torch.Tensor: + x0, y0, x1, y1 = crop_box + if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h: + return masks + # Coordinate transform masks + pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0) + pad = (x0, pad_x - x0, y0, pad_y - y0) + return torch.nn.functional.pad(masks, pad, value=0) + + +def remove_small_regions( + mask: np.ndarray, area_thresh: float, mode: str +) -> Tuple[np.ndarray, bool]: + """ + Removes small disconnected regions and holes in a mask. Returns the + mask and an indicator of if the mask has been modified. + """ + import cv2 # type: ignore + + assert mode in ["holes", "islands"] + correct_holes = mode == "holes" + working_mask = (correct_holes ^ mask).astype(np.uint8) + n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8) + sizes = stats[:, -1][1:] # Row 0 is background label + small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh] + if len(small_regions) == 0: + return mask, False + fill_labels = [0] + small_regions + if not correct_holes: + fill_labels = [i for i in range(n_labels) if i not in fill_labels] + # If every region is below threshold, keep largest + if len(fill_labels) == 0: + fill_labels = [int(np.argmax(sizes)) + 1] + mask = np.isin(regions, fill_labels) + return mask, True + + +def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]: + from pycocotools import mask as mask_utils # type: ignore + + h, w = uncompressed_rle["size"] + rle = mask_utils.frPyObjects(uncompressed_rle, h, w) + rle["counts"] = rle["counts"].decode("utf-8") # Necessary to serialize with json + return rle + + +def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor: + """ + Calculates boxes in XYXY format around masks. Return [0,0,0,0] for + an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4. + """ + # torch.max below raises an error on empty inputs, just skip in this case + if torch.numel(masks) == 0: + return torch.zeros(*masks.shape[:-2], 4, device=masks.device) + + # Normalize shape to CxHxW + shape = masks.shape + h, w = shape[-2:] + if len(shape) > 2: + masks = masks.flatten(0, -3) + else: + masks = masks.unsqueeze(0) + + # Get top and bottom edges + in_height, _ = torch.max(masks, dim=-1) + in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :] + bottom_edges, _ = torch.max(in_height_coords, dim=-1) + in_height_coords = in_height_coords + h * (~in_height) + top_edges, _ = torch.min(in_height_coords, dim=-1) + + # Get left and right edges + in_width, _ = torch.max(masks, dim=-2) + in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :] + right_edges, _ = torch.max(in_width_coords, dim=-1) + in_width_coords = in_width_coords + w * (~in_width) + left_edges, _ = torch.min(in_width_coords, dim=-1) + + # If the mask is empty the right edge will be to the left of the left edge. + # Replace these boxes with [0, 0, 0, 0] + empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges) + out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1) + out = out * (~empty_filter).unsqueeze(-1) + + # Return to original shape + if len(shape) > 2: + out = out.reshape(*shape[:-2], 4) + else: + out = out[0] + + return out diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/onnx.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/onnx.py new file mode 100644 index 0000000000000000000000000000000000000000..4297b31291e036700d6ad0b818afb7dd72da3054 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/onnx.py @@ -0,0 +1,144 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +from torch.nn import functional as F + +from typing import Tuple + +from ..modeling import Sam +from .amg import calculate_stability_score + + +class SamOnnxModel(nn.Module): + """ + This model should not be called directly, but is used in ONNX export. + It combines the prompt encoder, mask decoder, and mask postprocessing of Sam, + with some functions modified to enable model tracing. Also supports extra + options controlling what information. See the ONNX export script for details. + """ + + def __init__( + self, + model: Sam, + return_single_mask: bool, + use_stability_score: bool = False, + return_extra_metrics: bool = False, + ) -> None: + super().__init__() + self.mask_decoder = model.mask_decoder + self.model = model + self.img_size = model.image_encoder.img_size + self.return_single_mask = return_single_mask + self.use_stability_score = use_stability_score + self.stability_score_offset = 1.0 + self.return_extra_metrics = return_extra_metrics + + @staticmethod + def resize_longest_image_size( + input_image_size: torch.Tensor, longest_side: int + ) -> torch.Tensor: + input_image_size = input_image_size.to(torch.float32) + scale = longest_side / torch.max(input_image_size) + transformed_size = scale * input_image_size + transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64) + return transformed_size + + def _embed_points(self, point_coords: torch.Tensor, point_labels: torch.Tensor) -> torch.Tensor: + point_coords = point_coords + 0.5 + point_coords = point_coords / self.img_size + point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords) + point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding) + + point_embedding = point_embedding * (point_labels != -1) + point_embedding = point_embedding + self.model.prompt_encoder.not_a_point_embed.weight * ( + point_labels == -1 + ) + + for i in range(self.model.prompt_encoder.num_point_embeddings): + point_embedding = point_embedding + self.model.prompt_encoder.point_embeddings[ + i + ].weight * (point_labels == i) + + return point_embedding + + def _embed_masks(self, input_mask: torch.Tensor, has_mask_input: torch.Tensor) -> torch.Tensor: + mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(input_mask) + mask_embedding = mask_embedding + ( + 1 - has_mask_input + ) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1) + return mask_embedding + + def mask_postprocessing(self, masks: torch.Tensor, orig_im_size: torch.Tensor) -> torch.Tensor: + masks = F.interpolate( + masks, + size=(self.img_size, self.img_size), + mode="bilinear", + align_corners=False, + ) + + prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size) + masks = masks[..., : int(prepadded_size[0]), : int(prepadded_size[1])] + + orig_im_size = orig_im_size.to(torch.int64) + h, w = orig_im_size[0], orig_im_size[1] + masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False) + return masks + + def select_masks( + self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Determine if we should return the multiclick mask or not from the number of points. + # The reweighting is used to avoid control flow. + score_reweight = torch.tensor( + [[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)] + ).to(iou_preds.device) + score = iou_preds + (num_points - 2.5) * score_reweight + best_idx = torch.argmax(score, dim=1) + masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1) + iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1) + + return masks, iou_preds + + @torch.no_grad() + def forward( + self, + image_embeddings: torch.Tensor, + point_coords: torch.Tensor, + point_labels: torch.Tensor, + mask_input: torch.Tensor, + has_mask_input: torch.Tensor, + orig_im_size: torch.Tensor, + ): + sparse_embedding = self._embed_points(point_coords, point_labels) + dense_embedding = self._embed_masks(mask_input, has_mask_input) + + masks, scores = self.model.mask_decoder.predict_masks( + image_embeddings=image_embeddings, + image_pe=self.model.prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embedding, + dense_prompt_embeddings=dense_embedding, + ) + + if self.use_stability_score: + scores = calculate_stability_score( + masks, self.model.mask_threshold, self.stability_score_offset + ) + + if self.return_single_mask: + masks, scores = self.select_masks(masks, scores, point_coords.shape[1]) + + upscaled_masks = self.mask_postprocessing(masks, orig_im_size) + + if self.return_extra_metrics: + stability_scores = calculate_stability_score( + upscaled_masks, self.model.mask_threshold, self.stability_score_offset + ) + areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1) + return upscaled_masks, scores, stability_scores, areas, masks + + return upscaled_masks, scores, masks diff --git a/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/transforms.py b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..3ad346661f84b0647026e130a552c4b38b83e2ac --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/segment_anything/utils/transforms.py @@ -0,0 +1,102 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torch.nn import functional as F +from torchvision.transforms.functional import resize, to_pil_image # type: ignore + +from copy import deepcopy +from typing import Tuple + + +class ResizeLongestSide: + """ + Resizes images to longest side 'target_length', as well as provides + methods for resizing coordinates and boxes. Provides methods for + transforming both numpy array and batched torch tensors. + """ + + def __init__(self, target_length: int) -> None: + self.target_length = target_length + + def apply_image(self, image: np.ndarray) -> np.ndarray: + """ + Expects a numpy array with shape HxWxC in uint8 format. + """ + target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length) + return np.array(resize(to_pil_image(image), target_size)) + + def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray: + """ + Expects a numpy array of length 2 in the final dimension. Requires the + original image size in (H, W) format. + """ + old_h, old_w = original_size + new_h, new_w = self.get_preprocess_shape( + original_size[0], original_size[1], self.target_length + ) + coords = deepcopy(coords).astype(float) + coords[..., 0] = coords[..., 0] * (new_w / old_w) + coords[..., 1] = coords[..., 1] * (new_h / old_h) + return coords + + def apply_boxes(self, boxes: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray: + """ + Expects a numpy array shape Bx4. Requires the original image size + in (H, W) format. + """ + boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size) + return boxes.reshape(-1, 4) + + def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor: + """ + Expects batched images with shape BxCxHxW and float format. This + transformation may not exactly match apply_image. apply_image is + the transformation expected by the model. + """ + # Expects an image in BCHW format. May not exactly match apply_image. + target_size = self.get_preprocess_shape(image.shape[0], image.shape[1], self.target_length) + return F.interpolate( + image, target_size, mode="bilinear", align_corners=False, antialias=True + ) + + def apply_coords_torch( + self, coords: torch.Tensor, original_size: Tuple[int, ...] + ) -> torch.Tensor: + """ + Expects a torch tensor with length 2 in the last dimension. Requires the + original image size in (H, W) format. + """ + old_h, old_w = original_size + new_h, new_w = self.get_preprocess_shape( + original_size[0], original_size[1], self.target_length + ) + coords = deepcopy(coords).to(torch.float) + coords[..., 0] = coords[..., 0] * (new_w / old_w) + coords[..., 1] = coords[..., 1] * (new_h / old_h) + return coords + + def apply_boxes_torch( + self, boxes: torch.Tensor, original_size: Tuple[int, ...] + ) -> torch.Tensor: + """ + Expects a torch tensor with shape Bx4. Requires the original image + size in (H, W) format. + """ + boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size) + return boxes.reshape(-1, 4) + + @staticmethod + def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]: + """ + Compute the output size given input size and target long side length. + """ + scale = long_side_length * 1.0 / max(oldh, oldw) + newh, neww = oldh * scale, oldw * scale + neww = int(neww + 0.5) + newh = int(newh + 0.5) + return (newh, neww) diff --git a/ArtiAgent - DefectFill/src/segment_anything/setup.cfg b/ArtiAgent - DefectFill/src/segment_anything/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..0eee130ba71d14ec260d33a8ebd96a6491079a54 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/setup.cfg @@ -0,0 +1,11 @@ +[isort] +line_length=100 +multi_line_output=3 +include_trailing_comma=True +known_standard_library=numpy,setuptools +skip_glob=*/__init__.py +known_myself=segment_anything +known_third_party=matplotlib,cv2,torch,torchvision,pycocotools,onnx,black,isort +no_lines_before=STDLIB,THIRDPARTY +sections=FUTURE,STDLIB,THIRDPARTY,MYSELF,FIRSTPARTY,LOCALFOLDER +default_section=FIRSTPARTY diff --git a/ArtiAgent - DefectFill/src/segment_anything/setup.py b/ArtiAgent - DefectFill/src/segment_anything/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0986317eb576a14ec774205c88fdee3cc6c0b3 --- /dev/null +++ b/ArtiAgent - DefectFill/src/segment_anything/setup.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from setuptools import find_packages, setup + +setup( + name="segment_anything", + version="1.0", + install_requires=[], + packages=find_packages(exclude="notebooks"), + extras_require={ + "all": ["matplotlib", "pycocotools", "opencv-python", "onnx", "onnxruntime"], + "dev": ["flake8", "isort", "black", "mypy"], + }, +) diff --git a/ArtiAgent - DefectFill/src/utils.py b/ArtiAgent - DefectFill/src/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cee815f375a99785bd044c566882b889834b82ba --- /dev/null +++ b/ArtiAgent - DefectFill/src/utils.py @@ -0,0 +1,182 @@ +import torch +import torch.nn.functional as F +import os +from typing import Optional, Dict, Any + +def save_checkpoint(model, optimizer, step, path): + """ + Save model checkpoint - Includes LoRA weights and learnable embeddings (Textual Inversion) + + Args: + model: The DefectFill model instance + optimizer: Optimizer state + step: Current training step + path: File path to save the checkpoint + """ + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(path), exist_ok=True) + + # Extract LoRA weights specifically from the UNet and Text Encoder + checkpoint = { + "step": step, + "text_encoder_lora": {k: v for k, v in model.pipeline.text_encoder.state_dict().items() if "lora" in k}, + "unet_lora": {k: v for k, v in model.pipeline.unet.state_dict().items() if "lora" in k}, + "optimizer": optimizer.state_dict() if optimizer is not None else None, + } + + # Save the learnable embedding (Textual Inversion component) + if hasattr(model, 'placeholder_token_id'): + token_embeds = model.pipeline.text_encoder.get_input_embeddings().weight.data + checkpoint["learned_embedding"] = token_embeds[model.placeholder_token_id].clone() + checkpoint["placeholder_token"] = model.placeholder_token + checkpoint["placeholder_token_id"] = model.placeholder_token_id + print(f"[Checkpoint] Saving learnable embedding: {model.placeholder_token} (id={model.placeholder_token_id})") + + torch.save(checkpoint, path) + print(f"Checkpoint saved to {path}") + +def load_checkpoint(model, optimizer, path) -> int: + """ + Load model checkpoint - Restores LoRA weights and learnable embeddings + + Args: + model: The DefectFill model instance + optimizer: Optimizer to load state into + path: Path to the checkpoint file + + Returns: + Current step retrieved from the checkpoint + """ + # Check if checkpoint exists + if not os.path.exists(path): + print(f"Checkpoint {path} not found, starting from scratch") + return 0 + + # Load checkpoint to CPU first to avoid VRAM spikes + checkpoint = torch.load(path, map_location='cpu') + + # Load text encoder LoRA weights + text_encoder_sd = model.pipeline.text_encoder.state_dict() + for k, v in checkpoint["text_encoder_lora"].items(): + if k in text_encoder_sd: + text_encoder_sd[k] = v.to(text_encoder_sd[k].device) + model.pipeline.text_encoder.load_state_dict(text_encoder_sd) + + # Load UNet LoRA weights + unet_sd = model.pipeline.unet.state_dict() + for k, v in checkpoint["unet_lora"].items(): + if k in unet_sd: + unet_sd[k] = v.to(unet_sd[k].device) + model.pipeline.unet.load_state_dict(unet_sd) + + # Load the learnable embedding (Textual Inversion) + if "learned_embedding" in checkpoint and hasattr(model, 'placeholder_token_id'): + learned_emb = checkpoint["learned_embedding"] + token_embeds = model.pipeline.text_encoder.get_input_embeddings().weight.data + token_embeds[model.placeholder_token_id] = learned_emb.to(token_embeds.device) + print(f"[Checkpoint] Loaded learnable embedding: {model.placeholder_token} (id={model.placeholder_token_id})") + + # Load optimizer state if provided + if optimizer is not None and "optimizer" in checkpoint: + optimizer.load_state_dict(checkpoint["optimizer"]) + + print(f"Checkpoint loaded from {path}") + return checkpoint["step"] + + +def compute_spatial_lpips(lpips_model, img1, img2, mask, smooth_boundary=True): + """ + Calculates the Perceptual Distance specifically within the masked region using Spatial LPIPS + + Args: + lpips_model: LPIPS model instance initialized with spatial=True + img1: Reference image [B, 3, H, W], range [-1, 1] + img2: Comparison image [B, 3, H, W], range [-1, 1] + mask: Defect mask [B, 1, H, W], range [0, 1] + smooth_boundary: Whether to blur the mask edges to avoid boundary artifacts + + Returns: + lpips_score: LPIPS score for the masked region (scalar) + """ + # 1. Compute spatial LPIPS map (pixel-wise perceptual distance) + lpips_map = lpips_model(img1, img2) # Output shape: [B, 1, H', W'] + + # 2. Resize the mask to match the LPIPS output resolution + mask_resized = F.interpolate( + mask, + size=lpips_map.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # 3. Optional: Smooth boundary edges (Gaussian-like blur via AvgPool) + if smooth_boundary: + mask_smoothed = F.avg_pool2d( + F.pad(mask_resized, (2, 2, 2, 2), mode='replicate'), + kernel_size=5, stride=1 + ) + else: + mask_smoothed = mask_resized + + # 4. Mask-weighted summation + weighted_sum = (lpips_map * mask_smoothed).sum(dim=(2, 3)) + mask_sum = mask_smoothed.sum(dim=(2, 3)) + 1e-8 + + # 5. Return normalized LPIPS score + return (weighted_sum / mask_sum).mean() + + +def compute_spatial_lpips_batch(lpips_model, reference, samples, mask, smooth_boundary=True): + """ + Batch calculation of Spatial LPIPS - Evaluates all generated samples at once (FP16 compatible) + + This utilizes the "paired batch" mode of LPIPS by expanding the reference image + to match the number of samples, allowing parallel evaluation in a single forward pass. + + Args: + lpips_model: Spatial LPIPS model instance (spatial=True) + reference: Single reference image [1, 3, H, W], range [-1, 1] + samples: Multiple generated samples [N, 3, H, W], range [-1, 1] + mask: Single defect mask [1, 1, H, W], range [0, 1] + smooth_boundary: Enable mask boundary smoothing + + Returns: + lpips_scores: A tensor of [N] LPIPS scores + """ + num_samples = samples.shape[0] + + # LPIPS model expects FP32 input; cast based on model parameters + lpips_dtype = next(lpips_model.parameters()).dtype + + # Expand reference and mask to match sample count N + reference_expanded = reference.repeat(num_samples, 1, 1, 1).to(dtype=lpips_dtype) + samples_for_lpips = samples.to(dtype=lpips_dtype) + mask_expanded = mask.repeat(num_samples, 1, 1, 1) + + # Parallel forward pass for all N pairs + lpips_maps = lpips_model(reference_expanded, samples_for_lpips) # [N, 1, H', W'] + + # Match mask size to LPIPS feature map resolution + mask_resized = F.interpolate( + mask_expanded.to(dtype=lpips_maps.dtype), + size=lpips_maps.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # Boundary smoothing + if smooth_boundary: + mask_smoothed = F.avg_pool2d( + F.pad(mask_resized, (2, 2, 2, 2), mode='replicate'), + kernel_size=5, stride=1 + ) + else: + mask_smoothed = mask_resized + + # Calculate weighted scores per sample + weighted_sum = (lpips_maps * mask_smoothed).sum(dim=(2, 3)) # [N, 1] + mask_sum = mask_smoothed.sum(dim=(2, 3)) + 1e-8 # [N, 1] + + lpips_scores = (weighted_sum / mask_sum).squeeze(1) # [N] + + return lpips_scores \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..4584eb21010e8b5cbe02be0d4f09b92a09597a50 --- /dev/null +++ b/README.md @@ -0,0 +1,10 @@ +--- +title: Agentic Defect Synthesis +emoji: ๐Ÿ‘€ +colorFrom: pink +colorTo: green +sdk: static +pinned: false +--- + +Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..3d7e14915ce064bdf3a739adbc5d9ee230aecffd --- /dev/null +++ b/app.py @@ -0,0 +1,1348 @@ +""" +DefectFill & DefectDiffu Unified Web UI - Flask Backend +========================================================= +Supports both DefectFill (few-shot LoRA training) and DefectDiffu (text-guided generation). + +Run from project root: + python defectfill_webui/app.py + +Requirements: + pip install flask flask-cors pillow numpy +""" + +import os +import sys +import json +import time +import shutil +import threading +import subprocess +import traceback +from pathlib import Path +from datetime import datetime +from PIL import Image, ImageDraw +import numpy as np +import uuid +import io +import re + + +# ============================================================================= +# Configuration +# ============================================================================= + +BASE_DIR = Path(__file__).parent +PROJECT_ROOT = BASE_DIR.parent # ArtiAgent_with_WebUI/ + +# DefectFill paths +DEFECTFILL_ROOT = PROJECT_ROOT / "ArtiAgent - DefectFill" +TRAIN_PY = DEFECTFILL_ROOT / 'engine' / 'DefectFill' / 'train.py' +TRAINING_DATA_DIR = DEFECTFILL_ROOT / 'src' / 'data' +CHECKPOINT_DIR = DEFECTFILL_ROOT / 'engine' / 'DefectFill' / 'checkpoints' +TRAINED_GALLERY_DIR = DEFECTFILL_ROOT / "storage" / "trained_products_gallery" +TRAINED_GALLERY_DIR.mkdir(parents=True, exist_ok=True) +TEMP_IMAGE_DIR = DEFECTFILL_ROOT / 'src' / 'data' / 'tempImage' +TEMP_IMAGE_DIR.mkdir(parents=True, exist_ok=True) + +# DefectDiffu paths +DEFECTDIFFU_ROOT = PROJECT_ROOT / "ArtiAgent - DefectDiffu" +DEFECTDIFFU_ENGINE = PROJECT_ROOT / "engine" / "DefectDiffu" + +# Web UI dirs +UPLOAD_FOLDER = BASE_DIR / "uploads" +OUTPUT_DIR = BASE_DIR / "outputs" +CONFIG_DIR = BASE_DIR / "config" + +# OpenCV availability check +try: + import cv2 + CV2_AVAILABLE = True +except ImportError: + CV2_AVAILABLE = False + print("[WARNING] opencv-python not installed. Image alignment unavailable.") + +# Ensure directories exist +for d in [UPLOAD_FOLDER, OUTPUT_DIR, CONFIG_DIR]: + d.mkdir(parents=True, exist_ok=True) + +# Domain instructions config +DOMAIN_CONFIG_PATH = CONFIG_DIR / "domain_instructions.json" +if not DOMAIN_CONFIG_PATH.exists(): + with open(DOMAIN_CONFIG_PATH, "w") as f: + json.dump({"examples": {}, "active": {}}, f, indent=2) + +# DefectDiffu config +DEFECTDIFFU_CONFIG_PATH = CONFIG_DIR / "defectdiffu_config.json" +if not DEFECTDIFFU_CONFIG_PATH.exists(): + with open(DEFECTDIFFU_CONFIG_PATH, "w") as f: + json.dump({"ckpt_path": "", "vae_path": "", "vlm_model": "gemma3:12b"}, f, indent=2) + +# Setup sys.path for imports +sys.path.insert(0, str(DEFECTFILL_ROOT / 'src' / 'segment_anything')) +sys.path.insert(0, str(DEFECTFILL_ROOT / 'src')) +sys.path.insert(0, str(DEFECTFILL_ROOT / 'pipeline')) +sys.path.insert(0, str(DEFECTFILL_ROOT)) +sys.path.insert(0, str(DEFECTDIFFU_ROOT)) +sys.path.insert(0, str(DEFECTDIFFU_ROOT / 'pipeline')) + +# DefectFill imports +try: + from artiagent_orchestrator import ArtiAgentOrchestrator as DefectFillOrchestrator + from align_image_to_reference import align_image_to_reference + from detect_similar_product import check_product_training_status +except Exception as e: + print(f"[WARNING] DefectFill imports failed: {e}") + DefectFillOrchestrator = None + align_image_to_reference = None + check_product_training_status = None + +# DefectDiffu imports +try: + from artiagent_orchestrator import ArtiAgentOrchestrator as DefectDiffuOrchestrator +except Exception as e: + print(f"[WARNING] DefectDiffu imports failed: {e}") + DefectDiffuOrchestrator = None + +from flask import Flask, render_template, request, jsonify, send_from_directory +from flask_cors import CORS + +app = Flask(__name__, template_folder=str(BASE_DIR / "templates"), static_folder=str(BASE_DIR / "static")) +app.config["MAX_CONTENT_LENGTH"] = 32 * 1024 * 1024 # 32MB upload limit +CORS(app) + +# ============================================================================= +# Global State +# ============================================================================= + +training_state = { + "running": False, + "process": None, + "object_class": None, + "defect_type": None, + "output_dir": None, + "log_lines": [], + "start_time": None, + "completed": False, + "error": None +} + +training_lock = threading.Lock() + +# --- Generation Job Tracking --- +generation_jobs = {} +jobs_lock = threading.Lock() + + +class ProgressCapture: + """Captures BOTH stdout and stderr, parses [Agent] Step X: markers.""" + def __init__(self, job_id): + self.job_id = job_id + self._buffer = "" + self._step_re = re.compile(r'\[Agent\] Step (\d+):\s*(.+)') + self._defect_re = re.compile(r'\[Agent\] Defect (\d+)/(\d+):') + + def write(self, s): + self._buffer += s + while '\n' in self._buffer: + line, self._buffer = self._buffer.split('\n', 1) + self._process_line(line) + + def _process_line(self, line): + m = self._step_re.search(line) + if m: + with jobs_lock: + if self.job_id in generation_jobs: + generation_jobs[self.job_id]['step'] = int(m.group(1)) + generation_jobs[self.job_id]['step_text'] = m.group(2).strip() + d = self._defect_re.search(line) + if d: + with jobs_lock: + if self.job_id in generation_jobs: + job = generation_jobs[self.job_id] + job['defect_current'] = int(d.group(1)) + job['defect_total'] = int(d.group(2)) + # --- ADD THIS BLOCK BELOW --- + if 'unit_total' in job and job['unit_total'] > 0: + defects_per_image = job['defect_total'] + image_current = job['image_current'] + # Formula: (Current Image - 1) * Defects per Image + Current Defect + unit_current = (image_current - 1) * defects_per_image + job['defect_current'] + job['unit_current'] = min(unit_current, job['unit_total']) + if '[Agent] ERROR' in line or 'Traceback (most recent call last):' in line: + with jobs_lock: + if self.job_id in generation_jobs: + generation_jobs[self.job_id]['has_error_trace'] = True + + def flush(self): + pass + + def isatty(self): + return False + + +class TeeCapture: + """Duplicates output to capture AND original stream.""" + def __init__(self, capture, original): + self.capture = capture + self.original = original + + def write(self, s): + self.capture.write(s) + self.original.write(s) + + def flush(self): + self.capture.flush() + self.original.flush() + + def isatty(self): + return False + + +def run_generation_job(job_id, params, mode='defectfill'): + """Background thread: runs generation + captures progress.""" + with jobs_lock: + generation_jobs[job_id]['status'] = 'running' + generation_jobs[job_id]['step'] = 1 + generation_jobs[job_id]['step_text'] = 'Planning defects from product description...' + generation_jobs[job_id]['started_at'] = time.time() + + generation_jobs[job_id]['unit_total'] = len(params['clean_paths']) * params['num_defects'] + generation_jobs[job_id]['unit_current'] = 0 + + capture = ProgressCapture(job_id) + old_stdout = sys.stdout + old_stderr = sys.stderr + sys.stdout = TeeCapture(capture, old_stdout) + sys.stderr = TeeCapture(capture, old_stderr) + + try: + if mode == 'defectfill': + orchestrator = DefectFillOrchestrator( + device=params['device'], + output_dir=str(params['gen_output_dir']), + checkpoint_dir=str(params['ckpt_dir']), + object_class=params['object_class'], + defect_type=params['defect_type'] or "", + valid_object_classes=params['valid_object_classes'], + valid_defect_types=params['valid_defect_types'], + image_size=params['image_size'], + num_steps=params['num_steps'], + guidance_scale=params['guidance_scale'], + domain_hint=params.get('domain_hint', '') + ) + else: # defectdiffu + orchestrator = DefectDiffuOrchestrator( + device=params['device'], + output_dir=str(params['gen_output_dir']), + vlm_model=params.get('vlm_model', 'gemma3:12b'), + defectdiffu_ckpt=params['defectdiffu_ckpt'], + vae_path=params['vae_path'], + image_size=params['image_size'], + num_steps=params['num_steps'] + ) + + all_output_images = [] + all_grouped_results = [] + total_images = len(params['clean_paths']) + + for img_idx, clean_path in enumerate(params['clean_paths']): + with jobs_lock: + generation_jobs[job_id]['image_current'] = img_idx + 1 + generation_jobs[job_id]['image_total'] = total_images + generation_jobs[job_id]['step_text'] = f'Image {img_idx+1}/{total_images}: Planning defects...' + + if mode == 'defectfill': + result = orchestrator.run( + product_description=params['product_desc'], + image_path=clean_path, + num_defects=params['num_defects'], + defect_type=params['defect_type'], + object_class=params['object_class'] + ) + else: # defectdiffu + result = orchestrator.run( + product_description=params['product_desc'], + image_path=clean_path, + max_defects=params['num_defects'], + defect_type=params['defect_type'] + ) + + source_name = Path(clean_path).name + clean_rel = Path(clean_path).relative_to(BASE_DIR) + clean_url = f"/api/output/image/{clean_rel}" + + defects_for_image = [] + + for r in result.get("results", []): + if r.get("success") and "output_dir" in r: + out_d = Path(r["output_dir"]) + if out_d.exists(): + mask_file = None + output_file = None + blended_file = None + raw_file = None + for img_file in sorted(out_d.glob("*.png")): + fname = img_file.name.lower() + if "mask" in fname: + mask_file = img_file + elif "blended" in fname: + blended_file = img_file + elif "raw_defectdiffu" in fname or "output" in fname: + raw_file = img_file + elif "defect" in fname and blended_file is None: + output_file = img_file + + rel_path = img_file.relative_to(BASE_DIR) + all_output_images.append({ + "filename": img_file.name, + "url": f"/api/output/image/{rel_path}", + "defect_type": r.get("defect_type", "unknown"), + "object_class": r.get("object_class", params.get('object_class', 'unknown')), + "source_index": img_idx, + "source_name": source_name + }) + + # DefectDiffu has different output files + if mode == 'defectdiffu': + defects_for_image.append({ + "defect_type": r.get("defect_type", "unknown"), + "defect_label": f"Defect {len(defects_for_image) + 1}: {r.get('defect_type', 'unknown')}", + "clean_image": { + "filename": source_name, + "url": clean_url + }, + "output_image": { + "filename": blended_file.name if blended_file else (output_file.name if output_file else "output.png"), + "url": f"/api/output/image/{(blended_file or output_file or raw_file).relative_to(BASE_DIR)}" + } if (blended_file or output_file or raw_file) else None, + "raw_patch_image": { + "filename": raw_file.name if raw_file else "raw.png", + "url": f"/api/output/image/{raw_file.relative_to(BASE_DIR)}" + } if raw_file else None, + "mask_image": { + "filename": mask_file.name if mask_file else "mask.png", + "url": f"/api/output/image/{mask_file.relative_to(BASE_DIR)}" + } if mask_file else None + }) + elif mask_file and (output_file or blended_file): + defects_for_image.append({ + "defect_type": r.get("defect_type", "unknown"), + "defect_label": f"Defect {len(defects_for_image) + 1}: {r.get('defect_type', 'unknown')}", + "clean_image": { + "filename": source_name, + "url": clean_url + }, + "mask_image": { + "filename": mask_file.name, + "url": f"/api/output/image/{mask_file.relative_to(BASE_DIR)}" + }, + "output_image": { + "filename": output_file.name if output_file else blended_file.name, + "url": f"/api/output/image/{(output_file or blended_file).relative_to(BASE_DIR)}" + } + }) + + if defects_for_image: + all_grouped_results.append({ + "source_index": img_idx, + "source_name": source_name, + "defect_count": len(defects_for_image), + "defects": defects_for_image + }) + + orchestrator.cleanup() + sys.stdout = old_stdout + sys.stderr = old_stderr + + elapsed = time.time() - generation_jobs[job_id]['started_at'] + + with jobs_lock: + generation_jobs[job_id]['status'] = 'completed' + generation_jobs[job_id]['step'] = 5 + generation_jobs[job_id]['step_text'] = 'Generation complete' + generation_jobs[job_id]['elapsed_time'] = elapsed + generation_jobs[job_id]['result'] = { + 'success': True, + 'grouped_results': all_grouped_results, + 'output_images': all_output_images, + 'experiment_id': params.get('object_class', 'defectdiffu'), + 'product_type': params.get('object_class', 'defectdiffu'), + 'elapsed_time': elapsed, + 'results_summary': [{"type": img['defect_type'], "success": True} for img in all_output_images], + 'output_dir': str(params['gen_output_dir'].relative_to(BASE_DIR)) + } + except Exception as e: + traceback.print_exc() + with jobs_lock: + generation_jobs[job_id]['status'] = 'error' + generation_jobs[job_id]['error'] = str(e) + generation_jobs[job_id]['step_text'] = f'Error: {str(e)}' + finally: + sys.stdout = old_stdout + sys.stderr = old_stderr + + +# ============================================================================= +# Helpers +# ============================================================================= + +def allowed_file(filename, exts={".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}): + return Path(filename).suffix.lower() in exts + + +def get_checkpoints(): + """Discover trained checkpoints: object_class -> [defect_types].""" + mapping = {} + if not CHECKPOINT_DIR.exists(): + return mapping + for obj_dir in CHECKPOINT_DIR.iterdir(): + if not obj_dir.is_dir(): + continue + if obj_dir.name == 'checkpoints_old': + continue + defects = [] + for defect_dir in obj_dir.iterdir(): + if not defect_dir.is_dir(): + continue + ckpt = defect_dir / "checkpoints" / "checkpoint_final.pt" + ckpt_alt = defect_dir / "checkpoint_final.pt" + if ckpt.exists() or ckpt_alt.exists(): + defects.append(defect_dir.name) + if defects: + mapping[obj_dir.name] = defects + return mapping + + +def generate_mask_from_rect(image_path, rect, output_mask_path): + """Generate a binary mask from rectangle coordinates.""" + img = Image.open(image_path).convert("RGB") + w, h = img.size + mask = Image.new("L", (w, h), 0) + draw = ImageDraw.Draw(mask) + + x1 = max(0, int(rect["x"])) + y1 = max(0, int(rect["y"])) + x2 = min(w, int(rect["x"] + rect["width"])) + y2 = min(h, int(rect["y"] + rect["height"])) + + draw.rectangle([x1, y1, x2, y2], fill=255) + mask.save(output_mask_path) + return {"x1": x1, "y1": y1, "x2": x2, "y2": y2} + + +def stream_training_logs(process): + """Read training subprocess output in background thread.""" + global training_state + try: + for line in iter(process.stdout.readline, ""): + line_stripped = line.rstrip() + if line_stripped: + with training_lock: + training_state["log_lines"].append(line_stripped) + training_state["log_lines"] = training_state["log_lines"][-2000:] + + process.stdout.close() + process.wait() + + with training_lock: + training_state["running"] = False + training_state["completed"] = (process.returncode == 0) + + if process.returncode == 0: + object_class = training_state.get("object_class") + if object_class: + sample_img = None + good_dir = TRAINING_DATA_DIR / object_class / "test" / "good" + if good_dir.exists(): + for f in sorted(good_dir.iterdir()): + if f.is_file() and allowed_file(f.name): + sample_img = f + break + if not sample_img: + prod_dir = TRAINING_DATA_DIR / object_class + if prod_dir.exists(): + for f in sorted(prod_dir.rglob("*")): + if f.is_file() and allowed_file(f.name) and "mask" not in f.name.lower(): + sample_img = f + break + if sample_img: + save_to_trained_gallery(str(sample_img), object_class) + else: + print(f"[Gallery Warning] No clean sample image found for '{object_class}'.") + else: + training_state["error"] = f"Training exited with code {process.returncode}" + + except Exception as e: + with training_lock: + training_state["running"] = False + training_state["error"] = f"Log thread crashed: {str(e)}" + + +def save_to_trained_gallery(clean_image_path: str, product_name: str): + """Copies 1 clean image to the gallery named after the product for DINOv2 matching.""" + try: + for existing_file in TRAINED_GALLERY_DIR.iterdir(): + if existing_file.is_file() and existing_file.stem == product_name: + print(f"[Gallery Info] Reference image for '{product_name}' already exists in gallery. Skipping copy.") + return True + + src = Path(clean_image_path) + if not src.exists(): + print(f"[Gallery Error] Clean image not found at: {clean_image_path}") + return False + + ext = src.suffix.lower() + if ext not in ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff']: + ext = '.png' + + dst = TRAINED_GALLERY_DIR / f"{product_name}{ext}" + shutil.copy2(src, dst) + print(f"[Gallery Success] Added reference image to gallery: {dst}") + return True + except Exception as e: + print(f"[Gallery Error] Failed to save image to gallery: {e}") + return False + + +# ============================================================================= +# Routes - Pages +# ============================================================================= + +@app.route("/") +def index(): + return render_template("index.html") + + +# ============================================================================= +# Routes - DefectDiffu Config +# ============================================================================= + +@app.route("/api/defectdiffu/config", methods=["GET"]) +def get_defectdiffu_config(): + """Get saved DefectDiffu model paths.""" + try: + with open(DEFECTDIFFU_CONFIG_PATH, "r") as f: + config = json.load(f) + return jsonify({"success": True, "config": config}) + except Exception as e: + return jsonify({"success": False, "error": str(e)}) + + +@app.route("/api/defectdiffu/config", methods=["POST"]) +def save_defectdiffu_config(): + """Save DefectDiffu model paths.""" + data = request.get_json() + config = { + "ckpt_path": data.get("ckpt_path", "").strip(), + "vae_path": data.get("vae_path", "").strip(), + "vlm_model": data.get("vlm_model", "gemma3:12b").strip() + } + try: + with open(DEFECTDIFFU_CONFIG_PATH, "w") as f: + json.dump(config, f, indent=2) + return jsonify({"success": True, "message": "DefectDiffu config saved."}) + except Exception as e: + return jsonify({"success": False, "error": str(e)}) + + +# ============================================================================= +# Routes - Training Data Upload & Mask (DefectFill) +# ============================================================================= + +@app.route("/api/training/good-images/count", methods=["GET"]) +def get_good_image_count(): + """Get count of existing good images for a product.""" + object_class = request.args.get("object_class", "").strip() + if not object_class: + return jsonify({"success": False, "error": "Object class required."}) + + good_dir = TRAINING_DATA_DIR / object_class / "test" / "good" + count = 0 + if good_dir.exists(): + count = len([f for f in good_dir.iterdir() if f.is_file() and allowed_file(f.name)]) + + return jsonify({ + "success": True, + "count": count, + "path": str(good_dir.relative_to(DEFECTFILL_ROOT)) + }) + + +@app.route("/api/training/upload-good-images", methods=["POST"]) +def upload_good_images(): + """Upload good (non-defective) images to temp staging area.""" + object_class = request.form.get("object_class", "").strip() + if not object_class: + return jsonify({"success": False, "error": "Product name is required."}) + + if "images" not in request.files: + return jsonify({"success": False, "error": "No images provided."}) + + files = request.files.getlist("images") + target_dir = TEMP_IMAGE_DIR / object_class / "test" / "good" + if target_dir.exists(): + shutil.rmtree(target_dir) + target_dir.mkdir(parents=True, exist_ok=True) + + saved = [] + for file in files: + if file and allowed_file(file.filename): + safe_name = Path(file.filename).name + save_path = target_dir / safe_name + counter = 1 + while save_path.exists(): + stem = Path(file.filename).stem + suffix = Path(file.filename).suffix + save_path = target_dir / f"{stem}_{counter:02d}{suffix}" + counter += 1 + file.save(save_path) + + img = cv2.imread(str(save_path)) + if img is not None: + resized = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LANCZOS4) + cv2.imwrite(str(save_path), resized) + + rel = save_path.relative_to(TEMP_IMAGE_DIR).as_posix() + saved.append({ + "filename": save_path.name, + "path": rel, + "url": f"/api/temp/image/{rel}" + }) + + return jsonify({ + "success": True, + "saved": saved, + "count": len(saved), + "total_in_folder": len([f for f in target_dir.iterdir() if f.is_file() and allowed_file(f.name)]), + "target_dir": str(target_dir.relative_to(DEFECTFILL_ROOT).as_posix()) + }) + + +@app.route("/api/training/good-image/") +def serve_good_image(subpath): + """Serve uploaded good images.""" + return send_from_directory(TRAINING_DATA_DIR, subpath) + + +@app.route("/api/training/upload-images", methods=["POST"]) +def upload_training_images(): + """Upload defect images to temp staging area.""" + object_class = request.form.get("object_class", "").strip() + defect_type = request.form.get("defect_type", "").strip() + + if not object_class or not defect_type: + return jsonify({"success": False, "error": "Product name and defect type are required."}) + + if "images" not in request.files: + return jsonify({"success": False, "error": "No images provided."}) + + files = request.files.getlist("images") + if len(files) < 1: + return jsonify({"success": False, "error": f"At least 1 defect images required. Received {len(files)}."}) + + target_dir = TEMP_IMAGE_DIR / object_class / "train" / "defective" / defect_type + if target_dir.exists(): + shutil.rmtree(target_dir) + target_dir.mkdir(parents=True, exist_ok=True) + + saved = [] + for file in files: + if file and allowed_file(file.filename): + filename = file.filename + safe_name = Path(filename).name + save_path = target_dir / safe_name + counter = 1 + while save_path.exists(): + stem = Path(filename).stem + suffix = Path(filename).suffix + save_path = target_dir / f"{stem}_{counter:02d}{suffix}" + counter += 1 + file.save(save_path) + rel = save_path.relative_to(TEMP_IMAGE_DIR).as_posix() + saved.append({ + "filename": save_path.name, + "path": rel, + "url": f"/api/temp/image/{rel}" + }) + + return jsonify({ + "success": True, + "saved": saved, + "count": len(saved), + "target_dir": str(target_dir.relative_to(DEFECTFILL_ROOT).as_posix()) + }) + + +@app.route("/api/training/image/") +def serve_training_image(subpath): + """Serve uploaded training images.""" + return send_from_directory(TRAINING_DATA_DIR, subpath) + + +@app.route("/api/temp/image/") +def serve_temp_image(subpath): + """Serve images from the temp staging directory.""" + return send_from_directory(TEMP_IMAGE_DIR, subpath) + + +@app.route("/api/training/save-mask", methods=["POST"]) +def save_mask(): + """Save user-drawn rectangle mask for a training image.""" + data = request.get_json() + image_path = data.get("image_path", "") + rect = data.get("rect", {}) + + if not image_path or not rect: + return jsonify({"success": False, "error": "Missing image_path or rect data."}) + + full_path = TRAINING_DATA_DIR / image_path + if not full_path.exists(): + return jsonify({"success": False, "error": "Image not found."}) + + rel_path = Path(image_path) + new_parts = [] + for part in rel_path.parts: + if part == "defective": + new_parts.append("defective_masks") + else: + new_parts.append(part) + + mask_path = TRAINING_DATA_DIR / Path(*new_parts) + mask_path.parent.mkdir(parents=True, exist_ok=True) + + coords = generate_mask_from_rect(full_path, rect, mask_path) + + return jsonify({ + "success": True, + "mask_path": str(mask_path.relative_to(TRAINING_DATA_DIR)), + "coords": coords + }) + + +@app.route("/api/training/align", methods=["POST"]) +def align_images(): + """Align all temp images for a product to a reference image, then move to final.""" + if not CV2_AVAILABLE: + return jsonify({"success": False, "error": "OpenCV not installed. Cannot align images."}) + + object_class = request.form.get("object_class", "").strip() + reference = request.files.get("reference") + + if not object_class or not reference: + return jsonify({"success": False, "error": "Object class and reference image are required."}) + + ref_dir = TEMP_IMAGE_DIR / object_class / "reference" + ref_dir.mkdir(parents=True, exist_ok=True) + ref_path = ref_dir / "reference.png" + reference.save(ref_path) + + obj_temp_dir = TEMP_IMAGE_DIR / object_class + aligned = [] + errors = [] + + for img_path in sorted(obj_temp_dir.rglob("*")): + if not img_path.is_file(): + continue + if "reference" in str(img_path.relative_to(obj_temp_dir)).split(os.sep): + continue + if not allowed_file(img_path.name): + continue + + rel = img_path.relative_to(obj_temp_dir) + final_path = TRAINING_DATA_DIR / object_class / rel + final_path.parent.mkdir(parents=True, exist_ok=True) + + try: + align_image_to_reference(str(ref_path), str(img_path), str(final_path)) + img_arr = cv2.imread(str(final_path)) + if img_arr is not None: + img_arr_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4) + cv2.imwrite(str(final_path), img_arr_resized) + + aligned.append({ + "filename": final_path.name, + "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()), + "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}" + }) + except Exception as e: + errors.append({"file": str(rel), "error": str(e)}) + shutil.copy2(str(img_path), str(final_path)) + img_arr = cv2.imread(str(final_path)) + if img_arr is not None: + img_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4) + cv2.imwrite(str(final_path), img_resized) + + aligned.append({ + "filename": final_path.name, + "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()), + "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}" + }) + + return jsonify({ + "success": True, + "aligned_count": len(aligned), + "error_count": len(errors), + "errors": errors, + "images": aligned + }) + + +@app.route("/api/training/skip-align", methods=["POST"]) +def skip_align(): + """Copy all temp images directly to final location without alignment.""" + data = request.get_json() + object_class = data.get("object_class", "").strip() + + if not object_class: + return jsonify({"success": False, "error": "Object class required."}) + + obj_temp_dir = TEMP_IMAGE_DIR / object_class + moved = [] + + if obj_temp_dir.exists(): + for img_path in sorted(obj_temp_dir.rglob("*")): + if not img_path.is_file(): + continue + if "reference" in str(img_path.relative_to(obj_temp_dir)).split(os.sep): + continue + + rel = img_path.relative_to(obj_temp_dir) + final_path = TRAINING_DATA_DIR / object_class / rel + final_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(str(img_path), str(final_path)) + + img_arr = cv2.imread(str(final_path)) + if img_arr is not None: + img_resized = cv2.resize(img_arr, (512, 512), interpolation=cv2.INTER_LANCZOS4) + cv2.imwrite(str(final_path), img_resized) + + moved.append({ + "filename": final_path.name, + "path": str(final_path.relative_to(TRAINING_DATA_DIR).as_posix()), + "url": f"/api/training/image/{final_path.relative_to(TRAINING_DATA_DIR).as_posix()}" + }) + + return jsonify({"success": True, "moved_count": len(moved), "images": moved}) + + +# ============================================================================= +# Routes - Domain Instructions +# ============================================================================= + +@app.route("/api/domain-instructions", methods=["GET"]) +def get_domain_instructions(): + """Get saved domain instructions.""" + try: + with open(DOMAIN_CONFIG_PATH, "r") as f: + config = json.load(f) + return jsonify({"success": True, "config": config}) + except Exception as e: + return jsonify({"success": False, "error": str(e)}) + + +@app.route("/api/domain-instructions", methods=["POST"]) +def save_domain_instructions(): + """Save domain-specific instructions for prompts.""" + data = request.get_json() + object_class = data.get("object_class", "default").strip() + instructions = data.get("instructions", "").strip() + + try: + with open(DOMAIN_CONFIG_PATH, "r") as f: + config = json.load(f) + + if "active" not in config: + config["active"] = {} + + # ========================================================== + # NEW APPEND LOGIC STARTS HERE + # ========================================================== + existing_instructions = config["active"].get(object_class, "") + + if instructions.strip(): + # If the user wrote something in the textarea, append it to the old instructions + if existing_instructions.strip(): + # Use double newlines to clearly separate the old and new contexts + config["active"][object_class] = existing_instructions.strip() + "\n\n" + instructions.strip() + else: + # If there were no old instructions, just use the new ones + config["active"][object_class] = instructions.strip() + else: + # If the user left the textarea EMPTY, we do NOT erase old instructions. + # We simply keep the existing instructions untouched. + pass + # ========================================================== + # NEW APPEND LOGIC ENDS HERE + # ========================================================== + + with open(DOMAIN_CONFIG_PATH, "w") as f: + json.dump(config, f, indent=2) + + return jsonify({"success": True, "message": f"Instructions saved for '{object_class}'."}) + except Exception as e: + return jsonify({"success": False, "error": str(e)}) + + +# ============================================================================= +# Routes - Training Execution (DefectFill) +# ============================================================================= + +@app.route("/api/check-product-exists", methods=["GET"]) +def check_product_exists(): + """Check if a product (object_class) already exists in training data.""" + object_class = request.args.get("object_class", "").strip() + if not object_class: + return jsonify({"success": False, "error": "Product name is required."}) + + product_dir = TRAINING_DATA_DIR / object_class + exists = product_dir.exists() and product_dir.is_dir() + + existing_defect_types = [] + sample_image_url = None + + if exists: + defective_dir = product_dir / "defective" + if defective_dir.exists(): + for defect_dir in sorted(defective_dir.iterdir()): + if defect_dir.is_dir(): + existing_defect_types.append(defect_dir.name) + if not sample_image_url: + for img_file in sorted(defect_dir.iterdir()): + if allowed_file(img_file.name): + rel_path = img_file.relative_to(BASE_DIR) + sample_image_url = f"/api/training/image/{rel_path}" + break + + return jsonify({ + "success": True, + "exists": exists, + "existing_defect_types": existing_defect_types, + "sample_image_url": sample_image_url + }) + + +@app.route("/api/training/start", methods=["POST"]) +def start_training(): + """Start few-shot training with train.py.""" + global training_state + + data = request.get_json() + object_class = data.get("object_class", "").strip() + defect_type = data.get("defect_type", "").strip() + lora_rank = data.get("lora_rank", 8) + lora_alpha = data.get("lora_alpha", 16) + max_steps = data.get("max_train_steps", 1500) + batch_size = data.get("batch_size", 2) + gradient_accum = data.get("gradient_accumulation_steps", 2) + lambda_defect = data.get("lambda_defect", 0.5) + lambda_obj = data.get("lambda_obj", 0.2) + lambda_attn = data.get("lambda_attn", 0.05) + alpha = data.get("alpha", 0.3) + + if not object_class or not defect_type: + return jsonify({"success": False, "error": "Object class and defect type are required."}) + + with training_lock: + if training_state["running"]: + return jsonify({"success": False, "error": "Training is already in progress."}) + + output_dir = CHECKPOINT_DIR / object_class / defect_type + output_dir.mkdir(parents=True, exist_ok=True) + + train_script = TRAIN_PY + if not train_script.exists(): + train_script = BASE_DIR.parent / "train.py" + + cmd = [ + sys.executable, '-u', str(train_script), + "--data_dir", str(TRAINING_DATA_DIR), + "--object_class", object_class, + "--defect_type", defect_type, + "--output_dir", str(output_dir), + "--lora_rank", str(lora_rank), + "--lora_alpha", str(lora_alpha), + "--max_train_steps", str(max_steps), + "--batch_size", str(batch_size), + "--gradient_accumulation_steps", str(gradient_accum), + "--lambda_defect", str(lambda_defect), + "--lambda_obj", str(lambda_obj), + "--lambda_attn", str(lambda_attn), + "--alpha", str(alpha), + "--save_steps", str(max(500, max_steps // 3)), + "--lr_warmup_steps", str(min(100, max_steps // 10)), + "--dilate_mask", "False", + "--seed", str(int(time.time() * 1000) % (2**31)) + ] + + try: + with training_lock: + training_state["running"] = True + training_state["object_class"] = object_class + training_state["defect_type"] = defect_type + training_state["output_dir"] = str(output_dir) + training_state["log_lines"] = [] + training_state["start_time"] = datetime.now().isoformat() + training_state["completed"] = False + training_state["error"] = None + + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=1, + universal_newlines=True, + cwd=str(DEFECTFILL_ROOT) + ) + + with training_lock: + training_state["process"] = process + + log_thread = threading.Thread(target=stream_training_logs, args=(process,)) + log_thread.daemon = True + log_thread.start() + + return jsonify({ + "success": True, + "message": "Training started.", + "command": " ".join(cmd), + "output_dir": str(output_dir.relative_to(DEFECTFILL_ROOT)) + }) + + except Exception as e: + with training_lock: + training_state["running"] = False + training_state["error"] = str(e) + return jsonify({"success": False, "error": str(e)}) + + +@app.route("/api/training/status", methods=["GET"]) +def training_status(): + """Get current training status and logs.""" + with training_lock: + state = dict(training_state) + state["process"] = None + return jsonify({"success": True, "status": state}) + + +@app.route("/api/training/stop", methods=["POST"]) +def stop_training(): + """Stop running training.""" + global training_state + with training_lock: + if training_state["process"] and training_state["process"].poll() is None: + training_state["process"].terminate() + training_state["running"] = False + training_state["error"] = "Training stopped by user." + return jsonify({"success": True, "message": "Training terminated."}) + return jsonify({"success": False, "error": "No training is running."}) + + +# ============================================================================= +# Routes - Checkpoints & Generation (DefectFill) +# ============================================================================= + +@app.route("/api/detect-similar-product", methods=["POST"]) +def detect_similar_product(): + """Endpoint to identify if a clean product image has already been trained.""" + if "image" not in request.files: + return jsonify({"success": False, "error": "No clean product image uploaded."}) + + file = request.files["image"] + temp_path = UPLOAD_FOLDER / f"temp_detect_{uuid.uuid4().hex}{Path(file.filename).suffix}" + + try: + file.save(temp_path) + threshold = float(request.form.get("threshold", 0.85)) + result = check_product_training_status( + target_image_path=str(temp_path), + gallery_dir=TRAINED_GALLERY_DIR, + similarity_threshold=threshold + ) + return jsonify({"success": True, "result": result}) + except Exception as e: + return jsonify({"success": False, "error": str(e)}) + finally: + if temp_path.exists(): + os.remove(temp_path) + + +@app.route("/api/checkpoints", methods=["GET"]) +def list_checkpoints(): + """List available trained checkpoints.""" + return jsonify({"success": True, "checkpoints": get_checkpoints()}) + + +@app.route("/api/generate", methods=["POST"]) +def generate_defect(): + """Start DefectFill generation in background and return a job_id for polling.""" + object_class = request.form.get("object_class", "").strip() + product_desc = request.form.get("product_desc", "").strip() + detection_mode = request.form.get("detection_mode", "dots") + + if not object_class or not product_desc: + return jsonify({"success": False, "error": "Object class and product description are required."}) + + defect_type = request.form.get("defect_type", "").strip() or None + num_defects = int(request.form.get("num_defects", 3)) + guidance_scale = float(request.form.get("guidance_scale", 8.0)) + num_steps = int(request.form.get("num_steps", 50)) + image_size = int(request.form.get("image_size", 512)) + device = request.form.get("device", "cuda") + + clean_files = request.files.getlist("clean_images") + if not clean_files: + return jsonify({"success": False, "error": "Clean images are required."}) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + allowed_exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp'} + clean_paths = [] + + for f in clean_files: + if not f or not f.filename: + continue + ext = os.path.splitext(f.filename)[1].lower() + if ext not in allowed_exts: + continue + safe_name = os.path.basename(f.filename) + clean_filename = f"clean_{timestamp}_{safe_name}" + clean_path = UPLOAD_FOLDER / clean_filename + f.save(clean_path) + clean_paths.append(str(clean_path)) + + if not clean_paths: + return jsonify({"success": False, "error": "No valid image files found in selected folder."}) + + ckpt_dir = CHECKPOINT_DIR + valid_types = get_checkpoints().get(object_class, []) + + if defect_type and defect_type not in valid_types: + return jsonify({ + "success": False, + "error": f"Defect type '{defect_type}' not found for '{object_class}'. Available: {valid_types}" + }) + + valid_object_classes = list(get_checkpoints().keys()) + valid_defect_types = get_checkpoints() + + gen_output_dir = OUTPUT_DIR / f"gen_{timestamp}" + gen_output_dir.mkdir(parents=True, exist_ok=True) + + try: + domain_hint = "" + if DOMAIN_CONFIG_PATH.exists(): + with open(DOMAIN_CONFIG_PATH, "r") as f: + dconfig = json.load(f) + domain_hint = dconfig.get("active", {}).get(object_class, "") + + # If detection_mode == "auto", do nothing - let the VLM infer entirely from the product description and domain context. + if detection_mode == "lines": + product_desc += " Focus defect placement on metallic pins, leads, or legs." + elif detection_mode == "dots": + product_desc += " Focus defect placement on solder ball array, BGA grid, or rounded pads." + elif detection_mode == "single_rounded": + product_desc += " Focus defect placement on a single, isolated rounded feature such as a mounting hole, circular pad, individual via, or isolated dot." + + # if domain_hint: + # product_desc += f" Domain context: {domain_hint}" + + job_params = { + 'device': device, + 'gen_output_dir': gen_output_dir, + 'ckpt_dir': ckpt_dir, + 'object_class': object_class, + 'defect_type': defect_type, + 'valid_object_classes': valid_object_classes, + 'valid_defect_types': valid_defect_types, + 'image_size': image_size, + 'num_steps': num_steps, + 'guidance_scale': guidance_scale, + 'product_desc': product_desc, + 'domain_hint': domain_hint, + 'clean_paths': clean_paths, + 'num_defects': num_defects, + } + + job_id = str(uuid.uuid4()) + with jobs_lock: + generation_jobs[job_id] = { + 'status': 'starting', + 'step': 0, + 'total_steps': 5, + 'step_text': 'Initializing...', + 'defect_current': 0, + 'defect_total': 0, + 'image_current': 0, + 'image_total': 0, + 'result': None, + 'error': None, + 'has_error_trace': False + } + + thread = threading.Thread(target=run_generation_job, args=(job_id, job_params, 'defectfill'), daemon=True) + thread.start() + + return jsonify({'success': True, 'job_id': job_id}) + + except Exception as e: + traceback.print_exc() + return jsonify({"success": False, "error": str(e)}) + + +# ============================================================================= +# Routes - DefectDiffu Generation +# ============================================================================= + +@app.route("/api/defectdiffu/generate", methods=["POST"]) +def generate_defectdiffu(): + """Start DefectDiffu generation in background and return a job_id for polling.""" + product_desc = request.form.get("product_desc", "").strip() + defect_type = request.form.get("defect_type", "").strip() or None + num_defects = int(request.form.get("num_defects", 3)) + image_size = int(request.form.get("image_size", 512)) + num_steps = int(request.form.get("num_steps", 50)) + device = request.form.get("device", "cuda") + + if not product_desc: + return jsonify({"success": False, "error": "Product description is required."}) + + # Load model config + try: + with open(DEFECTDIFFU_CONFIG_PATH, "r") as f: + dd_config = json.load(f) + except Exception as e: + return jsonify({"success": False, "error": f"Failed to load DefectDiffu config: {e}"}) + + ckpt_path = dd_config.get("ckpt_path", "") + vae_path = dd_config.get("vae_path", "") + vlm_model = dd_config.get("vlm_model", "gemma3:12b") + + if not ckpt_path or not vae_path: + return jsonify({"success": False, "error": "DefectDiffu model paths not configured. Please set them in the UI first."}) + + if not Path(ckpt_path).exists(): + return jsonify({"success": False, "error": f"Checkpoint not found: {ckpt_path}"}) + if not Path(vae_path).exists(): + return jsonify({"success": False, "error": f"VAE not found: {vae_path}"}) + + clean_files = request.files.getlist("clean_images") + if not clean_files: + return jsonify({"success": False, "error": "Clean images are required."}) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + allowed_exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp'} + clean_paths = [] + + for f in clean_files: + if not f or not f.filename: + continue + ext = os.path.splitext(f.filename)[1].lower() + if ext not in allowed_exts: + continue + safe_name = os.path.basename(f.filename) + clean_filename = f"clean_dd_{timestamp}_{safe_name}" + clean_path = UPLOAD_FOLDER / clean_filename + f.save(clean_path) + clean_paths.append(str(clean_path)) + + if not clean_paths: + return jsonify({"success": False, "error": "No valid image files found."}) + + gen_output_dir = OUTPUT_DIR / f"defectdiffu_{timestamp}" + gen_output_dir.mkdir(parents=True, exist_ok=True) + + try: + job_params = { + 'device': device, + 'gen_output_dir': gen_output_dir, + 'defectdiffu_ckpt': ckpt_path, + 'vae_path': vae_path, + 'vlm_model': vlm_model, + 'image_size': image_size, + 'num_steps': num_steps, + 'product_desc': product_desc, + 'clean_paths': clean_paths, + 'num_defects': num_defects, + 'defect_type': defect_type, + } + + job_id = str(uuid.uuid4()) + with jobs_lock: + generation_jobs[job_id] = { + 'status': 'starting', + 'step': 0, + 'total_steps': 5, + 'step_text': 'Initializing DefectDiffu...', + 'defect_current': 0, + 'defect_total': 0, + 'image_current': 0, + 'image_total': 0, + 'result': None, + 'error': None, + 'has_error_trace': False + } + + thread = threading.Thread(target=run_generation_job, args=(job_id, job_params, 'defectdiffu'), daemon=True) + thread.start() + + return jsonify({'success': True, 'job_id': job_id}) + + except Exception as e: + traceback.print_exc() + return jsonify({"success": False, "error": str(e)}) + + +@app.route('/api/generation-status/', methods=['GET']) +def generation_status(job_id): + with jobs_lock: + job = generation_jobs.get(job_id) + if not job: + return jsonify({'success': False, 'error': 'Job not found'}), 404 + return jsonify({'success': True, 'job': job}) + + +@app.route("/api/output/image/") +def serve_output_image(subpath): + """Serve generated output images.""" + return send_from_directory(BASE_DIR, subpath) + + +# ============================================================================= +# Routes - System Info +# ============================================================================= + +@app.route("/api/system/info", methods=["GET"]) +def system_info(): + """Get system information.""" + import torch + info = { + "cuda_available": torch.cuda.is_available(), + "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0, + "cuda_device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + "project_root": str(PROJECT_ROOT), + "defectfill_root": str(DEFECTFILL_ROOT), + "defectdiffu_root": str(DEFECTDIFFU_ROOT), + "checkpoint_dir": str(CHECKPOINT_DIR), + "training_data_dir": str(TRAINING_DATA_DIR) + } + return jsonify({"success": True, "info": info}) + + +# ============================================================================= +# Main +# ============================================================================= + +if __name__ == "__main__": + print("=" * 60) + print(" DefectFill & DefectDiffu Unified Web UI") + print("=" * 60) + print(f" Project root: {PROJECT_ROOT}") + print(f" DefectFill root: {DEFECTFILL_ROOT}") + print(f" DefectDiffu root: {DEFECTDIFFU_ROOT}") + print(f" Checkpoints: {CHECKPOINT_DIR}") + print(f" Training data: {TRAINING_DATA_DIR}") + print(f" Outputs: {OUTPUT_DIR}") + print("-" * 60) + print(" Open http://127.0.0.1:5000 in your browser") + print("=" * 60) + + app.run(host="0.0.0.0", port=5000, debug=False, threaded=True) diff --git a/config/defectdiffu_config.json b/config/defectdiffu_config.json new file mode 100644 index 0000000000000000000000000000000000000000..3982f9ba43702336fbf70fe271293412b19442cc --- /dev/null +++ b/config/defectdiffu_config.json @@ -0,0 +1,5 @@ +{ + "ckpt_path": "C:\\Users\\admin_mtds\\OneDrive\\Desktop\\ChinKuan\\ArtiAgent - Defect\\engine\\DefectDiffu\\checkpoint_old-4\\model_300.pth", + "vae_path": "C:\\Users\\admin_mtds\\OneDrive\\Desktop\\ChinKuan\\ArtiAgent - Defect\\engine\\DefectDiffu\\checkpoints\\sd-vae-ft-mse", + "vlm_model": "gemma3:12b" +} \ No newline at end of file diff --git a/config/domain_instructions.json b/config/domain_instructions.json new file mode 100644 index 0000000000000000000000000000000000000000..d62a1fb18c6ce4f361b3fbddb1442934064328c3 --- /dev/null +++ b/config/domain_instructions.json @@ -0,0 +1,7 @@ +{ + "examples": {}, + "active": { + "triac": "Combine is contains all other defect type under the same object class", + "metal_mounting_plate_assembly": "If the image is a metal mounting plate assembly (e.g., an aluminum plate with a central circular hole, multiple mounting holes, and Phillips-head screws):\n- The component consists of a flat metal plate, threaded mounting holes, Phillips-head screws, and possibly a secondary component (like a white ceramic block) attached via screws.\n- **CRITICAL CONSTRAINT**: You MUST exclusively propose defects regarding the **screws and screw holes**. Do NOT propose defects on the metal plate surface, the central circular cutout, or attached wires, unless they are directly related to the presence or absence of screws.\n- Acceptable screw defects include:\n 1. `missing_screw`: A screw is absent from a designated threaded mounting hole within the assembly pattern (e.g., a hole on the plate edge or on the attached white block that is left completely open).\n 2. `extra_screw`: A screw is present in a threaded hole that is not part of the required mounting pattern, or a screw placed in a non-threaded hole where it doesn't belong. For this defect type, you MUST only consider and propose locations that are in the upper half of the image (i.e., the top 50% of the vertical space). Do not propose any extra_screw defects in the lower half.\n- **IMPORTANT**: Depending on the defect type, you MUST choose a specific target_entity to guide the perception engine correctly:\n - If the defect_type is `missing_screw`, the target_entity MUST be **\"empty_screw_hole\"**. (The engine will look for a dark, empty circle).\n - If the defect_type is `extra_screw`, the target_entity MUST be **\"extra_screw_head\"**. (The engine will look for a bright, solid screw head).\n- The target_subentity must specify the exact location within the assembly (e.g., \"top-left corner\", \"white block center\").\n- The location_hint must state the spatial position clearly, referencing the plate's edges or the attached block (e.g., \"bottom-right corner of the plate\", \"middle-left of the white block\").\n- The `description` must clearly detail the visual evidence: e.g., \"The threaded hole in the top-left corner is completely empty, exposing the metallic threads, indicating a missing screw.\" or \"An unaccounted Phillips-head screw is visible in the bottom-right unthreaded hole.\"\n- The `defect_coverage_ratio` should be 1.0. This represents the area of a single screw head relative to the full image.\n- The `mask_shape` must be **\"rectangle\"** (as enforced by the system)." + } +} \ No newline at end of file diff --git a/index.html b/index.html new file mode 100644 index 0000000000000000000000000000000000000000..b0c4b3666032a737f3903db53e6a8a9272483e28 --- /dev/null +++ b/index.html @@ -0,0 +1,19 @@ + + + + + + My static Space + + + +
+

Welcome to your static Space!

+

You can modify this app directly by editing index.html in the Files and versions tab.

+

+ Also don't forget to check the + Spaces documentation. +

+
+ + diff --git a/requirement.txt.txt b/requirement.txt.txt new file mode 100644 index 0000000000000000000000000000000000000000..28dec92dbbef40901f0122808da82986340b1cad --- /dev/null +++ b/requirement.txt.txt @@ -0,0 +1,4 @@ +flask +flask-cors +pillow +numpy \ No newline at end of file diff --git a/static/css/style.css b/static/css/style.css new file mode 100644 index 0000000000000000000000000000000000000000..b818a265c0f87d57a03cb5084bea47d78b961c1d --- /dev/null +++ b/static/css/style.css @@ -0,0 +1,930 @@ +/* ============================================ + ArtiAgent Web UI - Unified Stylesheet + Supports both DefectFill & DefectDiffu + ============================================ */ + +:root { + --primary: #2563eb; + --primary-dark: #1d4ed8; + --primary-light: #dbeafe; + --success: #059669; + --success-light: #d1fae5; + --warning: #d97706; + --warning-light: #fef3c7; + --danger: #dc2626; + --danger-light: #fee2e2; + --dark: #1e293b; + --gray-50: #f8fafc; + --gray-100: #f1f5f9; + --gray-200: #e2e8f0; + --gray-300: #cbd5e1; + --gray-400: #94a3b8; + --gray-500: #64748b; + --gray-600: #475569; + --gray-700: #334155; + --gray-800: #1e293b; + --radius: 8px; + --shadow-sm: 0 1px 2px 0 rgb(0 0 0 / 0.05); + --shadow: 0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1); + --shadow-md: 0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1); + --shadow-lg: 0 10px 15px -3px rgb(0 0 0 / 0.1), 0 4px 6px -4px rgb(0 0 0 / 0.1); +} + +* { + margin: 0; + padding: 0; + box-sizing: border-box; +} + +body { + font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif; + background: var(--gray-50); + color: var(--gray-800); + line-height: 1.5; + min-height: 100vh; +} + +/* Header */ +.header { + background: var(--dark); + color: white; + padding: 1rem 2rem; + display: flex; + align-items: center; + justify-content: space-between; + box-shadow: var(--shadow-md); + position: sticky; + top: 0; + z-index: 100; +} + +.header h1 { + font-size: 1.25rem; + font-weight: 600; + display: flex; + align-items: center; + gap: 0.5rem; +} + +.header .badge { + background: var(--primary); + color: white; + padding: 0.25rem 0.75rem; + border-radius: 9999px; + font-size: 0.75rem; + font-weight: 500; +} + +.model-selector { + display: flex; + gap: 0.5rem; + background: rgba(255,255,255,0.1); + padding: 0.25rem; + border-radius: var(--radius); +} + +.model-btn { + background: transparent; + border: none; + color: rgba(255,255,255,0.7); + padding: 0.5rem 1rem; + border-radius: 6px; + cursor: pointer; + font-size: 0.875rem; + font-weight: 500; + transition: all 0.2s; +} + +.model-btn:hover { + color: white; + background: rgba(255,255,255,0.1); +} + +.model-btn.active { + background: var(--primary); + color: white; +} + +.model-btn:disabled { + opacity: 0.4; + cursor: not-allowed; +} + +/* Layout */ +.container { + max-width: 1400px; + margin: 0 auto; + padding: 2rem; +} + +.tabs { + display: flex; + gap: 0.25rem; + background: white; + padding: 0.5rem; + border-radius: var(--radius); + box-shadow: var(--shadow); + margin-bottom: 1.5rem; + width: fit-content; +} + +.tab-btn { + background: none; + border: none; + padding: 0.625rem 1.25rem; + border-radius: 6px; + cursor: pointer; + font-size: 0.875rem; + font-weight: 500; + color: var(--gray-500); + transition: all 0.2s; + display: flex; + align-items: center; + gap: 0.5rem; +} + +.tab-btn:hover { + color: var(--gray-700); + background: var(--gray-100); +} + +.tab-btn.active { + background: var(--primary-light); + color: var(--primary); +} + +.tab-content { + display: none; +} + +.tab-content.active { + display: block; +} + +/* Cards */ +.card { + background: white; + border-radius: var(--radius); + box-shadow: var(--shadow); + border: 1px solid var(--gray-200); + overflow: hidden; +} + +.card-header { + padding: 1.25rem 1.5rem; + border-bottom: 1px solid var(--gray-200); + display: flex; + align-items: center; + justify-content: space-between; +} + +.card-header h2 { + font-size: 1rem; + font-weight: 600; + color: var(--gray-800); + display: flex; + align-items: center; + gap: 0.5rem; +} + +.card-header .step-number { + background: var(--primary); + color: white; + width: 24px; + height: 24px; + border-radius: 50%; + display: flex; + align-items: center; + justify-content: center; + font-size: 0.75rem; + font-weight: 600; +} + +.card-body { + padding: 1.5rem; +} + +/* Grid layouts */ +.grid-2 { + display: grid; + grid-template-columns: repeat(2, 1fr); + gap: 1.5rem; +} + +.grid-3 { + display: grid; + grid-template-columns: repeat(3, 1fr); + gap: 1.5rem; +} + +@media (max-width: 1024px) { + .grid-2, .grid-3 { + grid-template-columns: 1fr; + } +} + +/* Forms */ +.form-group { + margin-bottom: 1.25rem; +} + +.form-group label { + display: block; + font-size: 0.875rem; + font-weight: 500; + color: var(--gray-700); + margin-bottom: 0.375rem; +} + +.form-group .hint { + font-size: 0.75rem; + color: var(--gray-400); + margin-top: 0.25rem; +} + +input[type="text"], +input[type="number"], +input[type="file"], +select, +textarea { + width: 100%; + padding: 0.625rem 0.875rem; + border: 1px solid var(--gray-300); + border-radius: var(--radius); + font-size: 0.875rem; + color: var(--gray-800); + background: white; + transition: border-color 0.2s, box-shadow 0.2s; +} + +input:focus, +select:focus, +textarea:focus { + outline: none; + border-color: var(--primary); + box-shadow: 0 0 0 3px var(--primary-light); +} + +textarea { + resize: vertical; + min-height: 120px; + font-family: inherit; +} + +input[type="file"] { + padding: 0.5rem; + cursor: pointer; +} + +/* Buttons */ +.btn { + display: inline-flex; + align-items: center; + justify-content: center; + gap: 0.5rem; + padding: 0.625rem 1.25rem; + border-radius: var(--radius); + font-size: 0.875rem; + font-weight: 500; + cursor: pointer; + border: none; + transition: all 0.2s; + white-space: nowrap; +} + +.btn:disabled { + opacity: 0.5; + cursor: not-allowed; +} + +.btn-primary { + background: var(--primary); + color: white; +} + +.btn-primary:hover:not(:disabled) { + background: var(--primary-dark); +} + +.btn-success { + background: var(--success); + color: white; +} + +.btn-success:hover:not(:disabled) { + background: #047857; +} + +.btn-danger { + background: var(--danger); + color: white; +} + +.btn-danger:hover:not(:disabled) { + background: #b91c1c; +} + +.btn-secondary { + background: var(--gray-100); + color: var(--gray-700); + border: 1px solid var(--gray-300); +} + +.btn-secondary:hover:not(:disabled) { + background: var(--gray-200); +} + +.btn-sm { + padding: 0.375rem 0.75rem; + font-size: 0.75rem; +} + +.btn-lg { + padding: 0.875rem 1.75rem; + font-size: 1rem; +} + +/* Upload zone */ +.upload-zone { + border: 2px dashed var(--gray-300); + border-radius: var(--radius); + padding: 2rem; + text-align: center; + background: var(--gray-50); + transition: all 0.2s; + cursor: pointer; +} + +.upload-zone:hover { + border-color: var(--primary); + background: var(--primary-light); +} + +.upload-zone.dragover { + border-color: var(--primary); + background: var(--primary-light); +} + +.upload-zone .icon { + font-size: 2rem; + margin-bottom: 0.5rem; + color: var(--gray-400); +} + +.upload-zone p { + color: var(--gray-500); + font-size: 0.875rem; +} + +/* Image gallery */ +.image-gallery { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); + gap: 1rem; + margin-top: 1rem; +} + +.image-item { + position: relative; + border-radius: var(--radius); + overflow: hidden; + border: 2px solid var(--gray-200); + aspect-ratio: 1; + background: var(--gray-100); +} + +.image-item img { + width: 100%; + height: 100%; + object-fit: cover; +} + +.image-item .overlay { + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: rgba(0,0,0,0.5); + display: flex; + align-items: center; + justify-content: center; + opacity: 0; + transition: opacity 0.2s; +} + +.image-item:hover .overlay { + opacity: 1; +} + +.image-item .badge-mask { + position: absolute; + top: 0.5rem; + right: 0.5rem; + background: var(--success); + color: white; + font-size: 0.625rem; + padding: 0.125rem 0.375rem; + border-radius: 4px; + font-weight: 600; +} + +/* Canvas container */ +.canvas-container { + position: relative; + border: 2px solid var(--gray-300); + border-radius: var(--radius); + overflow: hidden; + background: var(--gray-100); + display: flex; + align-items: center; + justify-content: center; + min-height: 400px; +} + +.canvas-container canvas { + max-width: 100%; + max-height: 600px; + cursor: crosshair; +} + +.canvas-toolbar { + display: flex; + gap: 0.5rem; + padding: 0.75rem; + background: var(--gray-100); + border-bottom: 1px solid var(--gray-200); + align-items: center; + flex-wrap: wrap; +} + +.canvas-toolbar span { + font-size: 0.75rem; + color: var(--gray-500); + margin-left: auto; +} + +/* Log viewer */ +.log-viewer { + background: var(--gray-800); + color: #e2e8f0; + font-family: "SF Mono", Monaco, "Cascadia Code", monospace; + font-size: 0.75rem; + line-height: 1.6; + padding: 1rem; + border-radius: var(--radius); + height: 400px; + overflow-y: auto; + white-space: pre-wrap; + word-break: break-all; +} + +.log-viewer .log-line { + padding: 0.125rem 0; +} + +.log-viewer .log-error { + color: #fca5a5; +} + +.log-viewer .log-success { + color: #86efac; +} + +.log-viewer .log-warn { + color: #fcd34d; +} + +/* Progress bar */ +.progress-bar { + width: 100%; + height: 8px; + background: var(--gray-200); + border-radius: 9999px; + overflow: hidden; +} + +.progress-bar .fill { + height: 100%; + background: var(--primary); + transition: width 0.3s ease; + border-radius: 9999px; +} + +/* Status badges */ +.status-badge { + display: inline-flex; + align-items: center; + gap: 0.375rem; + padding: 0.375rem 0.75rem; + border-radius: 9999px; + font-size: 0.75rem; + font-weight: 500; +} + +.status-badge.running { + background: var(--primary-light); + color: var(--primary); +} + +.status-badge.completed { + background: var(--success-light); + color: var(--success); +} + +.status-badge.error { + background: var(--danger-light); + color: var(--danger); +} + +.status-badge.idle { + background: var(--gray-100); + color: var(--gray-500); +} + +/* Preview modal */ +.modal-overlay { + position: fixed; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: rgba(0,0,0,0.75); + display: flex; + align-items: center; + justify-content: center; + z-index: 1000; + padding: 2rem; +} + +.modal-content { + background: white; + border-radius: var(--radius); + max-width: 90vw; + max-height: 90vh; + overflow: auto; + position: relative; +} + +.modal-close { + position: absolute; + top: 1rem; + right: 1rem; + background: rgba(0,0,0,0.5); + color: white; + border: none; + width: 32px; + height: 32px; + border-radius: 50%; + cursor: pointer; + font-size: 1rem; + display: flex; + align-items: center; + justify-content: center; +} + +.modal-content img { + max-width: 100%; + max-height: 80vh; + display: block; +} + +/* Generation results */ +.result-grid { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(280px, 1fr)); + gap: 1.5rem; + margin-top: 1.5rem; +} + +.result-card { + background: white; + border-radius: var(--radius); + overflow: hidden; + box-shadow: var(--shadow); + border: 1px solid var(--gray-200); +} + +.result-card img { + width: 100%; + height: 200px; + object-fit: cover; + cursor: pointer; + transition: transform 0.2s; +} + +.result-card img:hover { + transform: scale(1.02); +} + +.result-card .info { + padding: 1rem; +} + +.result-card .info h4 { + font-size: 0.875rem; + font-weight: 600; + margin-bottom: 0.25rem; +} + +.result-card .info p { + font-size: 0.75rem; + color: var(--gray-500); +} + +/* Domain instruction examples */ +.example-box { + background: var(--gray-50); + border: 1px solid var(--gray-200); + border-radius: var(--radius); + padding: 1rem; + margin-bottom: 1rem; +} + +.example-box h4 { + font-size: 0.875rem; + font-weight: 600; + margin-bottom: 0.5rem; + color: var(--gray-700); +} + +.example-box pre { + background: var(--gray-800); + color: #e2e8f0; + padding: 0.75rem; + border-radius: 6px; + font-size: 0.75rem; + overflow-x: auto; + white-space: pre-wrap; +} + +/* Alert messages */ +.alert { + padding: 1rem; + border-radius: var(--radius); + margin-bottom: 1rem; + font-size: 0.875rem; + display: flex; + align-items: flex-start; + gap: 0.75rem; +} + +.alert-info { + background: var(--primary-light); + color: var(--primary-dark); + border: 1px solid #bfdbfe; +} + +.alert-success { + background: var(--success-light); + color: var(--success); + border: 1px solid #a7f3d0; +} + +.alert-warning { + background: var(--warning-light); + color: var(--warning); + border: 1px solid #fde68a; +} + +.alert-danger { + background: var(--danger-light); + color: var(--danger); + border: 1px solid #fecaca; +} + +/* Utility */ +.hidden { + display: none !important; +} + +.mt-1 { margin-top: 0.5rem; } +.mt-2 { margin-top: 1rem; } +.mt-3 { margin-top: 1.5rem; } +.mb-1 { margin-bottom: 0.5rem; } +.mb-2 { margin-bottom: 1rem; } +.flex { display: flex; } +.flex-col { flex-direction: column; } +.items-center { align-items: center; } +.justify-between { justify-content: space-between; } +.gap-1 { gap: 0.5rem; } +.gap-2 { gap: 1rem; } +.w-full { width: 100%; } +.text-sm { font-size: 0.875rem; } +.text-xs { font-size: 0.75rem; } +.text-gray { color: var(--gray-500); } +.font-mono { font-family: monospace; } + +/* Spinner */ +@keyframes spin { + to { transform: rotate(360deg); } +} + +.spinner { + width: 16px; + height: 16px; + border: 2px solid rgba(255,255,255,0.3); + border-top-color: white; + border-radius: 50%; + animation: spin 0.8s linear infinite; + display: inline-block; +} + +.spinner-dark { + border-color: var(--gray-300); + border-top-color: var(--primary); +} + +/* Responsive */ +@media (max-width: 768px) { + .container { + padding: 1rem; + } + .header { + padding: 1rem; + flex-direction: column; + gap: 1rem; + } + .tabs { + width: 100%; + overflow-x: auto; + } +} + +/* ============================================ + Result Groups - Collapsible Accordion + ============================================ */ + +.result-group { + border: 1px solid var(--gray-200); + border-radius: var(--radius); + margin-bottom: 0.75rem; + overflow: hidden; + background: white; +} + +.result-group-header { + display: flex; + align-items: center; + padding: 0.875rem 1.25rem; + background: var(--gray-50); + cursor: pointer; + user-select: none; + transition: background 0.2s; + border-bottom: 1px solid transparent; +} + +.result-group-header:hover { + background: var(--gray-100); +} + +.result-group-header.active { + background: var(--primary-light); + border-bottom-color: var(--primary); +} + +.result-group-header .group-title { + font-weight: 600; + font-size: 0.875rem; + color: var(--gray-800); +} + +.result-group-header .group-meta { + font-size: 0.75rem; + color: var(--gray-500); + margin-left: 0.75rem; +} + +.result-group-header .toggle-icon { + font-size: 0.875rem; + color: var(--gray-600); + margin-left: auto; + transition: color 0.2s; +} + +.result-group-header:hover .toggle-icon { + color: var(--primary); +} + +.result-group-content { + display: none; + background: white; +} + +.result-group-content.expanded { + display: block; + padding: 1rem; +} + +.result-group-content .result-grid { + margin-top: 0; + grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); + gap: 1rem; +} + +/* ============================================ + Result Defect Groups - Second Level + ============================================ */ + +.result-defect-group { + border: 1px solid var(--gray-200); + border-radius: var(--radius); + margin-bottom: 0.5rem; + overflow: hidden; +} + +.result-defect-header { + display: flex; + align-items: center; + padding: 0.625rem 1rem; + background: var(--gray-50); + cursor: pointer; + user-select: none; + transition: background 0.2s; + border-bottom: 1px solid transparent; +} + +.result-defect-header:hover { + background: var(--gray-100); +} + +.result-defect-header.active { + background: #e0e7ff; + border-bottom-color: var(--primary); +} + +.result-defect-header .defect-title { + font-weight: 600; + font-size: 0.8125rem; + color: var(--gray-800); +} + +.result-defect-header .defect-meta { + font-size: 0.75rem; + color: var(--gray-500); + margin-left: 0.75rem; +} + +.result-defect-content { + display: none; + background: white; + padding: 0.75rem; +} + +.result-defect-content.expanded { + display: block; +} + +/* ============================================ + Result Trio Grid - Clean / Mask / Output + ============================================ */ + +.result-trio-grid { + display: grid; + grid-template-columns: repeat(4, 1fr); + gap: 0.75rem; +} + +.result-trio-card { + background: white; + border-radius: var(--radius); + overflow: hidden; + border: 1px solid var(--gray-200); + text-align: center; +} + +.result-trio-card .trio-label { + font-size: 0.75rem; + font-weight: 600; + color: var(--gray-600); + padding: 0.375rem; + background: var(--gray-100); + border-bottom: 1px solid var(--gray-200); +} + +.result-trio-card img { + width: 100%; + height: 160px; + object-fit: cover; + cursor: pointer; + transition: transform 0.2s; +} + +.result-trio-card img:hover { + transform: scale(1.02); +} + +.result-trio-card .trio-filename { + font-size: 0.6875rem; + color: var(--gray-500); + padding: 0.375rem; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + +@media (max-width: 768px) { + .result-trio-grid { + grid-template-columns: 1fr; + } +} diff --git a/static/js/app.js b/static/js/app.js new file mode 100644 index 0000000000000000000000000000000000000000..1f668039a6a7a2e50998ad3a6cadf9aaf96248a3 --- /dev/null +++ b/static/js/app.js @@ -0,0 +1,1647 @@ +/** + * ArtiAgent Web UI - Unified Frontend Logic + * Supports both DefectFill & DefectDiffu + * ========================================== + */ + +// ============================================================================= +// Global State +// ============================================================================= + +const state = { + currentModel: 'defectfill', + activeTab: 'training', + uploadedImages: [], + currentMaskImage: null, + drawnRect: null, + drawnRectNatural: null, + isDrawing: false, + trainingPollInterval: null, + checkpoints: {}, + generationResults: [], + savedMasks: new Set(), + ddConfig: {}, +}; + +// ============================================================================= +// Initialization +// ============================================================================= + +document.addEventListener('DOMContentLoaded', () => { + initTabs(); + initModelSelector(); + initTrainingWizard(); + initGenerationForm(); + initDefectDiffu(); + initCanvas(); + loadCheckpoints(); + loadSystemInfo(); + loadDDConfig(); + setupDomainExamples(); +}); + +// ============================================================================= +// Navigation & Tabs +// ============================================================================= + +function initTabs() { + const tabBtns = document.querySelectorAll('.tab-btn'); + const tabContents = document.querySelectorAll('.tab-content'); + + tabBtns.forEach(btn => { + btn.addEventListener('click', () => { + const target = btn.dataset.tab; + state.activeTab = target; + + tabBtns.forEach(b => b.classList.remove('active')); + tabContents.forEach(c => c.classList.remove('active')); + + btn.classList.add('active'); + document.getElementById(`tab-${target}`).classList.add('active'); + }); + }); +} + +function initModelSelector() { + const modelBtns = document.querySelectorAll('.model-btn'); + modelBtns.forEach(btn => { + btn.addEventListener('click', () => { + const model = btn.dataset.model; + state.currentModel = model; + + modelBtns.forEach(b => b.classList.remove('active')); + btn.classList.add('active'); + + const defectfillTabs = document.getElementById('defectfill-tabs'); + const defectdiffuTabs = document.getElementById('defectdiffu-tabs'); + + if (model === 'defectfill') { + defectfillTabs.classList.remove('hidden'); + defectdiffuTabs.classList.add('hidden'); + } else { + defectfillTabs.classList.add('hidden'); + defectdiffuTabs.classList.remove('hidden'); + } + }); + }); +} + +// ============================================================================= +// Training Wizard (DefectFill) +// ============================================================================= + +function initTrainingWizard() { + const btnStep1 = document.getElementById('btn-step1'); + btnStep1.addEventListener('click', () => { + const objClass = document.getElementById('train-object-class').value.trim(); + const defectType = document.getElementById('train-defect-type').value.trim(); + + if (!objClass || !defectType) { + showAlert('step1-alert', 'Please enter both product name and defect type.', 'danger'); + return; + } + + document.getElementById('step2').classList.remove('hidden'); + document.getElementById('step2').scrollIntoView({ behavior: 'smooth' }); + checkGoodImages(objClass); + }); + + const goodUploadZone = document.getElementById('good-upload-zone'); + const goodFileInput = document.getElementById('good-images'); + + goodUploadZone.addEventListener('click', () => goodFileInput.click()); + goodUploadZone.addEventListener('dragover', (e) => { + e.preventDefault(); + goodUploadZone.classList.add('dragover'); + }); + goodUploadZone.addEventListener('dragleave', () => goodUploadZone.classList.remove('dragover')); + goodUploadZone.addEventListener('drop', (e) => { + e.preventDefault(); + goodUploadZone.classList.remove('dragover'); + handleGoodFileSelect(e.dataTransfer.files); + }); + + goodFileInput.addEventListener('change', (e) => handleGoodFileSelect(e.target.files)); + document.getElementById('btn-upload-good').addEventListener('click', () => { + const files = goodFileInput.files; + if (files.length === 0) { + showAlert('step2-alert', 'Please select files first.', 'warning'); + return; + } + handleGoodFileSelect(files); + }); + document.getElementById('btn-use-existing-good').addEventListener('click', () => { + document.getElementById('step3').classList.remove('hidden'); + document.getElementById('step3').scrollIntoView({ behavior: 'smooth' }); + }); + document.getElementById('btn-proceed-good').addEventListener('click', () => { + document.getElementById('step3').classList.remove('hidden'); + document.getElementById('step3').scrollIntoView({ behavior: 'smooth' }); + }); + + const defectUploadZone = document.getElementById('defect-upload-zone'); + const defectFileInput = document.getElementById('defect-images'); + + defectUploadZone.addEventListener('click', () => defectFileInput.click()); + defectUploadZone.addEventListener('dragover', (e) => { + e.preventDefault(); + defectUploadZone.classList.add('dragover'); + }); + defectUploadZone.addEventListener('dragleave', () => defectUploadZone.classList.remove('dragover')); + defectUploadZone.addEventListener('drop', (e) => { + e.preventDefault(); + defectUploadZone.classList.remove('dragover'); + handleFileSelect(e.dataTransfer.files); + }); + + defectFileInput.addEventListener('change', (e) => handleFileSelect(e.target.files)); + document.getElementById('btn-step3-upload').addEventListener('click', () => { + const files = defectFileInput.files; + if (files.length === 0) { + showAlert('step3-alert', 'Please select files first.', 'warning'); + return; + } + handleFileSelect(files); + }); + + document.getElementById('btn-align-yes').addEventListener('click', () => { + document.getElementById('align-reference-section').classList.remove('hidden'); + }); + document.getElementById('btn-align-no').addEventListener('click', skipAlignment); + + const refUploadZone = document.getElementById('reference-upload-zone'); + const refFileInput = document.getElementById('reference-image'); + refUploadZone.addEventListener('click', () => refFileInput.click()); + refFileInput.addEventListener('change', (e) => { + if (e.target.files && e.target.files.length > 0) { + renderReferenceImage(e.target.files[0]); + } + }); + document.getElementById('btn-align-proceed').addEventListener('click', runAlignment); + + document.getElementById('btn-save-mask').addEventListener('click', saveMask); + document.getElementById('btn-skip-mask').addEventListener('click', skipMask); + document.getElementById('btn-prev-image').addEventListener('click', prevMaskImage); + document.getElementById('btn-next-image').addEventListener('click', nextMaskImage); + document.getElementById('btn-clear-rect').addEventListener('click', clearRect); + document.getElementById('btn-done-masks').addEventListener('click', finishMaskStep); + + document.getElementById('btn-save-domain').addEventListener('click', saveDomainInstructions); + document.getElementById('btn-load-domain-example').addEventListener('click', loadDomainExample); + + document.getElementById('btn-start-training').addEventListener('click', startTraining); + document.getElementById('btn-stop-training').addEventListener('click', stopTraining); +} + +function handleFileSelect(files) { + const objClass = document.getElementById('train-object-class').value.trim(); + const defectType = document.getElementById('train-defect-type').value.trim(); + + if (!objClass || !defectType) { + showAlert('step3-alert', 'Please complete Step 1 first.', 'warning'); + return; + } + + if (files.length < 1) { + showAlert('step3-alert', 'At least 1 images required for few-shot training.', 'warning'); + return; + } + + const formData = new FormData(); + formData.append('object_class', objClass); + formData.append('defect_type', defectType); + + for (const file of files) { + if (file.type.startsWith('image/')) { + formData.append('images', file); + } + } + + const btn = document.getElementById('btn-step3-upload'); + btn.disabled = true; + btn.innerHTML = ' Uploading...'; + + fetch('/api/training/upload-images', { + method: 'POST', + body: formData + }) + .then(r => r.json()) + .then(data => { + btn.disabled = false; + btn.innerHTML = '๐Ÿ“ค Upload Images'; + + if (data.success) { + state.uploadedImages = data.saved; + renderUploadedImages(); + showAlert('step3-alert', `Uploaded ${data.count} images successfully.`, 'success'); + document.getElementById('step3-5').classList.remove('hidden'); + document.getElementById('step3-5').scrollIntoView({ behavior: 'smooth' }); + } else { + showAlert('step3-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + btn.innerHTML = '๐Ÿ“ค Upload Images'; + showAlert('step3-alert', 'Upload failed: ' + err.message, 'danger'); + }); +} + +function checkProductExists(objClass, defectType) { + fetch(`/api/check-product-exists?object_class=${encodeURIComponent(objClass)}`) + .then(r => r.json()) + .then(data => { + if (data.success && data.exists) { + showDuplicateProductModal(objClass, data.existing_defect_types, data.sample_image_url, defectType); + } else { + proceedToStep2(); + } + }) + .catch(err => { + console.error('Error checking product existence:', err); + proceedToStep2(); + }); +} + +function showDuplicateProductModal(productName, existingDefectTypes, sampleImageUrl, defectType) { + const modal = document.getElementById('duplicate-product-modal'); + document.getElementById('dup-modal-product-name').textContent = productName; + document.getElementById('dup-modal-defect-types').textContent = + existingDefectTypes.length > 0 ? existingDefectTypes.join(', ') : 'None'; + + const imgContainer = document.getElementById('dup-modal-image-container'); + const img = document.getElementById('dup-modal-sample-image'); + if (sampleImageUrl) { + img.src = sampleImageUrl; + imgContainer.classList.remove('hidden'); + } else { + imgContainer.classList.add('hidden'); + } + + modal.classList.remove('hidden'); + + document.getElementById('btn-add-defect-type').onclick = () => { + modal.classList.add('hidden'); + proceedToStep2(); + }; + + document.getElementById('btn-new-product').onclick = () => { + modal.classList.add('hidden'); + promptForNewProductName(productName); + }; + + document.getElementById('btn-close-dup-modal').onclick = () => { + modal.classList.add('hidden'); + }; + + modal.onclick = (e) => { + if (e.target === modal) modal.classList.add('hidden'); + }; +} + +function promptForNewProductName(originalName) { + const newName = prompt( + `Product "${originalName}" already exists.\n\nPlease enter a new, unique product name:` + ); + if (!newName || !newName.trim()) { + showAlert('step1-alert', 'A new product name is required to proceed.', 'warning'); + return; + } + const trimmed = newName.trim(); + + fetch(`/api/check-product-exists?object_class=${encodeURIComponent(trimmed)}`) + .then(r => r.json()) + .then(data => { + if (data.success && data.exists) { + showAlert('step1-alert', `Product name "${trimmed}" also already exists. Please choose another name.`, 'danger'); + setTimeout(() => promptForNewProductName(originalName), 300); + } else { + document.getElementById('train-object-class').value = trimmed; + showAlert('step1-alert', `Product name updated to "${trimmed}". Proceeding...`, 'success'); + proceedToStep2(); + } + }) + .catch(err => { + showAlert('step1-alert', 'Error checking new name: ' + err.message, 'danger'); + }); +} + +function checkGoodImages(objectClass) { + fetch(`/api/training/good-images/count?object_class=${encodeURIComponent(objectClass)}`) + .then(r => r.json()) + .then(data => { + const statusDiv = document.getElementById('good-images-status'); + const useExistingBtn = document.getElementById('btn-use-existing-good'); + const proceedBtn = document.getElementById('btn-proceed-good'); + + if (data.success && data.count > 0) { + statusDiv.className = 'alert alert-info'; + statusDiv.innerHTML = `${data.count} good image(s) already exist for "${escapeHtml(objectClass)}".
You can upload more or proceed with existing ones.`; + statusDiv.classList.remove('hidden'); + useExistingBtn.classList.remove('hidden'); + proceedBtn.classList.remove('hidden'); + } else { + statusDiv.className = 'alert alert-warning'; + statusDiv.innerHTML = `No good images found for "${escapeHtml(objectClass)}".
Please upload at least 1 clean (non-defective) image.`; + statusDiv.classList.remove('hidden'); + useExistingBtn.classList.add('hidden'); + proceedBtn.classList.add('hidden'); + } + }) + .catch(err => { + showAlert('step2-alert', 'Failed to check good images: ' + err.message, 'danger'); + }); +} + +function handleGoodFileSelect(files) { + const objClass = document.getElementById('train-object-class').value.trim(); + + if (!objClass) { + showAlert('step2-alert', 'Please complete Step 1 first.', 'warning'); + return; + } + + const formData = new FormData(); + formData.append('object_class', objClass); + + for (const file of files) { + if (file.type.startsWith('image/')) { + formData.append('images', file); + } + } + + const btn = document.getElementById('btn-upload-good'); + btn.disabled = true; + btn.innerHTML = ' Uploading...'; + + fetch('/api/training/upload-good-images', { + method: 'POST', + body: formData + }) + .then(r => r.json()) + .then(data => { + btn.disabled = false; + btn.innerHTML = '๐Ÿ“ค Upload Good Images'; + + if (data.success) { + renderGoodImages(data.saved); + showAlert('step2-alert', `Uploaded ${data.count} good images. Total in folder: ${data.total_in_folder}`, 'success'); + document.getElementById('btn-proceed-good').classList.remove('hidden'); + document.getElementById('btn-use-existing-good').classList.remove('hidden'); + } else { + showAlert('step2-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + btn.innerHTML = '๐Ÿ“ค Upload Good Images'; + showAlert('step2-alert', 'Upload failed: ' + err.message, 'danger'); + }); +} + +function renderGoodImages(images) { + const gallery = document.getElementById('good-images-gallery'); + images.forEach(img => { + const div = document.createElement('div'); + div.className = 'image-item'; + div.innerHTML = `${img.filename}`; + gallery.appendChild(div); + }); +} + +function proceedToStep2() { + document.getElementById('step2').classList.remove('hidden'); + document.getElementById('step2').scrollIntoView({ behavior: 'smooth' }); +} + +function renderUploadedImages() { + const gallery = document.getElementById('defect-upload-gallery'); + gallery.innerHTML = ''; + + state.uploadedImages.forEach((img, idx) => { + const div = document.createElement('div'); + div.className = 'image-item'; + div.innerHTML = ` + ${img.filename} + +
+ +
+ `; + gallery.appendChild(div); + }); + + updateMaskCounter(); +} + +function renderReferenceImage(file) { + const gallery = document.getElementById('reference-gallery'); + gallery.innerHTML = ''; + const div = document.createElement('div'); + div.className = 'image-item'; + const url = URL.createObjectURL(file); + div.innerHTML = `Reference`; + gallery.appendChild(div); +} + +function skipAlignment() { + const objClass = document.getElementById('train-object-class').value.trim(); + const defectType = document.getElementById('train-defect-type').value.trim(); + const btn = document.getElementById('btn-align-no'); + btn.disabled = true; + + fetch('/api/training/skip-align', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ object_class: objClass }) + }) + .then(r => r.json()) + .then(data => { + btn.disabled = false; + if (data.success) { + if (data.images) { + const defectImages = data.images.filter(img => + img.path.includes(`/defective/${encodeURIComponent(defectType)}`) + ); + + if (defectImages.length > 0) { + state.uploadedImages = defectImages; + renderUploadedImages(); + } + } + showAlert('step3-5-alert', `Moved ${data.moved_count} images to final location.`, 'success'); + proceedToStep4(); + } else { + showAlert('step3-5-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + showAlert('step3-5-alert', 'Error: ' + err.message, 'danger'); + }); +} + +function runAlignment() { + const objClass = document.getElementById('train-object-class').value.trim(); + const defectType = document.getElementById('train-defect-type').value.trim(); + const refInput = document.getElementById('reference-image'); + + if (!refInput.files || refInput.files.length === 0) { + showAlert('step3-5-alert', 'Please upload a reference image first.', 'warning'); + return; + } + + const formData = new FormData(); + formData.append('object_class', objClass); + formData.append('reference', refInput.files[0]); + + const btn = document.getElementById('btn-align-proceed'); + const spinner = document.getElementById('align-spinner'); + btn.disabled = true; + spinner.classList.remove('hidden'); + + fetch('/api/training/align', { + method: 'POST', + body: formData + }) + .then(r => r.json()) + .then(data => { + btn.disabled = false; + spinner.classList.add('hidden'); + + if (data.success) { + if (data.images) { + const defectImages = data.images.filter(img => + img.path.includes(`/defective/${encodeURIComponent(defectType)}`) + ); + + if (defectImages.length > 0) { + state.uploadedImages = defectImages; + renderUploadedImages(); + } + } + + let msg = `Aligned ${data.aligned_count} images successfully.`; + if (data.error_count > 0) { + msg += ` ${data.error_count} failed and were copied as-is.`; + } + showAlert('step3-5-alert', msg, 'success'); + proceedToStep4(); + } else { + showAlert('step3-5-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + spinner.classList.add('hidden'); + showAlert('step3-5-alert', 'Alignment failed: ' + err.message, 'danger'); + }); +} + +function proceedToStep4() { + document.getElementById('step4').classList.remove('hidden'); + document.getElementById('step4').scrollIntoView({ behavior: 'smooth' }); + initMaskCanvas(0); +} + +// ============================================================================= +// Canvas Mask Drawing +// ============================================================================= + +function initCanvas() { + // Canvas is initialized per-image +} + +function initMaskCanvas(index) { + if (index < 0 || index >= state.uploadedImages.length) return; + + state.currentMaskImage = index; + state.drawnRect = null; + state.drawnRectNatural = null; + + const imgData = state.uploadedImages[index]; + const canvas = document.getElementById('mask-canvas'); + const ctx = canvas.getContext('2d'); + const img = new Image(); + + img.onload = () => { + const maxW = 800; + const maxH = 600; + let w = img.width; + let h = img.height; + + if (w > maxW || h > maxH) { + const ratio = Math.min(maxW / w, maxH / h); + w = Math.floor(w * ratio); + h = Math.floor(h * ratio); + } + + canvas.width = w; + canvas.height = h; + canvas.dataset.naturalWidth = img.width; + canvas.dataset.naturalHeight = img.height; + + ctx.drawImage(img, 0, 0, w, h); + updateMaskCounter(); + document.getElementById('step4-canvas').classList.remove('hidden'); + }; + + img.src = imgData.url; + + let startX, startY; + + function getCanvasCoords(e) { + const rect = canvas.getBoundingClientRect(); + const scaleX = canvas.width / rect.width; + const scaleY = canvas.height / rect.height; + return { + x: (e.clientX - rect.left) * scaleX, + y: (e.clientY - rect.top) * scaleY + }; + } + + canvas.onmousedown = (e) => { + const pos = getCanvasCoords(e); + startX = pos.x; + startY = pos.y; + state.isDrawing = true; + state.drawnRect = null; + state.drawnRectNatural = null; + }; + + canvas.onmousemove = (e) => { + if (!state.isDrawing) return; + const pos = getCanvasCoords(e); + const x = pos.x; + const y = pos.y; + + ctx.drawImage(img, 0, 0, canvas.width, canvas.height); + + ctx.strokeStyle = '#2563eb'; + ctx.lineWidth = 2; + ctx.setLineDash([5, 5]); + ctx.strokeRect(startX, startY, x - startX, y - startY); + + ctx.fillStyle = 'rgba(37, 99, 235, 0.2)'; + ctx.fillRect(startX, startY, x - startX, y - startY); + }; + + canvas.onmouseup = (e) => { + if (!state.isDrawing) return; + state.isDrawing = false; + + const pos = getCanvasCoords(e); + const endX = pos.x; + const endY = pos.y; + + const x = Math.min(startX, endX); + const y = Math.min(startY, endY); + const width = Math.abs(endX - startX); + const height = Math.abs(endY - startY); + + if (width > 10 && height > 10) { + state.drawnRect = { x, y, width, height }; + + const scaleX = img.width / canvas.width; + const scaleY = img.height / canvas.height; + state.drawnRectNatural = { + x: Math.round(x * scaleX), + y: Math.round(y * scaleY), + width: Math.round(width * scaleX), + height: Math.round(height * scaleY) + }; + } + + ctx.setLineDash([]); + }; + + canvas.onmouseleave = () => { + if (state.isDrawing) { + state.isDrawing = false; + ctx.drawImage(img, 0, 0, canvas.width, canvas.height); + if (state.drawnRect) { + ctx.strokeStyle = '#2563eb'; + ctx.lineWidth = 2; + ctx.strokeRect(state.drawnRect.x, state.drawnRect.y, + state.drawnRect.width, state.drawnRect.height); + } + } + }; +} + +function toggleInputMode() { + const mode = document.querySelector('input[name="input_mode"]:checked').value; + const singleInput = document.getElementById('gen-clean-image'); + const folderInput = document.getElementById('gen-clean-folder'); + const singleDiv = document.getElementById('single-image-input'); + const folderDiv = document.getElementById('folder-input'); + const hint = document.getElementById('input-hint'); + + if (mode === 'folder') { + singleDiv.classList.add('hidden'); + folderDiv.classList.remove('hidden'); + singleInput.value = ''; + hint.textContent = 'Select a folder containing clean (non-defective) product images. All images will be processed.'; + } else { + singleDiv.classList.remove('hidden'); + folderDiv.classList.add('hidden'); + folderInput.value = ''; + hint.textContent = 'Upload a clean (non-defective) image of the product.'; + } +} + +function toggleDDInputMode() { + const mode = document.querySelector('input[name="dd_input_mode"]:checked').value; + const singleDiv = document.getElementById('dd-single-image-input'); + const folderDiv = document.getElementById('dd-folder-input'); + const singleInput = document.getElementById('dd-clean-image'); + const folderInput = document.getElementById('dd-clean-folder'); + const hint = document.getElementById('dd-input-hint'); + + if (mode === 'folder') { + singleDiv.classList.add('hidden'); + folderDiv.classList.remove('hidden'); + singleInput.value = ''; + hint.textContent = 'Select a folder containing clean (non-defective) product images. All images will be processed.'; + } else { + singleDiv.classList.remove('hidden'); + folderDiv.classList.add('hidden'); + folderInput.value = ''; + hint.textContent = 'Upload a clean (non-defective) image of the product.'; + } +} + +function clearRect() { + state.drawnRect = null; + state.drawnRectNatural = null; + const canvas = document.getElementById('mask-canvas'); + const ctx = canvas.getContext('2d'); + const img = new Image(); + img.onload = () => { + ctx.drawImage(img, 0, 0, canvas.width, canvas.height); + }; + if (state.uploadedImages[state.currentMaskImage]) { + img.src = state.uploadedImages[state.currentMaskImage].url; + } +} + +function saveMask() { + if (!state.drawnRectNatural) { + showAlert('step4-alert', 'Please draw a rectangle on the image first.', 'warning'); + return; + } + + const imgData = state.uploadedImages[state.currentMaskImage]; + + fetch('/api/training/save-mask', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + image_path: imgData.path, + rect: state.drawnRectNatural + }) + }) + .then(r => r.json()) + .then(data => { + if (data.success) { + state.savedMasks.add(state.currentMaskImage); + const badge = document.getElementById(`mask-badge-${state.currentMaskImage}`); + if (badge) badge.classList.remove('hidden'); + updateMaskCounter(); + + const totalSaved = state.savedMasks.size; + const remaining = state.uploadedImages.length - totalSaved; + + if (totalSaved >= 3) { + showAlert('step4-alert', + `Mask saved! Please click the Done button to continue to Step 5 (Remaining unmasked images: ${remaining})`, + 'success', + 0); + } else { + showAlert('step4-alert', 'Mask saved!', 'success', 1200); + } + + if (state.currentMaskImage < state.uploadedImages.length - 1) { + setTimeout(() => nextMaskImage(), 500); + } + } else { + showAlert('step4-alert', data.error, 'danger'); + } + }) + .catch(err => { + showAlert('step4-alert', 'Error: ' + err.message, 'danger'); + }); +} + +function skipMask() { + if (state.currentMaskImage < state.uploadedImages.length - 1) { + nextMaskImage(); + } +} + +function prevMaskImage() { + if (state.currentMaskImage > 0) { + initMaskCanvas(state.currentMaskImage - 1); + } +} + +function nextMaskImage() { + if (state.currentMaskImage < state.uploadedImages.length - 1) { + initMaskCanvas(state.currentMaskImage + 1); + } +} + +function updateMaskCounter() { + const counter = document.getElementById('mask-counter'); + const total = state.uploadedImages.length; + const masked = state.savedMasks.size; + const current = state.currentMaskImage !== null ? state.currentMaskImage + 1 : 0; + const filename = state.uploadedImages[state.currentMaskImage]?.filename || ''; + counter.textContent = `Image ${current} of ${total}: ${filename} | Masks saved: ${masked}/${total}`; +} + +function finishMaskStep() { + if (state.savedMasks.size < 1) { + showAlert('step4-alert', `You must save masks for at least 1 images. Currently saved: ${state.savedMasks.size}.`, 'warning'); + return; + } + document.getElementById('step5').classList.remove('hidden'); + document.getElementById('step5').scrollIntoView({ behavior: 'smooth' }); + document.getElementById('step4-alert').classList.add('hidden'); +} + +// ============================================================================= +// Domain Instructions +// ============================================================================= + +function setupDomainExamples() { + const examples = { + triac: `If the image is a standard TO-220 package (like a Triac): +- The component consists of three metallic pins/leads, a black plastic mold body, a metal mounting tab with a mounting hole, and printed silkscreen text. +- CRITICAL CONSTRAINT: You MUST exclusively propose defects located on the metallic pins/leads. Do NOT propose any defects on the black plastic body, the metal mounting tab, the mounting hole, or the silkscreen text. +- Acceptable pin/lead defects include: pin bending (crooked leads), lead deformation, surface scratches on the metallic pins, or contamination/solder dross on the lead surfaces. +- Use "3 pin" as target_entity rather than specific individual pins. +- The description must clearly detail exactly which pin is affected and how it is deformed or damaged (e.g., "The right pin is bent sharply outwards near the mid-section"). +- The defect_coverage_ratio should be 0.4 - 0.60.`, + + xray_pcb: `If the image is an X-ray or CT scan of a PCB/IC package: +- The die appears as a large dark or light rectangular region, usually in the center. +- Solder balls/BGA appear as an array of circular dots. +- Defects in X-ray appear as grayscale anomalies (darker or brighter regions), NOT as circuit traces, lines, or wire patterns. +- A "xray_die" defect should look like: a natural void, crack, or delamination INSIDE the die material โ€” blobby, irregular, or circular grayscale spots. +- A "xray_solder" defect should look like: missing solder, solder bridges, or voids in solder balls. +- NEVER describe the defect as "circuit traces", "wiring", "grid lines", or "PCB patterns". The defect must look like a manufacturing flaw, not an overlaid schematic. +- The defect_coverage_ratio should be 1.0.`, + + vcsel: `If the image is a VCSEL laser diode: +- The component has a metallic can package (TO-can), a glass lens cap on top, and electrode bars. +- Defects should focus on: lens cap cracks, scratches on the glass dome, electrode bar discontinuity, or metallic can surface damage. +- Use "lens cap" for lens-related defects and "electrode bars" for electrode-related defects. +- Avoid targeting the emission aperture directly with large coverage masks.`, + + generic: `If the image shows an electronic component: +- Identify the main functional areas: metallic leads/pins, plastic or ceramic body, and any optical or contact surfaces. +- Propose defects that are realistic manufacturing flaws: scratches, cracks, contamination, bent leads, missing material, or discoloration. +- Use specific terminology for metallic parts (e.g., "metallic pins", "copper leads", "gold-plated contacts") rather than generic terms like "leg" or "metal part". +- The defect must be visually distinct and not confused with normal component features.` + }; + + window.domainExamples = examples; +} + +function loadDomainExample() { + const select = document.getElementById('domain-example-select'); + const key = select.value; + if (window.domainExamples[key]) { + document.getElementById('domain-instructions').value = window.domainExamples[key]; + } +} + +function saveDomainInstructions() { + const objClass = document.getElementById('train-object-class').value.trim() || 'default'; + const instructions = document.getElementById('domain-instructions').value.trim(); + + fetch('/api/domain-instructions', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + object_class: objClass, + instructions: instructions + }) + }) + .then(r => r.json()) + .then(data => { + if (data.success) { + showAlert('step5-alert', data.message, 'success'); + document.getElementById('step6').classList.remove('hidden'); + document.getElementById('step6').scrollIntoView({ behavior: 'smooth' }); + } else { + showAlert('step5-alert', data.error, 'danger'); + } + }) + .catch(err => { + showAlert('step5-alert', 'Error: ' + err.message, 'danger'); + }); +} + +// ============================================================================= +// Training Execution (DefectFill) +// ============================================================================= + +function startTraining() { + const objClass = document.getElementById('train-object-class').value.trim(); + const defectType = document.getElementById('train-defect-type').value.trim(); + + const config = { + object_class: objClass, + defect_type: defectType, + lora_rank: parseInt(document.getElementById('cfg-lora-rank').value) || 8, + lora_alpha: parseInt(document.getElementById('cfg-lora-alpha').value) || 16, + max_train_steps: parseInt(document.getElementById('cfg-max-steps').value) || 1500, + batch_size: parseInt(document.getElementById('cfg-batch-size').value) || 2, + gradient_accumulation_steps: parseInt(document.getElementById('cfg-grad-accum').value) || 2, + lambda_defect: parseFloat(document.getElementById('cfg-lambda-defect').value) || 0.5, + lambda_obj: parseFloat(document.getElementById('cfg-lambda-obj').value) || 0.2, + lambda_attn: parseFloat(document.getElementById('cfg-lambda-attn').value) || 0.05, + alpha: parseFloat(document.getElementById('cfg-alpha').value) || 0.3 + }; + + const btn = document.getElementById('btn-start-training'); + btn.disabled = true; + btn.innerHTML = ' Starting Training...'; + + fetch('/api/training/start', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(config) + }) + .then(r => r.json()) + .then(data => { + if (data.success) { + document.getElementById('training-monitor').classList.remove('hidden'); + document.getElementById('training-monitor').scrollIntoView({ + behavior: 'smooth', + block: 'start' + }); + showAlert('step6-alert', 'Training started! Monitor logs below.', 'success'); + startLogPolling(); + } else { + btn.disabled = false; + btn.innerHTML = '๐Ÿš€ Start Training'; + showAlert('step6-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + btn.innerHTML = '๐Ÿš€ Start Training'; + showAlert('step6-alert', 'Error: ' + err.message, 'danger'); + }); +} + +function stopTraining() { + fetch('/api/training/stop', { method: 'POST' }) + .then(r => r.json()) + .then(data => { + showAlert('step6-alert', data.message || data.error, data.success ? 'warning' : 'danger'); + }); +} + +function startLogPolling() { + if (state.trainingPollInterval) clearInterval(state.trainingPollInterval); + + const logViewer = document.getElementById('training-log'); + const statusBadge = document.getElementById('training-status-badge'); + const progressBar = document.getElementById('training-progress'); + const btnStart = document.getElementById('btn-start-training'); + const btnStop = document.getElementById('btn-stop-training'); + + btnStop.classList.remove('hidden'); + + state.trainingPollInterval = setInterval(() => { + fetch('/api/training/status') + .then(r => r.json()) + .then(data => { + if (!data.success) return; + + const status = data.status; + + statusBadge.className = 'status-badge'; + if (status.running) { + statusBadge.classList.add('running'); + statusBadge.innerHTML = 'โ— Running'; + } else if (status.error) { + statusBadge.classList.add('error'); + statusBadge.innerHTML = 'โ— Error'; + } else if (status.completed) { + statusBadge.classList.add('completed'); + statusBadge.innerHTML = 'โ— Completed'; + } else { + statusBadge.classList.add('idle'); + statusBadge.innerHTML = 'โ— Idle'; + } + + if (status.log_lines && status.log_lines.length > 0) { + logViewer.innerHTML = status.log_lines.map(line => { + let cls = 'log-line'; + if (line.includes('Error') || line.includes('FAILED')) cls += ' log-error'; + else if (line.includes('complete') || line.includes('saved')) cls += ' log-success'; + else if (line.includes('Warning')) cls += ' log-warn'; + return `
${escapeHtml(line)}
`; + }).join(''); + logViewer.scrollTop = logViewer.scrollHeight; + } + + const progress = estimateProgress(status.log_lines); + progressBar.style.width = progress + '%'; + + if (!status.running && status.start_time) { + clearInterval(state.trainingPollInterval); + btnStart.disabled = false; + btnStart.innerHTML = '๐Ÿš€ Start Training'; + btnStop.classList.add('hidden'); + + if (status.completed) { + showAlert('step6-alert', 'Training completed successfully! You can now generate defects.', 'success'); + loadCheckpoints(); + } else if (status.error) { + showAlert('step6-alert', 'Training failed: ' + status.error, 'danger'); + } + } + }); + }, 2000); +} + +function estimateProgress(logLines) { + if (!logLines || logLines.length === 0) return 0; + const lastLines = logLines.slice(-20).join(' '); + const match = lastLines.match(/(\d+)\s*\/\s*(\d+)/); + if (match) { + const current = parseInt(match[1]); + const total = parseInt(match[2]); + if (total > 0) return Math.min(100, (current / total) * 100); + } + if (lastLines.includes('complete') || lastLines.includes('saved')) return 100; + return 0; +} + +// ============================================================================= +// Generation (DefectFill) +// ============================================================================= + +function initGenerationForm() { + document.getElementById('gen-form').addEventListener('submit', (e) => { + e.preventDefault(); + generateDefect(); + }); + + document.getElementById('btn-refresh-checkpoints').addEventListener('click', loadCheckpoints); + + initDetectSimilarProduct(); +} + +function initDetectSimilarProduct() { + const detectUploadZone = document.getElementById('detect-upload-zone'); + const detectFileInput = document.getElementById('detect-product-image'); + const btnDetect = document.getElementById('btn-detect-product'); + const preview = document.getElementById('detect-image-preview'); + + if (!detectUploadZone || !btnDetect) return; + + detectUploadZone.addEventListener('click', () => detectFileInput.click()); + + detectFileInput.addEventListener('change', (e) => { + if (e.target.files && e.target.files[0]) { + const file = e.target.files[0]; + preview.innerHTML = ``; + preview.classList.remove('hidden'); + } + }); + + btnDetect.addEventListener('click', () => { + const files = detectFileInput.files; + if (!files || files.length === 0) { + showAlert('detect-product-alert', 'Please upload 1 clean product image first.', 'warning'); + return; + } + + const formData = new FormData(); + formData.append('image', files[0]); + + btnDetect.disabled = true; + btnDetect.innerHTML = ' Detecting similarity...'; + + fetch('/api/detect-similar-product', { + method: 'POST', + body: formData + }) + .then(r => r.json()) + .then(data => { + btnDetect.disabled = false; + btnDetect.innerHTML = '๐Ÿ” Verify Product Model'; + + if (data.success) { + const res = data.result; + if (res.is_trained) { + const select = document.getElementById('gen-object-class'); + select.value = res.matched_product; + select.dispatchEvent(new Event('change')); + + showAlert( + 'detect-product-alert', + `โœ… Match Found! Recognized as ${escapeHtml(res.matched_product)} (Similarity: ${(res.confidence * 100).toFixed(1)}%). Object Class dropdown updated automatically.`, + 'success' + ); + } else { + const alertHtml = ` +
+ โš ๏ธ Product Not Trained
+ ${escapeHtml(res.message)}
+ +
+ `; + const alertDiv = document.getElementById('detect-product-alert'); + alertDiv.className = 'alert alert-warning mt-2'; + alertDiv.innerHTML = alertHtml; + alertDiv.classList.remove('hidden'); + } + } else { + showAlert('detect-product-alert', data.error, 'danger'); + } + }) + .catch(err => { + btnDetect.disabled = false; + btnDetect.innerHTML = '๐Ÿ” Verify Product Model'; + showAlert('detect-product-alert', 'Detection error: ' + err.message, 'danger'); + }); + }); +} + +function switchToTrainingTab() { + const trainingTabBtn = document.querySelector('.tab-btn[data-tab="training"]'); + if (trainingTabBtn) { + trainingTabBtn.click(); + window.scrollTo({ top: 0, behavior: 'smooth' }); + } +} + +function loadCheckpoints() { + fetch('/api/checkpoints') + .then(r => r.json()) + .then(data => { + if (data.success) { + state.checkpoints = data.checkpoints; + renderCheckpointSelector(); + } + }) + .catch(console.error); +} + +function renderCheckpointSelector() { + const select = document.getElementById('gen-object-class'); + const defectSelect = document.getElementById('gen-defect-type'); + + const currentObj = select.value; + + select.innerHTML = ''; + Object.keys(state.checkpoints).forEach(objClass => { + const opt = document.createElement('option'); + opt.value = objClass; + opt.textContent = objClass; + select.appendChild(opt); + }); + + if (currentObj && state.checkpoints[currentObj]) { + select.value = currentObj; + updateDefectTypeOptions(currentObj); + } + + select.addEventListener('change', (e) => { + updateDefectTypeOptions(e.target.value); + }); +} + +function updateDefectTypeOptions(objClass) { + const defectSelect = document.getElementById('gen-defect-type'); + defectSelect.innerHTML = ''; + + if (state.checkpoints[objClass]) { + state.checkpoints[objClass].forEach(dt => { + const opt = document.createElement('option'); + opt.value = dt; + opt.textContent = dt; + defectSelect.appendChild(opt); + }); + } +} + +function generateDefect() { + const form = document.getElementById('gen-form'); + const formData = new FormData(form); + + const objClass = formData.get('object_class'); + if (!objClass) { + showAlert('gen-alert', 'Please select a trained product.', 'warning'); + return; + } + + const btn = document.getElementById('btn-generate'); + btn.disabled = true; + btn.innerHTML = ' Generating...'; + + fetch('/api/generate', { + method: 'POST', + body: formData + }) + .then(r => r.json()) + .then(data => { + if (data.success) { + startGenerationPolling(data.job_id, 'defectfill'); + } else { + btn.disabled = false; + btn.innerHTML = 'โœจ Generate Synthetic Defects'; + showAlert('gen-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + btn.innerHTML = 'โœจ Generate Synthetic Defects'; + showAlert('gen-alert', 'Generation failed: ' + err.message, 'danger'); + }); +} + +// ============================================================================= +// DefectDiffu +// ============================================================================= + +function initDefectDiffu() { + // Load config on init + loadDDConfig(); + + // Save config button + document.getElementById('btn-save-dd-config').addEventListener('click', saveDDConfig); + + // Generation form + document.getElementById('dd-gen-form').addEventListener('submit', (e) => { + e.preventDefault(); + generateDefectDiffu(); + }); +} + +function loadDDConfig() { + fetch('/api/defectdiffu/config') + .then(r => r.json()) + .then(data => { + if (data.success && data.config) { + state.ddConfig = data.config; + document.getElementById('dd-ckpt-path').value = data.config.ckpt_path || ''; + document.getElementById('dd-vae-path').value = data.config.vae_path || ''; + document.getElementById('dd-vlm-model').value = data.config.vlm_model || 'gemma3:12b'; + } + }) + .catch(console.error); +} + +function saveDDConfig() { + const ckptPath = document.getElementById('dd-ckpt-path').value.trim(); + const vaePath = document.getElementById('dd-vae-path').value.trim(); + const vlmModel = document.getElementById('dd-vlm-model').value.trim() || 'gemma3:12b'; + + if (!ckptPath || !vaePath) { + showAlert('dd-config-alert', 'Both checkpoint path and VAE path are required.', 'warning'); + return; + } + + fetch('/api/defectdiffu/config', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + ckpt_path: ckptPath, + vae_path: vaePath, + vlm_model: vlmModel + }) + }) + .then(r => r.json()) + .then(data => { + if (data.success) { + state.ddConfig = { ckpt_path: ckptPath, vae_path: vaePath, vlm_model: vlmModel }; + showAlert('dd-config-alert', data.message, 'success'); + } else { + showAlert('dd-config-alert', data.error, 'danger'); + } + }) + .catch(err => { + showAlert('dd-config-alert', 'Error saving config: ' + err.message, 'danger'); + }); +} + +function generateDefectDiffu() { + const form = document.getElementById('dd-gen-form'); + const formData = new FormData(form); + + const productDesc = formData.get('product_desc'); + if (!productDesc) { + showAlert('dd-gen-alert', 'Product description is required.', 'warning'); + return; + } + + // Check config is set + if (!state.ddConfig.ckpt_path || !state.ddConfig.vae_path) { + showAlert('dd-gen-alert', 'Please configure DefectDiffu model paths first.', 'warning'); + document.getElementById('dd-config-card').scrollIntoView({ behavior: 'smooth' }); + return; + } + + const btn = document.getElementById('btn-dd-generate'); + btn.disabled = true; + btn.innerHTML = ' Generating...'; + + fetch('/api/defectdiffu/generate', { + method: 'POST', + body: formData + }) + .then(r => r.json()) + .then(data => { + if (data.success) { + startGenerationPolling(data.job_id, 'defectdiffu'); + } else { + btn.disabled = false; + btn.innerHTML = 'โœจ Generate with DefectDiffu'; + showAlert('dd-gen-alert', data.error, 'danger'); + } + }) + .catch(err => { + btn.disabled = false; + btn.innerHTML = 'โœจ Generate with DefectDiffu'; + showAlert('dd-gen-alert', 'Generation failed: ' + err.message, 'danger'); + }); +} + +// ============================================================================= +// Generation Polling (Unified for both models) +// ============================================================================= + +function startGenerationPolling(jobId, mode) { + const btn = mode === 'defectfill' + ? document.getElementById('btn-generate') + : document.getElementById('btn-dd-generate'); + const resultsContainer = mode === 'defectfill' + ? document.getElementById('gen-results') + : document.getElementById('dd-gen-results'); + const summaryDiv = mode === 'defectfill' + ? document.getElementById('gen-summary') + : document.getElementById('dd-gen-summary'); + const alertId = mode === 'defectfill' ? 'gen-alert' : 'dd-gen-alert'; + + resultsContainer.innerHTML = ` +
+

โณ Generating Synthetic Defects...

+
+
+
+
+ Initializing... + +
+
+
Waiting for generation to start...
+
+
+ `; + + const pollInterval = setInterval(() => { + fetch(`/api/generation-status/${jobId}`) + .then(r => r.json()) + .then(data => { + if (!data.success) { + clearInterval(pollInterval); + return; + } + + const job = data.job; + let progress = 0; + if (job.unit_total && job.unit_current !== undefined) { + progress = Math.min(100, (job.unit_current / job.unit_total) * 100); + } else { + progress = Math.min(100, (job.step / job.total_steps) * 100); + } + + document.getElementById(`gen-progress-bar-${mode}`).style.width = progress + '%'; + document.getElementById(`gen-step-text-${mode}`).textContent = + `Step ${job.step}/${job.total_steps}: ${job.step_text}`; + + const defectCounter = document.getElementById(`gen-defect-counter-${mode}`); + let counterText = ''; + if (job.image_total > 0) { + counterText += `Image ${job.image_current}/${job.image_total}`; + } + if (job.defect_total > 0) { + counterText += (counterText ? ' | ' : '') + `Defect ${job.defect_current}/${job.defect_total}`; + } + defectCounter.textContent = counterText; + + const logViewer = document.getElementById(`gen-log-viewer-${mode}`); + const line = document.createElement('div'); + line.className = 'log-line'; + line.textContent = `[${new Date().toLocaleTimeString()}] ${job.step_text}`; + logViewer.appendChild(line); + logViewer.scrollTop = logViewer.scrollHeight; + + if (job.status === 'completed') { + clearInterval(pollInterval); + btn.disabled = false; + btn.innerHTML = mode === 'defectfill' + ? 'โœจ Generate Synthetic Defects' + : 'โœจ Generate with DefectDiffu'; + renderGenerationResults(job.result, mode); + showAlert(alertId, `Generated ${job.result.output_images?.length || 0} images in ${formatElapsedTime(job.result.elapsed_time)}`, 'success'); + } else if (job.status === 'error') { + clearInterval(pollInterval); + btn.disabled = false; + btn.innerHTML = mode === 'defectfill' + ? 'โœจ Generate Synthetic Defects' + : 'โœจ Generate with DefectDiffu'; + showAlert(alertId, 'Generation failed: ' + job.error, 'danger'); + } + }) + .catch(console.error); + }, 10000); +} + +function renderGenerationResults(data, mode = 'defectfill') { + const container = mode === 'defectfill' + ? document.getElementById('gen-results') + : document.getElementById('dd-gen-results'); + const summary = mode === 'defectfill' + ? document.getElementById('gen-summary') + : document.getElementById('dd-gen-summary'); + + container.innerHTML = ''; + + const grouped = data.grouped_results; + let totalImages = 0; + if (grouped) { + grouped.forEach(g => { + totalImages += g.defects.length * 4; + }); + } + + if (grouped && grouped.length > 0) { + summary.innerHTML = ` +
+ Generation Complete
+ Experiment: ${data.experiment_id || 'N/A'} | + Product: ${data.product_type || 'N/A'} | + Time: ${formatElapsedTime(data.elapsed_time)} | + ${totalImages} image(s) generated +
+ `; + + grouped.forEach((group, groupIdx) => { + const groupId = `result-group-${mode}-${groupIdx}`; + + const groupEl = document.createElement('div'); + groupEl.className = 'result-group'; + + const header = document.createElement('div'); + header.className = 'result-group-header'; + header.onclick = () => toggleResultGroup(groupId); + header.innerHTML = ` + Source ${groupIdx + 1}: ${escapeHtml(group.source_name)} + ${group.defect_count} defect(s) + โ–ถ + `; + + const content = document.createElement('div'); + content.className = 'result-group-content'; + content.id = groupId; + + group.defects.forEach((defect, defectIdx) => { + const defectId = `${groupId}-defect-${defectIdx}`; + + const defectGroup = document.createElement('div'); + defectGroup.className = 'result-defect-group'; + + const defectHeader = document.createElement('div'); + defectHeader.className = 'result-defect-header'; + defectHeader.onclick = (e) => { + e.stopPropagation(); + toggleResultGroup(defectId); + }; + defectHeader.innerHTML = ` + ${escapeHtml(defect.defect_label)} + ${escapeHtml(defect.defect_type)} + โ–ถ + `; + + const defectContent = document.createElement('div'); + defectContent.className = 'result-defect-content'; + defectContent.id = defectId; + + const trioGrid = document.createElement('div'); + trioGrid.className = 'result-trio-grid'; + + const imageTypes = [ + { key: 'clean_image', label: 'Clean' }, + { key: 'mask_image', label: 'Mask' }, + { key: 'raw_patch_image', label: 'Raw Patch' }, + { key: 'output_image', label: 'Output' } + ]; + + imageTypes.forEach(type => { + const imgData = defect[type.key]; + if (imgData && imgData.url) { + const card = document.createElement('div'); + card.className = 'result-trio-card'; + card.innerHTML = ` +
${type.label}
+ ${imgData.filename} +
${escapeHtml(imgData.filename)}
+ Download + `; + trioGrid.appendChild(card); + } + }); + + defectContent.appendChild(trioGrid); + defectGroup.appendChild(defectHeader); + defectGroup.appendChild(defectContent); + content.appendChild(defectGroup); + }); + + groupEl.appendChild(header); + groupEl.appendChild(content); + container.appendChild(groupEl); + }); + } else if (data.output_images && data.output_images.length > 0) { + summary.innerHTML = ` +
+ Generation Complete
+ Experiment: ${data.experiment_id || 'N/A'} | + Product: ${data.product_type || 'N/A'} | + Time: ${formatElapsedTime(data.elapsed_time)} +
+ `; + + const grid = document.createElement('div'); + grid.className = 'result-grid'; + + data.output_images.forEach(img => { + const card = document.createElement('div'); + card.className = 'result-card'; + card.innerHTML = ` + ${img.filename} +
+

${img.defect_type}

+

${img.filename}

+ Download +
+ `; + grid.appendChild(card); + }); + + container.appendChild(grid); + + } else { + summary.innerHTML = ` +
+ No images were generated. Check the defect type and product description. +
+ `; + } +} + +function toggleResultGroup(groupId) { + const content = document.getElementById(groupId); + const icon = document.getElementById(`icon-${groupId}`); + const header = content.previousElementSibling; + + if (content.classList.contains('expanded')) { + content.classList.remove('expanded'); + icon.textContent = 'โ–ถ'; + header.classList.remove('active'); + } else { + content.classList.add('expanded'); + icon.textContent = 'โ–ผ'; + header.classList.add('active'); + } +} + + +// ============================================================================= +// Utilities +// ============================================================================= + +function showAlert(elementId, message, type, autoHideMs = null) { + const el = document.getElementById(elementId); + if (!el) return; + + if (el._hideTimeout) { + clearTimeout(el._hideTimeout); + el._hideTimeout = null; + } + + el.className = `alert alert-${type}`; + el.innerHTML = message; + el.classList.remove('hidden'); + + if (type === 'success') { + const delay = autoHideMs !== null ? autoHideMs : 0; + if (delay > 0) { + el._hideTimeout = setTimeout(() => { + el.classList.add('hidden'); + }, delay); + } + } +} + +function escapeHtml(text) { + const div = document.createElement('div'); + div.textContent = text; + return div.innerHTML; +} + +function previewImage(url) { + const overlay = document.createElement('div'); + overlay.className = 'modal-overlay'; + overlay.innerHTML = ` +
+ + Preview +
+ `; + overlay.addEventListener('click', (e) => { + if (e.target === overlay) overlay.remove(); + }); + document.body.appendChild(overlay); +} + +function loadSystemInfo() { + fetch('/api/system/info') + .then(r => r.json()) + .then(data => { + if (data.success) { + const info = data.info; + const badge = document.getElementById('cuda-badge'); + if (info.cuda_available) { + badge.textContent = `CUDA: ${info.cuda_device_name}`; + badge.style.background = 'var(--success)'; + } else { + badge.textContent = 'CPU Mode'; + badge.style.background = 'var(--warning)'; + } + } + }) + .catch(console.error); +} + +function formatElapsedTime(seconds) { + if (!seconds && seconds !== 0) return '0s'; + const mins = Math.floor(seconds / 60); + const secs = (seconds % 60).toFixed(1); + if (mins > 0) { + return `${mins}m ${secs}s`; + } + return `${secs}s`; +} diff --git a/style.css b/style.css new file mode 100644 index 0000000000000000000000000000000000000000..114adf441e9032febb46bc056b2a8bb651075f0d --- /dev/null +++ b/style.css @@ -0,0 +1,28 @@ +body { + padding: 2rem; + font-family: -apple-system, BlinkMacSystemFont, "Arial", sans-serif; +} + +h1 { + font-size: 16px; + margin-top: 0; +} + +p { + color: rgb(107, 114, 128); + font-size: 15px; + margin-bottom: 10px; + margin-top: 5px; +} + +.card { + max-width: 620px; + margin: 0 auto; + padding: 16px; + border: 1px solid lightgray; + border-radius: 16px; +} + +.card p:last-child { + margin-bottom: 0; +} diff --git a/templates/index.html b/templates/index.html new file mode 100644 index 0000000000000000000000000000000000000000..2d653e270561ffc27fd560a1c68143a02bda5dfa --- /dev/null +++ b/templates/index.html @@ -0,0 +1,645 @@ + + + + + + ArtiAgent Web UI โ€” Synthetic Defect Generation + + + + +
+

+ ๐Ÿ”ง ArtiAgent Web UI + Checking CUDA... +

+
+ + +
+
+ +
+ +
+
+ + +
+ + +
+
+ +
+ + +
+
+

1 Product & Defect Type

+
+
+
+ + +

This becomes the checkpoint folder name. Use lowercase with underscores.

+
+
+ + +

The specific defect category to learn (subfolder name).

+
+ + +
+
+ + + + + + + + + + + + + + + + + + +
+ + +
+ + +
+
+

๐Ÿ“‹ Quick Guide

+
+
+
    +
  1. Enter product & defect names
    2. +
  2. Upload โ‰ฅ1 defect images
    3. +
  3. Draw rectangle masks on defects
    4. +
  4. Add domain instructions (optional)
    5. +
  5. Start training & wait
    6. +
  6. Switch to "Generate Defects" tab
    7. +
+
+
+
+
+
+ + +
+
+ +
+
+

1 Product Verification (Detect Similar Product)

+
+
+

+ Not sure if this product model is trained? Upload 1 clean image to auto-detect and select the trained model. (Make sure the image is aligned, not random orientations) +

+ +
+
๐Ÿ”
+

Click to upload clean product image

+

PNG/JPG/BMP/TIF accepted

+
+ + + + + + + +
+
+ + +
+
+

2 Synthetic Defect Generation (DefectFill)

+ +
+
+
+
+ + +

Trained checkpoints auto-populate here. Train a product first.

+
+ +
+ + +

Leave empty to let the VLM choose the most plausible defect.

+
+ +
+ + +

Select "Auto" to let the VLM decide, or manually choose to guide the perception module.

+
+ +
+ + +

Be specific. This drives VLM defect planning and placement.

+
+ +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ +
+ +
+ + +
+
+ +
+
+ + +
+ +

Upload a clean (non-defective) image of the product.

+
+ +
+ + +
+ + +
+ +
+
+ + +
+
+
+

๐Ÿ–ผ๏ธ Generation Results

+
+
+
+
+

+ Generated images will appear here.
+ Fill the form and click Generate. +

+
+
+
+
+
+
+
+ + + +
+ + + + + + +