FaceRC

 avatarquoc14
Publica year ago16 Snippets
Search
Language
Sort by
📅 Created Date
Order

eval.py

pythona year ago
import os
import numpy as np
import pandas as pd
from PIL import Image
from huggingface_model_utils import load_model_from_local_path
from torchvision.transforms import Compose, ToTensor, Normalize
import torch
import inspect
from facenet_pytorch import MTCNN


# Load models
device = 'cuda' if torch.cuda.is_available() else 'cpu'
aligner = load_model_from_local_path('model/minchul/cvlface_DFA_resnet50').to(device)
fr_model = load_model_from_local_path('model/minchul/cvlface_adaface_vit_base_webface4m').to(device)

# Hàm chuyển ảnh PIL sang tensor
def pil_to_input(pil_image, device):
    trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
    return trans(pil_image).unsqueeze(0).to(device)

# Hàm trích xuất đặc trưng từ ảnh
def get_feat(pil_image):
    """Trích xuất đặc trưng từ ảnh PIL."""
    input_tensor = pil_to_input(pil_image, device)
    aligned_x, _, aligned_ldmks, _, _, _ = aligner(input_tensor)

    # Kiểm tra xem mô hình có cần keypoints không
    input_signature = inspect.signature(fr_model.model.net.forward)
    if input_signature.parameters.get('keypoints') is not None:
        feat = fr_model(aligned_x, aligned_ldmks)
    else:
        feat = fr_model(aligned_x)
    return feat

# Hàm tính toàn bộ cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu
def compute_cosine_similarity(feat_input, feats_db):
    """Tính toàn bộ cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu."""
    return torch.nn.functional.cosine_similarity(feat_input, feats_db)

# Hàm lấy ID bằng cách so sánh với cơ sở dữ liệu
def get_id(pil_image, feats_db, db_ids, threshold=0.3):
    """So sánh ảnh với cơ sở dữ liệu và trả về ID nếu tìm thấy."""
    # Trích xuất đặc trưng từ ảnh đầu vào
    feat_input = get_feat(pil_image)

    # Tính toán cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu
    similarities = compute_cosine_similarity(feat_input, feats_db)

    # Lấy ra giá trị cosine similarity lớn nhất và index tương ứng
    max_sim, idx_max = similarities.max(0)

    if max_sim.item() >= threshold:
        best_match_id = db_ids[idx_max.item()]
        return {"id": best_match_id, "status": "found"}
    else:
        return {"id": None, "status": "not_found"}

# Hàm chính
def main():
    # Đọc cơ sở dữ liệu enroll một lần
    enroll_csv = 'face_features_vit.csv'
    if not os.path.exists(enroll_csv):
        print(f"Không tìm thấy file cơ sở dữ liệu {enroll_csv}.")
        return

    # Load và chuyển đổi đặc trưng thành tensor một lần duy nhất
    enroll_data = pd.read_csv(enroll_csv, dtype={'id': str})
    
    # Chuyển đặc trưng từ chuỗi thành tensor để tính toán nhanh hơn
    feats_db = torch.stack([torch.tensor(eval(feat), device=device) for feat in enroll_data['feat'].values])
    db_ids = enroll_data['id'].tolist()

    # Khởi tạo các chỉ số
    TP, FP, FN, TN = 0, 0, 0, 0
    threshold = 0.5  # Ngưỡng cosine similarity

    # Đường dẫn tới thư mục inout
    folder_path2 = '/home/quoc14/Code/HeThongNhanDang/_vts_users_v1/vts_in_out_faces_v1'
    processed_folders = 0
    total_folders = len(os.listdir(folder_path2))

    # Duyệt qua các folder trong inout
    for folder in os.listdir(folder_path2):
        folder_path = os.path.join(folder_path2, folder)
        if os.path.isdir(folder_path):
            image_files = [f for f in os.listdir(folder_path) if f.endswith('.jpg')]
            for image_file in image_files:
                image_path = os.path.join(folder_path, image_file)
                pil_image = Image.open(image_path)

                # Nhận diện khuôn mặt bằng cách so sánh đặc trưng
                check_result = get_id(pil_image, feats_db, db_ids, threshold)
                
                if check_result['status'] == 'found':
                    matched_id = check_result['id']
                    if str(folder) == matched_id:
                        TP += 1  # Nhận diện đúng
                    else:
                        FP += 1  # Nhận diện nhầm
                else:
                    if str(folder) in db_ids:
                        FN += 1  # Lẽ ra phải nhận diện nhưng không
                    else:
                        TN += 1  # Không có trong hệ thống và hệ thống không nhận diện

            processed_folders += 1
            print(f"Đã xử lý {processed_folders}/{total_folders} folder")

    # Tính toán FNIR, FPIR, ACC
    FNIR = FN / (TP + FN) if (TP + FN) > 0 else 0
    FPIR = FP / (FP + TN) if (FP + TN) > 0 else 0
    ACC = (TP + TN) / (TP + TN + FP + FN) if (TP + TN + FP + FN) > 0 else 0

    # Lưu kết quả vào file CSV chỉ một lần
    results = [{
        'Model': 'VIT',
        'ACC': ACC,
        'FNIR': FNIR,
        'FPIR': FPIR
    }]
    df_results = pd.DataFrame(results)
    df_results.to_csv('evaluation_results_vts.csv', index=False)
    print(FN, FP, TP, TN)
    print(f"Accuracy: {ACC}, FNIR: {FNIR}, FPIR: {FPIR}")

