Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import os
- import torch
- import numpy as np
- import cv2
- from tqdm import tqdm
- from torchvision import models, transforms
- from sklearn.neighbors import KNeighborsClassifier
- # Cek apakah GPU tersedia dan kompatibel dengan CUDA 12.2
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- print(f"🚀 Menggunakan device: {device}")
- # Path dataset
- DATASET_PATH = "assets/datasets"
- TRAIN_PATH = os.path.join(DATASET_PATH, "train")
- TEST_PATH = os.path.join(DATASET_PATH, "test")
- # Load Pretrained Model (EfficientNet-B0 sebagai Feature Extractor)
- efficientnet = models.efficientnet_b0(weights="IMAGENET1K_V1")
- efficientnet = torch.nn.Sequential(*list(efficientnet.children())[:-1]) # Hapus layer FC
- efficientnet = efficientnet.to(device).eval()
- # Transformasi untuk preprocessing gambar
- transform = transforms.Compose([
- transforms.ToPILImage(),
- transforms.Resize((224, 224)),
- transforms.ToTensor(),
- transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
- ])
- def extract_features(image_paths):
- """Ekstrak fitur dari batch gambar menggunakan EfficientNet-B0."""
- batch_images = []
- valid_paths = []
- for image_path in image_paths:
- image = cv2.imread(image_path)
- if image is not None:
- image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
- image = transform(image)
- batch_images.append(image)
- valid_paths.append(image_path)
- if len(batch_images) == 0:
- return None, valid_paths
- batch_tensor = torch.stack(batch_images).to(device) # Konversi ke tensor batch
- with torch.no_grad():
- features = efficientnet(batch_tensor).squeeze().cpu().numpy()
- return features, valid_paths
- def load_dataset(dataset_path, batch_size=32):
- """Muat dataset dan ekstrak fitur dari setiap gambar dalam batch."""
- X, y = [], []
- image_paths = []
- labels = []
- for person_id in os.listdir(dataset_path):
- person_folder = os.path.join(dataset_path, person_id)
- if not os.path.isdir(person_folder):
- continue
- for file in os.listdir(person_folder):
- if file.endswith(('.jpg', '.jpeg', '.png')):
- img_path = os.path.join(person_folder, file)
- image_paths.append(img_path)
- labels.append(person_id)
- # Proses ekstraksi fitur dalam batch
- for i in tqdm(range(0, len(image_paths), batch_size), desc="📂 Ekstraksi Fitur", unit="batch"):
- batch_paths = image_paths[i:i+batch_size]
- features, valid_paths = extract_features(batch_paths)
- if features is not None:
- X.extend(features)
- y.extend(labels[i:i+len(valid_paths)])
- return np.array(X), np.array(y)
- # Load Data Train & Test dengan batch processing
- print("📂 Memuat dataset train...")
- X_train, y_train = load_dataset(TRAIN_PATH, batch_size=64)
- print("📂 Memuat dataset test...")
- X_test, y_test = load_dataset(TEST_PATH, batch_size=64)
- # Inisialisasi dan Latih KNN Classifier
- print("🚀 Melatih model Few-Shot Learning dengan KNN...")
- knn = KNeighborsClassifier(n_neighbors=3, metric="euclidean")
- knn.fit(X_train, y_train)
- # Evaluasi Model
- y_pred = knn.predict(X_test)
- accuracy = np.mean(y_pred == y_test)
- print(f"✅ Few-Shot Learning Accuracy: {accuracy:.2%}")
Add Comment
Please, Sign In to add comment
