Upload folder using huggingface_hub

src/ViTAE_library.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+from tqdm import tqdm
+import random
+import matplotlib.pyplot as plt
+
+class IrisTensorDataset(Dataset):
+    def __init__(self, file_list):
+        self.file_list = file_list
+        
+    def __len__(self):
+        return len(self.file_list)
+    
+    def __getitem__(self, idx):
+        tensor_path = self.file_list[idx]
+        img_tensor = torch.load(tensor_path)
+        return img_tensor, img_tensor
+    
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, in_channels=1, patch_size=16, embed_dim=256):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = x.flatten(2)
+        x = x.transpose(1, 2)
+        return x
+    
+class ViTAutoencoder(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, in_channels=1, embed_dim=256, latent_dim=512, num_heads=16, num_layers=6):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_patches = (img_size // patch_size) ** 2
+        self.embed_dim = embed_dim
+        
+        self.patch_embed = PatchEmbedding(in_channels, patch_size, embed_dim)
+        self.encoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
+        
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim*4, batch_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        
+        self.to_latent = nn.Linear(self.num_patches * embed_dim,latent_dim)
+        self.from_latent = nn.Linear(latent_dim,self.num_patches * embed_dim)
+        
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
+        decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim*4, batch_first=True)
+        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
+        
+        self.head = nn.Linear(embed_dim, patch_size * patch_size * in_channels)
+
+    def encode(self, x):
+        x = self.patch_embed(x) + self.encoder_pos_embed
+        x = self.transformer_encoder(x)
+        x = x.flatten(1)
+        latent = self.to_latent(x)
+        return latent
+
+    def decode(self, latent, batch_size):
+        x = self.from_latent(latent)
+        x = x.view(batch_size, self.num_patches, self.embed_dim)
+        x = x + self.decoder_pos_embed
+        x = self.transformer_decoder(x, x)
+        x = self.head(x)
+        
+        grid_size = int(self.num_patches ** 0.5)
+        x = x.view(batch_size, grid_size, grid_size, 1, self.patch_size, self.patch_size)
+        x = x.permute(0, 3, 1, 4, 2, 5).contiguous()
+        recon_img = x.view(batch_size, 1, grid_size * self.patch_size, grid_size * self.patch_size)
+        return recon_img
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        latent = self.encode(x)
+        recon_img = self.decode(latent, batch_size)
+        return recon_img
+    
+
+def trainModel(max_epoch, model, optimizer, criterion, train_loader, val_loader, train_loss_list, val_loss_list, SAVE_PATH, PATIENCE):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    for epoch in range(1, max_epoch + 1):
+        # -- Training Phase --
+        model.to(device)
+        model.train()
+        running_train_loss = 0.0
+        
+        train_bar = tqdm(train_loader, desc=f"Epoch {epoch}/{max_epoch} [Train]")
+        for inputs, targets in train_bar:
+            inputs, targets = inputs.to(device), targets.to(device)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+            loss.backward()
+            optimizer.step()
+            
+            running_train_loss += loss.item() * inputs.size(0)
+            train_bar.set_postfix({'Loss': f"{loss.item():.4f}"})
+            
+        epoch_train_loss = running_train_loss / (len(train_loader) * train_loader.batch_size)
+        train_loss_list.append(epoch_train_loss)
+        
+        # -- Validation Phase --
+        model.eval()
+        running_val_loss = 0.0
+        with torch.no_grad():
+            for inputs, targets in val_loader:
+                inputs, targets = inputs.to(device), targets.to(device)
+                outputs = model(inputs)
+                loss = criterion(outputs, targets)
+                running_val_loss += loss.item() * inputs.size(0)
+                
+        epoch_val_loss = running_val_loss / len(val_loader) * val_loader.batch_size
+        val_loss_list.append(epoch_val_loss)
+        
+        print(f"Summary - Train Loss: {epoch_train_loss:.5f} | Val Loss: {epoch_val_loss:.5f}")
+        
+        # -- Autosave & Early Stopping Logic --
+        if epoch_val_loss < best_val_loss:
+            best_val_loss = epoch_val_loss
+            epochs_without_improvement = 0
+            torch.save(model.state_dict(), SAVE_PATH)
+            print(f"--> Validation improved. Model saved to {SAVE_PATH}")
+        else:
+            epochs_without_improvement += 1
+            print(f"--> No improvement for {epochs_without_improvement} epoch(s).")
+            
+            if epochs_without_improvement >= PATIENCE:
+                print(f"\n[Early Stopping Triggered] Stopping at epoch {epoch} due to tolerance limits.")
+                break
+
+    plt.figure(figsize=(10, 6))
+    plt.plot(range(1, len(train_loss_list) + 1), train_loss_list, label='Training Loss', marker='o')
+    plt.plot(range(1, len(val_loss_list) + 1), val_loss_list, label='Validation Loss', marker='x')
+    plt.title('Training and Validation Loss over Epochs')
+    plt.xlabel('Epochs')
+    plt.ylabel('Mean Squared Error (MSE)')
+    plt.legend()
+    plt.grid(True)
+    plt.show()
+
+
+def ViTAEModelTest(data, model, type):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model.to(device)
+    criterion = nn.MSELoss()
+    if type.lower() == 'batch':
+        running_test_loss = 0.0
+        with torch.no_grad():
+            for inputs, targets in tqdm(data, desc="Evaluating Test Set"):
+                inputs, targets = inputs.to(device), targets.to(device)
+                outputs = model(inputs)
+                loss = criterion(outputs, targets)
+                running_test_loss += loss.item() * inputs.size(0)
+                
+        final_test_loss = running_test_loss / len(data) * data.batch_size
+        print(f"\nFinal Test Set Loss (MSE): {final_test_loss:.5f}")
+
+    elif type.lower() == 'single':
+        original_data=data.unsqueeze(0)
+        recon_data=model(original_data)
+
+        original_array=original_data.squeeze().detach().cpu().numpy()
+        recon_array=recon_data.squeeze().detach().cpu().numpy()
+        diff_img_from_file = original_array - recon_array
+        
+        fig, axes = plt.subplots(1, 3, figsize=(10, 5))
+
+        # 4. Plot the first image on the left axis (axes[0])
+        axes[0].imshow(original_array, cmap='gray')
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')  # Optional: hides grid lines and pixel numbers
+
+        # 5. Plot the second image on the right axis (axes[1])
+        axes[1].imshow(recon_array, cmap='gray')
+        axes[1].set_title('Reconstructed Image')
+        axes[1].axis('off')  # Optional: hides grid lines and pixel numbers
+
+        axes[2].imshow(diff_img_from_file, cmap='gray')
+        axes[2].set_title('Difference Image')
+        axes[2].axis('off')
+        # 6. Display the plot cleanly
+        plt.tight_layout()  # Prevents labels/titles from overlapping
+        plt.show()

