import cv2
import torch
import numpy as np
from torchvision import transforms
from PIL import Image
from u2net import U2NET  

# Load U2NET model
model = U2NET(3, 1)
model_path = 'checkpoints/u2net.pth' 
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()

# Tiền xử lý ảnh
def preprocess(img):
    transform = transforms.Compose([
        transforms.Resize((192, 256)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])])
    return transform(img).unsqueeze(0)

# Tạo mặt nạ từ ảnh đầu vào
def get_mask(image_path):
    image = Image.open(image_path).convert('RGB')
    input_tensor = preprocess(image)
    with torch.no_grad():
        d1, *_ = model(input_tensor)
        pred = d1[0][0].numpy()
        pred = (pred - pred.min()) / (pred.max() - pred.min())
        mask = cv2.resize(pred, image.size)
        return np.expand_dims(mask, axis=2)

# Áp dụng mặt nạ và đổi nền thành trắng
def apply_mask(img_path, output_path='output_2.png'):
    original = cv2.imread(img_path)
    original_rgb = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
    mask = get_mask(img_path)
    mask = cv2.resize(mask, (original.shape[1], original.shape[0]))
    mask = np.expand_dims(mask, axis=2)
    mask = np.repeat(mask, 3, axis=2)
    result = (original_rgb * mask + 255 * (1 - mask)).astype(np.uint8)
    result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
    result_bgr = cv2.resize(result_bgr, (192, 256))    
    cv2.imwrite(output_path, result_bgr)
    print(f"[✔] Đã lưu ảnh tại: {output_path}")

# === CHẠY THỬ ===
apply_mask('test_2.jpg')  # ⚠️ Thay bằng ảnh thật của bạn