import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import clip
import os
from tqdm import tqdm


# 1. 數據集定義
class CLIPDataset(Dataset):
    def __init__(self, image_dir, text_file, transform=None):
        """
        image_dir: 圖像文件夾路徑
        text_file: 文本文件路徑，每行格式為"圖像文件名,文本描述"
        """
        self.image_dir = image_dir
        self.transform = transform
        self.data = []
        with open(text_file, 'r', encoding='utf-8') as f:
            for line in f:
                img_name, text = line.strip().split(',', 1)
                self.data.append((img_name, text))


    def __len__(self):
        return len(self.data)


    def __getitem__(self, idx):
        img_name, text = self.data[idx]
        img_path = os.path.join(self.image_dir, img_name)
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, text


# 2. 數據增強與加載
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, cnotallow=0.2, saturatinotallow=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


# 假設訓練數據存放在./train_images，文本文件為./train_texts.txt
dataset = CLIPDataset(
    image_dir='./train_images',
    text_file='./train_texts.txt',
    transform=transform
)
dataloader = DataLoader(
    dataset,
    batch_size=256,  # 根據GPU顯存調整
    shuffle=True,
    num_workers=8,
    pin_memory=True
)


# 3. 加載預訓練編碼器（或初始化）
device = "cuda" if torch.cuda.is_available() else "cpu"
image_encoder, text_encoder = clip.load("ViT-B/32", device=device)


# 若從頭訓練，可初始化模型
# from clip.model import VisionTransformer, TextTransformer
# image_encoder = VisionTransformer(...).to(device)
# text_encoder = TextTransformer(...).to(device)


# 4. 定義損失函數與優(yōu)化器
temperature = 0.07  # 溫度參數
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(
    list(image_encoder.parameters()) + list(text_encoder.parameters()),
    lr=1e-4,
    weight_decay=1e-6  # L2正則化
)


# 學習率衰減
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)


# 5. 訓練循環(huán)
num_epochs = 30
for epoch in range(num_epochs):
    image_encoder.train()
    text_encoder.train()
    total_loss = 0.0


    for images, texts in tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}"):
        images = images.to(device)
        texts = clip.tokenize(texts).to(device)  # 文本tokenize


        # 編碼圖像和文本
        with torch.no_grad():  # 若凍結部分層，在此處設置
            image_features = image_encoder(images)
            text_features = text_encoder(texts)


        # L2歸一化
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)


        # 計算相似度矩陣
        logits = (image_features @ text_features.T) / temperature


        # 計算對比損失（雙向）
        labels = torch.arange(logits.shape[0], device=device)
        loss = (criterion(logits, labels) + criterion(logits.T, labels)) / 2


        # 反向傳播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


        total_loss += loss.item() * images.size(0)


    # 計算平均損失
    avg_loss = total_loss / len(dataset)
    print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")
    scheduler.step()


    # 保存模型
    if (epoch + 1) % 5 == 0:
        torch.save({
            'image_encoder': image_encoder.state_dict(),
            'text_encoder': text_encoder.state_dict(),
            'optimizer': optimizer.state_dict()
        }, f"clip_epoch_{epoch+1}.pth")


print("訓練完成！")

pip install torch ftfy regex tqdm  # 基礎依賴
pip install git+https://github.com/openai/CLIP.git  # CLIP庫
pip install gradio faiss-cpu  # 可視化與向量檢索

import clip
import torch
from PIL import Image


# 加載模型和預處理函數
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)


# 圖片預處理：縮放、歸一化等
def process_image(img_path):
    image = Image.open(img_path).convert("RGB")
    return preprocess(image).unsqueeze(0).to(device)


# 文本預處理：tokenize（轉換成模型可理解的格式）
def process_text(text):
    return clip.tokenize([text]).to(device)

import faiss
import os
import numpy as np


# 圖片文件夾路徑
img_dir = "path/to/your/images"
img_paths = [os.path.join(img_dir, f) for f in os.listdir(img_dir) if f.endswith(('jpg', 'png'))]


# 提取所有圖片的向量
img_features = []
for path in img_paths:
    img = process_image(path)
    with torch.no_grad():
        feat = model.encode_image(img)  # 圖像編碼
        feat /= feat.norm(dim=-1, keepdim=True)  # 歸一化
        img_features.append(feat.cpu().numpy())


# 構建FAISS索引
d = 512  # CLIP向量維度通常為512
index = faiss.IndexFlatL2(d)  # 用L2距離衡量相似度
index.add(np.vstack(img_features))  # 加入所有圖片向量

import gradio as gr


def search_images(query):
    # 文本編碼
    text = process_text(query)
    with torch.no_grad():
        text_feat = model.encode_text(text)
        text_feat /= text_feat.norm(dim=-1, keepdim=True)


    # 檢索最相似的5張圖
    _, indices = index.search(text_feat.cpu().numpy(), k=5)
    return [img_paths[i] for i in indices[0]]


# 搭建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("## ?? CLIP以文搜圖演示")
    query = gr.Textbox(label="輸入文本描述（如'a red dress'）")
    btn = gr.Button("搜索")
    gallery = gr.Gallery(label="檢索結果")
    btn.click(fn=search_images, inputs=query, outputs=gallery)


demo.launch()