src/test_model.py

@@ -0,0 +1,171 @@
+"""
+Iris Matching Model Tester
+===========================
+ทดสอบโมเดลจับคู่ม่านตา (ViTAutoencoder + MatchCNN_V2)
+
+วิธีใช้:
+    python test_model.py <ภาพม่านตา1> <ภาพม่านตา2>
+
+ตัวอย่าง:
+    python test_model.py sample_A.jpg sample_B.jpg
+
+ไฟล์ที่ต้องมีอยู่ในโฟลเดอร์เดียวกับสคริปต์นี้ (หรือแก้ path ด้านล่าง):
+    - best_match_cnn.pth       (มีอยู่แล้ว)
+    - best_vit_autoencoder.pth (ยังไม่มี - ต้องดาวน์โหลดจาก Drive: IrisPeriNectec/)
+    - ViTAE_library.py         (ยังไม่มี - ต้องดาวน์โหลดจาก Drive: IrisPeriNectec/)
+"""
+
+import os
+import sys
+import argparse
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+
+# ---------------------------------------------------------------------------
+# Config — แก้ path ตรงนี้ถ้าไฟล์ไม่ได้อยู่โฟลเดอร์เดียวกับสคริปต์
+# ---------------------------------------------------------------------------
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+AUTOENCODER_WEIGHTS = os.path.join(SCRIPT_DIR, "..", "Model", "best_vit_autoencoder.pth")
+MATCHCNN_WEIGHTS = os.path.join(SCRIPT_DIR, "..", "Model", "best_match_cnn.pth")
+VITAE_LIBRARY_DIR = SCRIPT_DIR  # โฟลเดอร์ที่มี ViTAE_library.py
+MATCH_THRESHOLD = 0.5
+
+
+# ---------------------------------------------------------------------------
+# MatchCNN_V2 — copy มาจาก NewGoodcnnV1.ipynb / Testcnn.ipynb ทุกตัวอักษร
+# ---------------------------------------------------------------------------
+class MatchCNN_V2(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.cnn_extractor = nn.Sequential(
+            nn.Conv2d(2, 32, 5, padding=2),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Flatten(),
+        )
+        self.fc_input_size = 100352
+
+        self.classifier = nn.Sequential(
+            nn.Linear(self.fc_input_size + 1, 256),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(256, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, input1, input2, cos_sim):
+        combined_imgs = torch.cat([input1, input2], dim=1)
+        cnn_features = self.cnn_extractor(combined_imgs)
+        fused = torch.cat([cnn_features, cos_sim], dim=1)
+        return self.classifier(fused).squeeze()
+
+
+def check_dependencies():
+    """เช็คว่าไฟล์ที่จำเป็นครบไหมก่อนรัน แจ้ง error ที่อ่านง่าย"""
+    missing = []
+    if not os.path.exists(MATCHCNN_WEIGHTS):
+        missing.append(f"  - {MATCHCNN_WEIGHTS}")
+    if not os.path.exists(AUTOENCODER_WEIGHTS):
+        missing.append(f"  - {AUTOENCODER_WEIGHTS}")
+    vitae_path = os.path.join(VITAE_LIBRARY_DIR, "ViTAE_library.py")
+    if not os.path.exists(vitae_path):
+        missing.append(f"  - {vitae_path}")
+
+    if missing:
+        print("❌ ไฟล์ที่จำเป็นยังไม่ครบ ขาดไฟล์ต่อไปนี้:")
+        for m in missing:
+            print(m)
+        print(
+            "\nไปดาวน์โหลดจาก Google Drive โฟลเดอร์ 'IrisPeriNectec' "
+            "(best_vit_autoencoder.pth และ ViTAE_library.py) "
+            "แล้ววางไว้ในโฟลเดอร์เดียวกับสคริปต์นี้ครับ"
+        )
+        sys.exit(1)
+
+
+def load_models(device):
+    check_dependencies()
+
+    sys.path.insert(0, VITAE_LIBRARY_DIR)
+    from ViTAE_library import ViTAutoencoder  # noqa: E402  (import ต้องมาหลัง sys.path.insert)
+
+    print("กำลังโหลด ViTAutoencoder ...")
+    autoencoder = ViTAutoencoder(
+        img_size=224,
+        patch_size=16,
+        in_channels=1,
+        embed_dim=256,
+        latent_dim=512,
+        num_heads=16,
+        num_layers=6,
+    ).to(device)
+    autoencoder.load_state_dict(torch.load(AUTOENCODER_WEIGHTS, map_location=device))
+    autoencoder.eval()
+
+    print("กำลังโหลด MatchCNN_V2 ...")
+    model = MatchCNN_V2().to(device)
+    model.load_state_dict(torch.load(MATCHCNN_WEIGHTS, map_location=device))
+    model.eval()
+
+    print("✅ โหลดโมเดลครบแล้ว\n")
+    return autoencoder, model
+
+
+def predict_pair(path1, path2, autoencoder, model, device, transform):
+    """คืนค่า match probability (0-1) ของภาพม่านตาสองรูป"""
+    for p in (path1, path2):
+        if not os.path.exists(p):
+            raise FileNotFoundError(f"ไม่พบไฟล์ภาพ: {p}")
+
+    img1 = transform(Image.open(path1).convert("L")).unsqueeze(0).to(device)
+    img2 = transform(Image.open(path2).convert("L")).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        t1 = autoencoder(img1)
+        t2 = autoencoder(img2)
+        cos_sim = F.cosine_similarity(t1.view(1, -1), t2.view(1, -1), dim=1).unsqueeze(1)
+        prob = model(t1, t2, cos_sim)
+
+    return prob.item()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ทดสอบโมเดลจับคู่ม่านตา")
+    parser.add_argument("image1", help="path ภาพม่านตารูปที่ 1")
+    parser.add_argument("image2", help="path ภาพม่านตารูปที่ 2")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"ใช้งานบน device: {device}\n")
+
+    autoencoder, model = load_models(device)
+
+    transform = transforms.Compose(
+        [transforms.Resize((224, 224)), transforms.ToTensor()]
+    )
+
+    prob = predict_pair(args.image1, args.image2, autoencoder, model, device, transform)
+
+    print("=" * 45)
+    print(f"ภาพ 1: {args.image1}")
+    print(f"ภาพ 2: {args.image2}")
+    print(f"Match probability: {prob:.4f}")
+    if prob > MATCH_THRESHOLD:
+        print("ผลลัพธ์: ✅ MATCH (น่าจะเป็นคนเดียวกัน)")
+    else:
+        print("ผลลัพธ์: ❌ NO MATCH (น่าจะคนละคน)")
+    print("=" * 45)
+
+
+if __name__ == "__main__":
+    main()

src/test_model_tpl.py

@@ -0,0 +1,173 @@
+"""
+Iris Matching Model Tester — TPL (Template) Pipeline
+======================================================
+ทดสอบโมเดลจับคู่ม่านตา เวอร์ชัน Template-based
+(ViTAutoencoder.encode() -> 512-dim template -> FC matcher)
+
+วิธีใช้:
+    python test_model_tpl.py <ภาพม่านตา1> <ภาพม่านตา2>
+
+ไฟล์ที่ต้องมีอยู่ในโฟลเดอร์เดียวกับสคริปต์นี้:
+    - scar_encoder_triplet.pth   (ViT encoder เทรนด้วย triplet loss)
+    - best_match_cnn(TPL).pth    (FC matcher — ชื่อไฟล์มีคลา์มีคหว้ง cnn ค้างมาจากเดิม
+                                   แต่ตัวจริงเป็น fully-connected)
+    - ViTAE_library.py           (นิยามคลาส ViTAutoencoder)
+
+หมายเหตุสำคัญ:
+    คลาส MatchFC_TPL ด้านล่างนี้ "reverse-engineer" มาจากการอ่าน shape ของ
+    state_dict ใน best_match_cnn(TPL).pth โดยตรง (ไม่มี notebook ต้นฉบับให้ดู)
+    โครงสร้างชั้น Linear ทุกชั้นตรงกับ state_dict 100% (ตรวจสอบแล้ว)
+    แต่ตัว activation function ระหว่างชั้น (ReLU) เป็นการอนุมานจากรูปแบบเดียวกับ
+    MatchCNN_V2 ตัวเดิมของผู้เขียนคนเดียวกัน ถ้าผลลัพธ์ดูแปลกๆ (เช่น probability
+    กระจุกที่ 0.5 หรือ 0/1 สุดขั้วผิดปกติ) ให้ลองเทียบกับ v2.ipynb/ANNtest.ipynb
+    ถ้ามีต้นฉบับจริง
+"""
+
+import os
+import sys
+import argparse
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+ENCODER_WEIGHTS = os.path.join(SCRIPT_DIR, "..", "Model", "scar_encoder_triplet.pth")
+MATCHER_WEIGHTS = os.path.join(SCRIPT_DIR, "..", "Model", "best_match_cnn(TPL).pth")
+VITAE_LIBRARY_DIR = SCRIPT_DIR
+MATCH_THRESHOLD = 0.5
+
+
+# ---------------------------------------------------------------------------
+# MatchFC_TPL — reverse-engineered จาก state_dict ของ best_match_cnn(TPL).pth
+#
+#   feature_processing.0.weight  (512, 1024)   -> Linear(1024, 512)
+#   feature_processing.2.weight  (256, 512)    -> Linear(512, 256)
+#   classifier.0.weight          (128, 257)    -> Linear(257, 128)
+#   classifier.3.weight          (1, 128)      -> Linear(128, 1)
+#
+#   1024 = template1(512) + template2(512) ต่อกัน
+#   257  = fused_feature(256) + cos_sim(1) ต่อกัน
+# ---------------------------------------------------------------------------
+class MatchFC_TPL(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.feature_processing = nn.Sequential(
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+        )
+        self.classifier = nn.Sequential(
+            nn.Linear(256 + 1, 128),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(128, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, template1, template2, cos_sim):
+        combined = torch.cat([template1, template2], dim=1)      # (B, 1024)
+        feat = self.feature_processing(combined)                  # (B, 256)
+        fused = torch.cat([feat, cos_sim], dim=1)                 # (B, 257)
+        return self.classifier(fused).squeeze()
+
+
+def check_dependencies():
+    missing = []
+    if not os.path.exists(ENCODER_WEIGHTS):
+        missing.append(f"  - {ENCODER_WEIGHTS}")
+    if not os.path.exists(MATCHER_WEIGHTS):
+        missing.append(f"  - {MATCHER_WEIGHTS}")
+    vitae_path = os.path.join(VITAE_LIBRARY_DIR, "ViTAE_library.py")
+    if not os.path.exists(vitae_path):
+        missing.append(f"  - {vitae_path}")
+
+    if missing:
+        print("❌ ไฟล์ที่จำเป็นยังไม่ครบ ขาดไฟล์ต่อไปนี้:")
+        for m in missing:
+            print(m)
+        sys.exit(1)
+
+
+def load_models(device):
+    check_dependencies()
+
+    sys.path.insert(0, VITAE_LIBRARY_DIR)
+    from ViTAE_library import ViTAutoencoder  # noqa: E402
+
+    print("กำลังโหลด ViT encoder (triplet) ...")
+    encoder = ViTAutoencoder(
+        img_size=224,
+        patch_size=16,
+        in_channels=1,
+        embed_dim=256,
+        latent_dim=512,
+        num_heads=16,
+        num_layers=6,
+    ).to(device)
+    encoder.load_state_dict(torch.load(ENCODER_WEIGHTS, map_location=device))
+    encoder.eval()
+
+    print("กำลังโหลด MatchFC_TPL ...")
+    matcher = MatchFC_TPL().to(device)
+    matcher.load_state_dict(torch.load(MATCHER_WEIGHTS, map_location=device))
+    matcher.eval()
+
+    print("✅ โหลดโมเดลครบแล้ว\n")
+    return encoder, matcher
+
+
+def predict_pair(path1, path2, encoder, matcher, device, transform):
+    for p in (path1, path2):
+        if not os.path.exists(p):
+            raise FileNotFoundError(f"ไม่พบไฟล์ภาพ: {p}")
+
+    img1 = transform(Image.open(path1).convert("L")).unsqueeze(0).to(device)
+    img2 = transform(Image.open(path2).convert("L")).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        # .encode() ให้ template 512 มิติตรงๆ (ไม่ผ่าน decode กลับเป็นภาพ)
+        t1 = encoder.encode(img1)   # (1, 512)
+        t2 = encoder.encode(img2)  # (1, 512)
+        cos_sim = F.cosine_similarity(t1, t2, dim=1).unsqueeze(1)  # (1, 1)
+        prob = matcher(t1, t2, cos_sim)
+
+    return prob.item()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ทดสอบโมเดลจับคู่ม่านตา (TPL pipeline)")
+    parser.add_argument("image1", help="path ภาพม่านตารูปที่ 1")
+    parser.add_argument("image2", help="path ภาพม่านตารูปที่ 2")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"ใช้งานบน device: {device}\n")
+
+    encoder, matcher = load_models(device)
+
+    transform = transforms.Compose(
+        [transforms.Resize((224, 224)), transforms.ToTensor()]
+    )
+
+    prob = predict_pair(args.image1, args.image2, encoder, matcher, device, transform)
+
+    print("=" * 45)
+    print(f"ภาพ 1: {args.image1}")
+    print(f"ภาพ 2: {args.image2}")
+    print(f"Match probability: {prob:.4f}")
+    if prob > MATCH_THRESHOLD:
+        print("ขลลัพธ์: ✅ MATCH (๙่าจะเป็นคนเดียวกัน)")
+    else:
+        print("ขลลัพธ์: ❌ NO MATCH (๙่าจะคนละคน)")
+    print("=" * 45)
+
+
+if __name__ == "__main__":
+    main()