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. 數(shù)據(jù)集定義
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. 數(shù)據(jù)增強與加載
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])
])


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


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


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


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


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


# 5. 訓(xùn)練循環(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():  # 若凍結(jié)部分層，在此處設(shè)置
            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("訓(xùn)練完成！")

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

import clip
import torch
from PIL import Image


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


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


# 文本預(yù)處理：tokenize（轉(zhuǎn)換成模型可理解的格式）
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())


# 構(gòu)建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="檢索結(jié)果")
    btn.click(fn=search_images, inputs=query, outputs=gallery)


demo.launch()