if __name__ == "__main__":
    main()

CodeAnhGao

pythona year ago
import gradio as gr
from face_engine.engine import FaceEngine
from PIL import Image
import torch
from torchvision.transforms import Compose, ToTensor, Normalize
import glob
from face_engine.huggingface_model_utils import load_model_from_local_path
import cv2
import numpy as np

model_extract = load_model_from_local_path("face_engine/models/minchul/cvlface_adaface_ir101_webface4m")
model_aligner = load_model_from_local_path("face_engine/models/minchul/cvlface_DFA_resnet50")

def pil_to_input(pil_image):
    trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
    return trans(pil_image).unsqueeze(0)

def get_feature(pil_image):
    input_tensor = pil_to_input(pil_image)
    aligned_x, _, aligned_ldmks, _, _, _ = model_aligner(input_tensor)
    print(aligned_x.shape)
    img = aligned_x.numpy()

    feature = model_extract(aligned_x)
    feature_normm = torch.nn.functional.normalize(feature)
    print(feature_normm.shape)
    return feature_normm, img
  

#imglink1 = "vts\\vts_in_out_faces_v1\\vts_in_out_faces_v1\\001067012926\\66a1b96dd53d525db9f90eb9.jpg"
#imglink2 = "vts\\vts_in_out_faces_v1\\vts_in_out_faces_v1\\001067012926\\66a1b277d53d525db9e959d8.jpg"

imglink1 = "D:\\_vts_users_v1\\vts_enroll_images_v1\\0a532a40-c8a7-4ad2-aad4-4368c65dec18\\0.jpg"
imglink2 = "D:\\_vts_users_v1\\vts_in_out_faces_v1\\040089018059\\66a9f5eed53d525db936c902.jpg"



img1 = Image.open(imglink1)
img2 = Image.open(imglink2)


f1, im1 = get_feature(img1)
f2, im2 = get_feature(img2)

im1 = np.transpose(im1[0], (1, 2, 0))
im2 = np.transpose(im2[0], (1, 2, 0))       

im1 = (im1+0.5)
im2 = (im2+0.5)
im1 = cv2.cvtColor(im1, cv2.COLOR_RGB2BGR)
im2 = cv2.cvtColor(im2, cv2.COLOR_RGB2BGR)
cv2.imshow('im1.jpg', im1)


distEclid = (f1 - f2).pow(2).sum().sqrt().item()
cossim = torch.nn.functional.cosine_similarity(f1, f2).item()

list_image = glob.glob("D:\\_vts_users_v1\\vts_in_out_faces_v1\\040089018059\\*.jpg")
for imglink in list_image:
    img2 = Image.open(imglink)
    f2, im2 = get_feature(img2)
    distEclid = (f1 - f2).pow(2).sum().sqrt().item()
    cossim = torch.nn.functional.cosine_similarity(f1, f2).item() #0.5
    print(imglink.split("\\")[-1].split('/')[-1], cossim, distEclid)

    im2 = np.transpose(im2[0], (1, 2, 0))
    im2 = cv2.cvtColor(im2, cv2.COLOR_RGB2BGR)
    cv2.imshow('im2.jpg', im2)
    cv2.waitKey()

baocao

plain_texta year ago
### **1. Giới thiệu hệ thống**

- Hệ thống nhận diện khuôn mặt của em gồm hai chức năng chính:
    - **Chức năng đăng ký (Enroll):** Nhận vào ảnh khuôn mặt và ID của người dùng. Hệ thống sẽ trích xuất đặc trưng của ảnh và lưu vào cơ sở dữ liệu để nhận diện trong tương lai.
        
        > Ảnh chụp màn hình chức năng này (để trống trong slide hộ mai thêm sau)
        > 
        
        ![Screenshot from 2024-10-22 06-51-32.png](https://prod-files-secure.s3.us-west-2.amazonaws.com/125df330-c655-4723-9084-11fea250a474/66d39997-7d0b-4e15-984f-f4a5f6bdf5b5/Screenshot_from_2024-10-22_06-51-32.png)
        
    - **Chức năng nhận diện (Identify):** Nhận vào ảnh khuôn mặt, trích xuất đặc trưng và so sánh với các đặc trưng đã có trong cơ sở dữ liệu, sau đó trả về ID tương ứng của khuôn mặt có độ tương đồng cao nhất.
        
        > Ảnh chụp màn hình chức năng này (để trống trong slide hộ mai thêm sau)
        > 
        
        ![Screenshot from 2024-10-22 06-51-48.png](https://prod-files-secure.s3.us-west-2.amazonaws.com/125df330-c655-4723-9084-11fea250a474/a58522b7-5ff0-426c-b620-fce5bd902bd7/Screenshot_from_2024-10-22_06-51-48.png)
        

---

### **2. Đánh giá độ chính xác của mô hình**

- Để đánh giá hiệu quả của mô hình trong hệ thống, em sử dụng **bộ dữ liệu VTS**. Đây là bộ dữ liệu với cấu trúc gồm hai phần chính:
    1. **Bộ dữ liệu Enroll:** Sử dụng để đăng ký các khuôn mặt và làm cơ sở dữ liệu cho hệ thống nhận diện.
    2. **Bộ dữ liệu Inout:** Sử dụng để kiểm tra và đánh giá khả năng nhận diện của mô hình.

---

### **3. Giới thiệu bộ dữ liệu VTS**

- **Bộ dữ liệu Enroll:**
    - Gồm 400 folder, mỗi folder tương ứng với một ID và chứa 5 ảnh khác nhau của khuôn mặt đó.
    - Đây là dữ liệu chính mà hệ thống sử dụng để nhận diện khi so sánh ảnh mới.
    
    > 
    > 
    > 
    > ![Screenshot from 2024-10-21 20-33-05.png](https://prod-files-secure.s3.us-west-2.amazonaws.com/125df330-c655-4723-9084-11fea250a474/2e859dbd-a99b-4ff9-a984-69311983bcfa/Screenshot_from_2024-10-21_20-33-05.png)
    > 
- **Bộ dữ liệu Inout:**
    - Gồm 500 folder, trong đó 400 folder có ID tương ứng trong cơ sở dữ liệu Enroll và 100 folder là những ID chưa có trong hệ thống.
    - Mục đích là kiểm tra khả năng nhận diện đúng khuôn mặt đã đăng ký và phát hiện các khuôn mặt mới chưa được đăng ký.
    
    > 
    > 
    > 
    > ![Screenshot from 2024-10-21 20-33-26.png](https://prod-files-secure.s3.us-west-2.amazonaws.com/125df330-c655-4723-9084-11fea250a474/f2241314-0311-43e9-a28d-1cce52b076e3/Screenshot_from_2024-10-21_20-33-26.png)
    > 

---

### **4. Demo hệ thống**

- **Thao tác đăng ký (Enroll):**
    - Em sẽ chọn một ảnh khuôn mặt
    - Sau khi đăng ký, hệ thống sẽ lưu đặc trưng của khuôn mặt này vào cơ sở dữ liệu để sử dụng trong quá trình nhận diện.
    
    > Thao tác demo: Tiến hành đăng ký một ảnh từ Enroll, hiển thị thông báo đã có trong csdl
    > 
- **Thao tác nhận diện (Identify):**
    - Em sẽ chọn một ảnh khuôn mặt của bọn em em và  từ bộ dữ liệu Inout để nhận diện.
    - Hệ thống sẽ so sánh đặc trưng của ảnh với cơ sở dữ liệu và trả về ID của ảnh có độ tương đồng cao nhất (nếu tìm thấy).
    
    > Thao tác demo: Nhận diện một khuôn mặt từ ảnh Inout, hệ thống sẽ trả về kết quả là ID tương ứng hoặc thông báo khuôn mặt chưa được đăng ký.
    > 

---

### **5. Đánh giá hiệu quả của mô hình qua các chỉ số**

- Sau khi demo hệ thống, em sẽ sử dụng các chỉ số sau để đánh giá hiệu quả mô hình:
    1. **Accuracy (ACC):** Độ chính xác của mô hình trong việc nhận diện đúng khuôn mặt.
    2. **False Negative Identification Rate (FNIR):** Tỷ lệ mô hình không nhận diện được khuôn mặt đã có trong cơ sở dữ liệu.
    3. **False Positive Identification Rate (FPIR):** Tỷ lệ mô hình nhận diện nhầm khuôn mặt chưa được đăng ký là đã đăng ký.
    
    > Giải thích từng chỉ số với ví dụ cụ thể, giúp người xem hiểu rõ hơn về cách đo lường hiệu quả của hệ thống.
    > 

---

### **6. Kết quả đánh giá**

- Dựa trên bộ dữ liệu VTS, hệ thống nhận diện khuôn mặt của em đạt được các chỉ số như sau:
    - **Accuracy (ACC):** 0.78
    - **FNIR: 0.2**
    - **FPIR:** 0.2
    
    > Hiển thị biểu đồ trực quan minh họa các chỉ số ACC, FNIR, FPIR để giúp người xem dễ hình dung và hiểu hơn về độ chính xác của hệ thống.
    > 

---

### **7. Kết luận**

- Hệ thống nhận diện khuôn mặt của em đạt được độ chính xác **Accuracy (ACC) = 0.78**, cho thấy hệ thống nhận diện đúng trong 78% các trường hợp.
- **FNIR = 0.2** và **FPIR = 0.2** cho thấy hệ thống bỏ sót 20% các khuôn mặt đã đăng ký và nhận nhầm 20% các khuôn mặt không có trong cơ sở dữ liệu.

Tuy nhiên, những kết quả này cũng phản ánh một phần hạn chế của hệ thống, có thể do **bộ dữ liệu VTS hiện tại chưa đủ phong phú**. Với chỉ 400 ID trong cơ sở dữ liệu đăng ký (Enroll) mỗi id có 5 ảnh , hệ thống có thể gặp khó khăn trong việc xử lý các trường hợp đa dạng hơn, dẫn đến tỷ lệ FNIR và FPIR cao hơn mong đợi.

- Để cải thiện, cần bổ sung nhiều dữ liệu hơn để giúp hệ thống nhận diện hiệu quả và chính xác hơn trong các điều kiện khác nhau.

Trong tương lai, em sẽ tập trung vào việc cải thiện thuật toán và mở rộng bộ dữ liệu để nâng cao độ chính xác và giảm thiểu các lỗi FNIR và FPIR, nhằm đưa hệ thống đến mức hiệu quả cao hơn trong thực tế.

---

main3.py

pythona year ago
import gradio as gr
from face_engine.engine import FaceEngine
from PIL import Image

# Danh sách mô hình cho người dùng lựa chọn
model_list = [
    "ir18_webface4m", 
    "ir50_webface4m", 
    "ir101_webface4m", 
    "ir101_webface12m", 
    "vit_base_kprpe_webface4m", 
    "vit_base_kprpe_webface12m", 
    "vit_base_webface4m"
]

# Khởi tạo FaceEngine
face_engine = FaceEngine()

# Hàm cập nhật mô hình dựa trên lựa chọn của người dùng
def update_model(model_name):
    face_engine.set_model(model_name)
    return f"Đã chuyển đổi sang mô hình: {model_name}"

# Giao diện chọn mô hình
model_selector = gr.Dropdown(
    choices=model_list, 
    label="Chọn mô hình trích xuất đặc trưng", 
    value="vit_base_kprpe_webface4m", 
    type="value"
)

# Hàm để đăng ký người dùng mới
def enroll_new_person(pil_image, user_id):
    # Kiểm tra xem khuôn mặt đã tồn tại trong hệ thống chưa
    check_result = face_engine.get_id(pil_image)
    
    if check_result['status'] == "not_found":
        # Nếu khuôn mặt chưa có trong hệ thống, tiến hành đăng ký
        face_engine.save_to_db(pil_image, user_id)
        return f"Khuôn mặt đã đăng ký thành công với ID là {user_id}."
    
    # Nếu khuôn mặt đã có trong hệ thống
    return f"Khuôn mặt đã tồn tại với ID là {check_result['id']}."

# Giao diện đăng ký (enroll)
enroll_interface = gr.Interface(
    fn=enroll_new_person, 
    inputs=[gr.Image(), gr.Textbox(label="Nhập ID (ví dụ: CCCD)")], 
    outputs="text",
    title="Đăng ký khuôn mặt mới",
    description="Tải lên ảnh và nhập ID của bạn để đăng ký vào hệ thống."
)

# Hàm để xác minh hai khuôn mặt có giống nhau không
def verify(pil_image_1, pil_image_2):
    cosine_similarity = face_engine.compute_cosine_similarity(pil_image_1, pil_image_2)
    if cosine_similarity > face_engine.threshold:
        return f"Hai ảnh là cùng một người, điểm tương đồng: {cosine_similarity:.2f}"
    return f"Hai ảnh là hai người khác nhau, điểm tương đồng: {cosine_similarity:.2f}"

# Giao diện xác minh (verify)
verification_interface = gr.Interface(
    fn=verify, 
    inputs=[gr.Image(), gr.Image()], 
    outputs="text",
    title="Xác minh khuôn mặt",
    description="Tải lên hai ảnh để xác minh xem chúng có phải cùng một người hay không."
)

# Hàm để nhận diện khuôn mặt từ ảnh đầu vào
def identify(pil_image):
    check_result = face_engine.get_id(pil_image)
    
    if check_result['status'] == "not_found":
        return "Khuôn mặt không tồn tại trong hệ thống."
    
    # Nếu tìm thấy khuôn mặt
    return f"Khuôn mặt được nhận diện với ID là {check_result['id']}."

# Giao diện nhận diện (identify)
identification_interface = gr.Interface(
    fn=identify, 
    inputs=gr.Image(), 
    outputs="text",
    title="Nhận diện khuôn mặt",
    description="Tải lên một ảnh để nhận diện khuôn mặt từ hệ thống."
)

# Giao diện chính với các tab
demo = gr.TabbedInterface(
    interface_list=[
        gr.Interface(fn=update_model, inputs=model_selector, outputs="text", title="Chọn Mô Hình"),
        enroll_interface,
        verification_interface,
        identification_interface
    ], 
    tab_names=["Chọn Mô Hình", "Enroll", "Verification", "Identification"]
)

# Khởi chạy ứng dụng Gradio
if __name__ == "__main__":
    demo.launch()

engine3.py

pythona year ago
import os
import torch
import pandas as pd
from torchvision.transforms import Compose, ToTensor, Normalize
import inspect
from PIL import Image
import gc
import sys
from face_engine.huggingface_model_utils import load_model_from_local_path

# Kiểm tra thiết bị (GPU hoặc CPU)
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

class ModelLoader:
    def __init__(self, model_path: str):
        """Khởi tạo và load mô hình từ đường dẫn đã chỉ định."""
        self.__model = load_model_from_local_path(f'/home/quoc14/Code/HeThongNhanDang/model/{model_path}').to(DEVICE)
    
    def instance(self):
        """Trả về mô hình đã load."""
        return self.__model

class FaceEngine:
    def __init__(self, csv_file='/home/quoc14/Code/HeThongNhanDang/face_features_400.csv') -> None:
        """Khởi tạo FaceEngine và thiết lập file CSV dùng để lưu và so sánh."""
        self.__database_path = csv_file  # File CSV mặc định để lưu đặc trưng khuôn mặt
        self.__available_models = {
            "ir18_webface4m": "minchul/cvlface_adaface_ir18_webface4m",
            "ir50_webface4m": "minchul/cvlface_adaface_ir50_webface4m",
            "ir101_webface4m": "minchul/cvlface_adaface_ir101_webface4m",
            "ir101_webface12m": "minchul/cvlface_adaface_ir101_webface12m",
            "vit_base_kprpe_webface4m": "minchul/cvlface_adaface_vit_base_kprpe_webface4m",
            "vit_base_kprpe_webface12m": "minchul/cvlface_adaface_vit_base_webface12m"
        }
        self._extractor = None
        self._aligner = ModelLoader("minchul/cvlface_DFA_resnet50").instance()
        self.__model_name = "ir101_webface4m"  # Mặc định sử dụng mô hình này
        self.threshold = 0.3  # Ngưỡng để nhận diện
        self.load_csv()  # Tải cơ sở dữ liệu từ file CSV và xử lý các tensor đặc trưng

    def load_csv(self):
        """Load file CSV chứa các đặc trưng khuôn mặt và chuyển đổi đặc trưng từ JSON sang tensor."""
        if not os.path.exists(self.__database_path):
            # Tạo file CSV nếu chưa tồn tại
            columns = ['id', 'feat', 'name', 'image_path', 'date_registered']
            df = pd.DataFrame(columns=columns)
            df.to_csv(self.__database_path, index=False)
            self.__db = pd.DataFrame(columns=columns)  # Khởi tạo database rỗng
        else:
            # Chỉ load các cột cần thiết để tiết kiệm bộ nhớ
            self.__db = pd.read_csv(self.__database_path, dtype={'id': str}, usecols=['id', 'feat'])
            # Chuyển đổi tất cả các đặc trưng từ chuỗi JSON sang tensor một lần
            def convert_to_tensor(feat_str):
                try:
                    # Chuyển chuỗi JSON thành danh sách sau đó thành tensor
                    feat_list = eval(feat_str)
                    if isinstance(feat_list, list):
                        return torch.tensor(feat_list, device=DEVICE)
                    else:
                        raise ValueError("Feature is not a valid list.")
                except Exception as e:
                    print(f"Error converting feature: {e}")
                    return None
            
            self.__db['feat_tensor'] = self.__db['feat'].apply(convert_to_tensor)
            # Loại bỏ các hàng có đặc trưng không hợp lệ
            self.__db = self.__db.dropna(subset=['feat_tensor'])


    def reset_model(self):
        """Reset mô hình hiện tại và giải phóng bộ nhớ."""
        if hasattr(self, '_extractor'):
            del self._extractor
        
        # Dọn dẹp bộ nhớ và giải phóng GPU
        gc.collect()
        torch.cuda.empty_cache()

        # Xóa các module đã tải từ mô hình trước
        modules_to_delete = [key for key in sys.modules if "model" in key or "transformers" in key]
        for module in modules_to_delete:
            del sys.modules[module]

    def load_model(self, model_key):
        """Load mô hình dựa trên từ khóa model_key."""
        model_path = self.__available_models[model_key]
        print(f"Đang tải mô hình: {model_key} từ {model_path}")
        return ModelLoader(model_path).instance()

    def set_model(self, model_key):
        """Đặt mô hình hiện tại dựa trên lựa chọn của người dùng."""
        if model_key in self.__available_models:
            self.reset_model()  # Reset mô hình trước khi chuyển đổi
            self._extractor = self.load_model(model_key)  # Tải mô hình mới
            self.__model_name = model_key  # Cập nhật tên mô hình hiện tại
            print(f"Đã chuyển đổi sang mô hình: {model_key}")
        else:
            print(f"Mô hình {model_key} không tồn tại!")

    # Hàm chuyển đổi ảnh PIL thành tensor và chuẩn hóa
    def pil_to_input(self, pil_image):
        trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
        return trans(pil_image).unsqueeze(0).to(DEVICE)

    def get_feat(self, pil_image):
        """Trích xuất đặc trưng từ ảnh."""
        input_tensor = self.pil_to_input(pil_image)
        aligned_x, _, aligned_ldmks, _, _, _ = self._aligner(input_tensor)

        # Kiểm tra mô hình xem có cần keypoints không
        input_signature = inspect.signature(self._extractor.model.net.forward)
        if input_signature.parameters.get('keypoints') is not None:
            feat = self._extractor(aligned_x, aligned_ldmks)
        else:
            feat = self._extractor(aligned_x)

        return feat

    def compute_cosine_similarity(self, feat_input, feats_db):
        """Tính toàn bộ cosine similarity giữa đặc trưng đầu vào và tất cả các đặc trưng trong cơ sở dữ liệu."""
        return torch.nn.functional.cosine_similarity(feat_input, feats_db)

    def get_id(self, pil_image):
        """So sánh ảnh với cơ sở dữ liệu và trả về ID nếu tìm thấy."""
        if not os.path.exists(self.__database_path):
            return {"id": None, "status": "not_found"}

        # Trích xuất đặc trưng của ảnh đầu vào
        feat_input = self.get_feat(pil_image)

        # Chuyển danh sách đặc trưng của cơ sở dữ liệu thành tensor để tính toán
        feats_db = torch.stack(self.__db['feat_tensor'].values.tolist()).to(DEVICE)

        # Tính toán cosine similarity giữa ảnh đầu vào và tất cả các ảnh trong cơ sở dữ liệu
        similarities = self.compute_cosine_similarity(feat_input, feats_db)

        # Lấy ra ảnh có cosine lớn nhất
        max_sim, idx_max = similarities.max(0)
        
        if max_sim.item() > self.threshold:
            # Đã nhận diện thành công
            best_match_id = self.__db.iloc[idx_max.item()]['id']
            return {"id": best_match_id, "status": "found"}
        else:
            return {"id": None, "status": "not_found"}
    def save_to_db(self, pil_image, user_id):
        """Lưu đặc trưng của ảnh mới với ID người dùng."""
        feature = self.get_feat(pil_image)

        if not isinstance(feature, torch.Tensor):
            return "Feature extraction failed, not saving to database."

        # Chuyển đổi tensor sang danh sách để lưu thành chuỗi JSON
        feature_list = feature.squeeze().cpu().detach().numpy().tolist()

        new_row = pd.DataFrame({
            'id': [user_id], 
            'name': ["quoc"], 
            'feat': [feature_list],  # Lưu đặc trưng dưới dạng danh sách JSON
             # Để trống nếu không có
            'image_path': ["quoc"],  # Để trống nếu không có
            'date_registered': [pd.Timestamp.now()]
        })
        self.__db = pd.concat([self.__db, new_row], ignore_index=True)
        self.__db.loc[self.__db['id'] == user_id, 'feat_tensor'] = self.__db.loc[self.__db['id'] == user_id, 'feat'].apply(lambda x: torch.tensor(x, device=DEVICE))


        # Thay vì ghi lại toàn bộ file, chỉ ghi thêm dòng mới
        with open(self.__database_path, 'a') as f:
            new_row.to_csv(f, header=False, index=False)
        
        return f"Data for user {user_id} saved successfully."

  • Total 16 snippets
  • 1
  • 2
  • 3
  • 4