裁剪實戰(zhàn))
找圖這件事說大不大說小也真不小。做設(shè)計的、寫文章的、搞電商的誰沒經(jīng)歷過在幾千張素材里翻來覆去、眼睛都快看瞎的時刻更別提那些我記得有一張圖特別合適但就是想不起來在哪個文件夾的崩潰瞬間。最近我花了兩周時間用YOLOv5加CLIP搭了一套本地找圖工具核心目標(biāo)就一個輸入一句話直接定位到圖還能順手把目標(biāo)區(qū)域裁出來。整個過程踩了不少坑也攢了一些在官方文檔里翻不到的經(jīng)驗今天全部攤開講。這套方案適合誰如果你手頭有大量圖片素材需要管理或者你是個開發(fā)者想了解視覺模型怎么落地到實際工具里再或者你只是好奇用文字搜圖到底是怎么實現(xiàn)的下面的內(nèi)容應(yīng)該都能給你一些可以直接抄作業(yè)的東西。我盡量不堆術(shù)語把每個環(huán)節(jié)為什么這么做、不這么做會怎樣都講清楚。1. 為什么是YOLOv5加CLIP這個組合1.1 兩個模型各管一攤分工明確先說清楚這兩個模型各自在干什么。YOLOv5負(fù)責(zé)找東西在哪它是一個目標(biāo)檢測模型你給它一張圖它能告訴你圖里有哪些物體、每個物體的邊界框坐標(biāo)是什么。CLIP負(fù)責(zé)理解你說的是什么它是一個圖文匹配模型能把文字和圖片映射到同一個向量空間里從而判斷一段文字和一張圖的相似度。為什么不是用一個模型搞定因為目標(biāo)檢測模型通常只能識別預(yù)設(shè)的幾十到幾百個類別你沒法用一句話讓它去找左上角那只橘貓旁邊的一個藍(lán)色杯子。而CLIP雖然能理解任意文字描述但它不擅長精確定位它只能告訴你這張圖和這句話像不像給不出精確的邊界框。所以把兩者串起來先用CLIP做粗篩從海量圖庫里找出可能相關(guān)的圖再用YOLOv5做精定位把目標(biāo)區(qū)域框出來并裁剪。這個流水線的設(shè)計思路是整個工具的核心骨架。我試過只用CLIP做檢索結(jié)果就是它能找到圖但你沒法自動裁剪因為CLIP不輸出坐標(biāo)。也試過只用YOLOv5那就只能按固定類別搜比如人車狗稍微復(fù)雜一點的描述就歇菜了。兩者結(jié)合之后體驗才真正到位。1.2 為什么不用更新的模型版本你可能會問YOLO都出到v8、v10了為什么還用v5CLIP也有更新的大版本為什么不用這個問題我認(rèn)真考慮過。YOLOv5雖然版本號看起來老但它的生態(tài)非常成熟文檔齊全導(dǎo)出ONNX、TensorRT的路徑都很順暢社區(qū)里能搜到的踩坑記錄也最多。對于一個要落地到本地工具的項目來說穩(wěn)定性比那一點點精度提升重要得多。CLIP這邊也是類似的情況?;A(chǔ)版CLIP的ViT-B/32模型只有約1.5億參數(shù)在消費級顯卡上就能跑得動推理速度也夠用。更大的模型確實精度更高但顯存占用和推理時間會成倍增長對于本地找圖這個場景來說性價比不劃算。我的建議是先用基礎(chǔ)版跑通全流程如果實測下來檢索準(zhǔn)確率不夠再考慮換更大的模型。1.3 整體流水線的數(shù)據(jù)流向把流程拆開看是這樣的用戶輸入一段文字描述系統(tǒng)先用CLIP的文本編碼器把文字轉(zhuǎn)成一個向量。同時圖庫里所有圖片在建立索引階段就已經(jīng)用CLIP的圖片編碼器轉(zhuǎn)成了向量存在一個向量數(shù)據(jù)庫里。檢索時計算文字向量和所有圖片向量的余弦相似度取Top-K張最相關(guān)的圖。接下來對這Top-K張圖逐張跑YOLOv5得到所有檢測框。然后關(guān)鍵的一步來了怎么判斷哪個框才是用戶想要的我的做法是把每個檢測框裁剪出來單獨送進(jìn)CLIP算一次圖文相似度得分最高的那個框就是目標(biāo)區(qū)域。最后按這個框的坐標(biāo)裁剪原圖輸出結(jié)果。整個鏈路聽起來不復(fù)雜但每個環(huán)節(jié)都有細(xì)節(jié)要摳。2. 環(huán)境搭建與依賴安裝的實操細(xì)節(jié)2.1 硬件選擇和顯存估算先說硬件。這套方案對顯卡有一定要求但不算離譜。我的測試環(huán)境是一張8GB顯存的顯卡跑ViT-B/32的CLIP加上YOLOv5s峰值顯存占用大約在3.5GB左右留了不少余量。如果你用的是6GB顯存的卡把YOLOv5換成nano版本YOLOv5nCLIP保持ViT-B/32也能跑起來峰值大概2.8GB。CPU能不能跑能但檢索速度會慢很多。我實測過純CPU環(huán)境對1000張圖建索引大概需要8到10分鐘檢索一次含YOLOv5推理大概3到5秒。如果你圖庫不大、對速度不敏感CPU方案也不是不能用。但如果圖庫上萬張還是建議上顯卡。內(nèi)存方面建議至少16GB。因為建索引的時候需要把圖片批量加載、預(yù)處理內(nèi)存不夠會頻繁觸發(fā)磁盤交換速度斷崖式下跌。2.2 Python環(huán)境和關(guān)鍵依賴版本Python版本我推薦3.8到3.10之間。3.11以上有些庫的wheel包還不全裝起來容易出問題。虛擬環(huán)境一定要建這個不用多說避免污染全局環(huán)境。關(guān)鍵依賴的版本搭配是個坑。我踩過的坑是torch版本和torchvision版本不匹配導(dǎo)致CLIP加載時報了一堆莫名其妙的錯誤。經(jīng)過幾次嘗試下面這套組合實測最穩(wěn)pip install torch1.13.1 torchvision0.14.1 pip install githttps://github.com/openai/CLIP.git pip install ultralytics8.0.200 pip install faiss-cpu1.7.4 pip install opencv-python4.8.1.78 pip install Pillow10.1.0 pip install numpy1.24.3這里解釋幾個選擇。faiss是Facebook開源的向量檢索庫用來做相似度搜索比用numpy暴力算快得多尤其是圖庫上萬張之后差距非常明顯。ultralytics這個包封裝了YOLOv5的推理接口用起來比直接調(diào)原始倉庫方便。opencv用來做圖片的讀取和裁剪。注意faiss-cpu和faiss-gpu只能裝一個別同時裝會沖突。如果你有顯卡并且圖庫很大可以裝faiss-gpu檢索速度能快好幾倍。2.3 模型權(quán)重下載與存放YOLOv5的權(quán)重文件在ultralytics包里會自動下載第一次運行時會從網(wǎng)上下載yolov5s.pt大概14MB。CLIP的權(quán)重也是首次運行時自動下載ViT-B/32大概350MB。這兩個下載過程在國內(nèi)網(wǎng)絡(luò)環(huán)境下可能會比較慢建議提前手動下載好放到緩存目錄里。CLIP的緩存目錄默認(rèn)在用戶主目錄下的.cache/clip文件夾YOLOv5的權(quán)重默認(rèn)在當(dāng)前工作目錄。如果你在離線環(huán)境部署需要提前把這些文件準(zhǔn)備好然后通過代碼指定本地路徑加載。3. 圖庫索引構(gòu)建把圖片變成可搜索的向量3.1 圖片預(yù)處理中的尺寸陷阱CLIP的圖片編碼器對輸入尺寸有要求ViT-B/32接受的是224x224的輸入。但這里有個容易忽略的點CLIP的預(yù)處理不只是縮放還包括中心裁剪和歸一化。如果你直接用自己的方式把圖片縮放到224x224而不走CLIP官方的preprocess流程檢索效果會明顯下降。正確的做法是使用CLIP自帶的preprocess函數(shù)import clip from PIL import Image model, preprocess clip.load(ViT-B/32, devicecuda) image Image.open(test.jpg).convert(RGB) image_input preprocess(image).unsqueeze(0).to(cuda)這個preprocess會做三件事把圖片短邊縮放到224、中心裁剪出224x224、按ImageNet的均值和標(biāo)準(zhǔn)差做歸一化。三步缺一不可。我一開始圖省事自己寫了個resize結(jié)果檢索準(zhǔn)確率掉了將近15個百分點后來換成官方預(yù)處理才恢復(fù)正常。3.2 批量編碼時的顯存控制建索引的時候如果一張一張圖編碼速度太慢。但一次性把所有圖都塞進(jìn)顯存又會爆顯存。我的做法是分批處理每批32張。這個數(shù)字不是拍腦袋定的我實測過批大小16的時候8GB顯存利用率只有60%左右浪費了批大小64的時候偶爾會爆32是一個比較穩(wěn)的平衡點。import torch import numpy as np def encode_images(image_paths, model, preprocess, batch_size32): all_features [] for i in range(0, len(image_paths), batch_size): batch_paths image_paths[i:ibatch_size] batch_images [] for p in batch_paths: try: img Image.open(p).convert(RGB) batch_images.append(preprocess(img)) except Exception as e: print(f跳過損壞圖片 {p}: {e}) continue if not batch_images: continue batch_tensor torch.stack(batch_images).to(cuda) with torch.no_grad(): features model.encode_image(batch_tensor) features features / features.norm(dim-1, keepdimTrue) all_features.append(features.cpu().numpy()) return np.concatenate(all_features, axis0)注意最后那行歸一化把特征向量除以它的L2范數(shù)。這一步很關(guān)鍵因為后面算余弦相似度的時候如果向量已經(jīng)歸一化了就只需要算點積速度快很多而且數(shù)值更穩(wěn)定。3.3 向量數(shù)據(jù)庫的選型和持久化圖庫不大的話直接用faiss的IndexFlatIP就行IP是內(nèi)積的意思配合歸一化后的向量等價于余弦相似度。這個索引結(jié)構(gòu)簡單、精確不需要訓(xùn)練。缺點是內(nèi)存占用大每張圖占512個float32也就是2KB。一萬張圖就是20MB十萬張圖200MB還能接受。如果圖庫超過百萬級別可以考慮IndexIVFFlat它通過聚類加速檢索但會損失一點點精度。對于個人找圖工具來說IndexFlatIP完全夠用。索引建好之后要持久化不然每次啟動都重新編碼一遍太浪費時間。faiss提供了write_index和read_indeximport faiss index faiss.IndexFlatIP(512) index.add(features) faiss.write_index(index, image_index.faiss) # 讀取 index faiss.read_index(image_index.faiss)同時還要把圖片路徑列表存下來因為faiss只存向量不存路徑。我用一個json文件存路徑列表順序和索引里的向量一一對應(yīng)。這里有個坑如果你中途往圖庫里加了新圖一定要同時更新索引和路徑列表否則順序錯亂檢索出來的圖就對不上了。4. 檢索與裁剪從文字到精確邊界框4.1 文本編碼的細(xì)節(jié)和常見誤區(qū)文本編碼看起來簡單但有幾個細(xì)節(jié)要注意。CLIP的文本編碼器對輸入的token長度有限制最多77個token。如果你的描述很長會被截斷。所以描述盡量精簡把核心關(guān)鍵詞放在前面。另外CLIP的文本編碼器輸出的特征也需要歸一化和圖片特征一樣。很多人忘了這一步導(dǎo)致相似度計算出來的值范圍不對排序就亂了。def encode_text(text, model): text_input clip.tokenize([text]).to(cuda) with torch.no_grad(): text_features model.encode_text(text_input) text_features text_features / text_features.norm(dim-1, keepdimTrue) return text_features.cpu().numpy()還有一個經(jīng)驗CLIP對英文的支持比中文好很多。如果你的描述是中文建議先翻譯成英文再編碼。我實測過同一張圖用中文描述檢索的準(zhǔn)確率比英文低大概20個百分點。這不是CLIP本身的問題而是訓(xùn)練數(shù)據(jù)里中文占比少。所以我的工具里加了一個簡單的翻譯層把中文描述轉(zhuǎn)成英文再送進(jìn)CLIP。4.2 相似度排序和Top-K的選擇檢索的時候把文本向量和索引里所有圖片向量做內(nèi)積得到一個分?jǐn)?shù)數(shù)組然后取分?jǐn)?shù)最高的K張圖。K取多少合適我的經(jīng)驗是取10到20之間。太小了可能漏掉目標(biāo)太大了后面YOLOv5推理的負(fù)擔(dān)重。def search_images(text, index, image_paths, model, top_k15): text_feat encode_text(text, model) scores, indices index.search(text_feat, top_k) results [] for score, idx in zip(scores[0], indices[0]): results.append((image_paths[idx], float(score))) return results這里返回的score是余弦相似度范圍在-1到1之間。一般來說分?jǐn)?shù)超過0.3就說明有一定相關(guān)性了超過0.5就是比較匹配了。但這個閾值不是絕對的取決于你的圖庫內(nèi)容和描述的具體程度。4.3 YOLOv5檢測與CLIP二次打分的配合拿到Top-K張圖之后逐張跑YOLOv5。YOLOv5會返回一堆檢測框每個框有坐標(biāo)、置信度和類別。但問題是YOLOv5檢測出來的框可能有很多個哪個才是用戶想要的我的做法是把每個檢測框裁剪出來單獨送進(jìn)CLIP算一次圖文相似度得分最高的那個框就是目標(biāo)。這個二次打分的思路是整個工具里最巧妙的一環(huán)。它把YOLOv5的定位能力和CLIP的理解能力真正結(jié)合起來了。from ultralytics import YOLO def detect_and_rank(image_path, text_feat, yolo_model, clip_model, preprocess): results yolo_model(image_path) boxes results[0].boxes best_box None best_score -1 img Image.open(image_path).convert(RGB) for box in boxes: x1, y1, x2, y2 box.xyxy[0].cpu().numpy() cropped img.crop((x1, y1, x2, y2)) cropped_input preprocess(cropped).unsqueeze(0).to(cuda) with torch.no_grad(): crop_feat clip_model.encode_image(cropped_input) crop_feat crop_feat / crop_feat.norm(dim-1, keepdimTrue) score float((crop_feat text_feat.T).cpu().numpy()[0][0]) if score best_score: best_score score best_box (x1, y1, x2, y2) return best_box, best_score這段代碼里有個細(xì)節(jié)裁剪出來的區(qū)域可能很小比如只有幾十個像素直接送進(jìn)CLIP效果不好。我的處理是如果裁剪區(qū)域小于64x64就把它放大到224x224再送進(jìn)去。這個放大操作對檢索準(zhǔn)確率有肉眼可見的提升。4.4 裁剪輸出的邊界處理裁剪的時候要注意邊界問題。YOLOv5返回的坐標(biāo)可能是浮點數(shù)而且可能超出圖片邊界。裁剪之前要把坐標(biāo) clamp 到圖片范圍內(nèi)否則PIL會報錯。def safe_crop(img, box): x1, y1, x2, y2 box w, h img.size x1 max(0, min(x1, w)) y1 max(0, min(y1, h)) x2 max(0, min(x2, w)) y2 max(0, min(y2, h)) if x2 x1 or y2 y1: return None return img.crop((x1, y1, x2, y2))另外有時候用戶想要的不是裁剪而是知道目標(biāo)在圖里的位置。所以我的工具里加了一個開關(guān)可以選擇只定位不裁剪或者定位并裁剪。定位模式下在原圖上畫一個矩形框輸出方便用戶確認(rèn)。5. 實測中暴露的問題和調(diào)優(yōu)經(jīng)驗5.1 檢索準(zhǔn)確率不夠時的排查思路工具跑通之后我發(fā)現(xiàn)有些描述檢索效果很好有些就很差。排查了一圈總結(jié)出幾個常見原因。第一個原因是描述太抽象。比如一張好看的圖這種描述CLIP根本沒法理解因為它沒有具體的視覺概念。解決辦法是引導(dǎo)用戶用具體的名詞和形容詞比如藍(lán)天下的白色建筑。第二個原因是圖庫里沒有相關(guān)內(nèi)容。這個聽起來像廢話但實際中經(jīng)常發(fā)生。用戶以為圖庫里有某張圖其實沒有檢索出來的都是勉強相關(guān)的。我的做法是在返回結(jié)果里顯示相似度分?jǐn)?shù)分?jǐn)?shù)低于0.2的直接標(biāo)注低置信度提醒用戶可能沒有匹配項。第三個原因是圖片本身質(zhì)量差。模糊、過暗、過曝的圖片CLIP編碼出來的特征質(zhì)量也差。對于這種情況可以在建索引階段做一個簡單的質(zhì)量過濾把分辨率低于一定閾值或者方差過低的圖片跳過。5.2 推理速度的優(yōu)化手段速度優(yōu)化有幾個方向。第一是模型量化把CLIP和YOLOv5都轉(zhuǎn)成FP16精度推理速度能提升30%到50%精度損失很小。第二是批處理檢索的時候如果Top-K是15可以把這15張圖攢成一個批次一起送進(jìn)YOLOv5比一張一張跑快很多。第三是索引層面的優(yōu)化。如果圖庫特別大可以用faiss的GPU版本檢索速度從毫秒級降到微秒級。不過對于個人工具來說CPU版的faiss已經(jīng)夠快了瓶頸主要在YOLOv5推理上。我實測的數(shù)據(jù)1000張圖的圖庫建索引CLIP編碼耗時約2分鐘檢索一次CLIP文本編碼faiss搜索YOLOv5推理15張圖耗時約1.5秒。10000張圖的圖庫建索引約20分鐘檢索耗時約2秒因為faiss搜索本身很快多出來的時間主要是YOLOv5推理的圖片數(shù)量沒變。5.3 幾個容易忽略的工程細(xì)節(jié)第一個細(xì)節(jié)是圖片路徑的處理。Windows和Linux的路徑分隔符不一樣如果代碼里硬編碼了反斜杠換到Linux上就掛了。統(tǒng)一用os.path或者pathlib來處理。第二個細(xì)節(jié)是中文路徑。OpenCV讀取中文路徑會失敗需要用numpy先讀文件再解碼?;蛘吒纱嘤肞IL讀圖PIL對中文路徑的支持更好。第三個細(xì)節(jié)是內(nèi)存泄漏。如果在循環(huán)里反復(fù)創(chuàng)建PIL Image對象而不釋放內(nèi)存會慢慢漲上去。我的做法是每處理完一批圖片手動調(diào)用gc.collect()雖然有點粗暴但確實有效。第四個細(xì)節(jié)是異常處理。圖庫里難免有損壞的圖片讀取的時候要try-except包起來跳過壞圖繼續(xù)處理而不是讓整個程序崩掉。6. 這套方案還能怎么擴展6.1 加入以圖搜圖的能力現(xiàn)在的工具是文字搜圖但其實CLIP的圖片編碼器也可以用來做以圖搜圖。用戶上傳一張參考圖編碼成向量然后在索引里搜索最相似的圖片。這個功能加進(jìn)來只需要幾行代碼因為圖片編碼的流程已經(jīng)現(xiàn)成了。擴展之后工具就同時支持文字搜圖和以圖搜圖兩種模式。實際使用中以圖搜圖適合我有一張類似的圖想找更多同風(fēng)格的文字搜圖適合我腦子里有個畫面想找對應(yīng)的圖。6.2 多標(biāo)簽檢索和組合條件CLIP的一個局限是它只能處理單一描述。如果你想搜藍(lán)天并且有建筑并且是日落時分單一向量可能表達(dá)不了這么復(fù)雜的組合。一個思路是把描述拆成多個短句分別編碼然后對相似度分?jǐn)?shù)做加權(quán)平均。另一個思路是用更高級的模型比如支持多標(biāo)簽的檢索方案。我目前的做法是讓用戶用逗號分隔多個關(guān)鍵詞每個關(guān)鍵詞單獨編碼然后取平均向量。實測下來比用一整句話編碼效果略好尤其是當(dāng)關(guān)鍵詞之間關(guān)聯(lián)性不強的時候。6.3 和本地文件管理器的聯(lián)動工具最終是要落到實際使用場景里的。我的做法是檢索結(jié)果直接輸出圖片的完整路徑并且提供一個在文件管理器中打開的按鈕。這樣用戶找到圖之后可以立刻在系統(tǒng)里進(jìn)行后續(xù)操作比如復(fù)制、移動、重命名。另外我還加了一個批量導(dǎo)出功能把檢索到的圖片和裁剪結(jié)果復(fù)制到一個指定文件夾里方便用戶集中處理。這個功能在做素材整理的時候特別實用。整套工具從構(gòu)思到跑通大概花了兩周時間其中大部分時間不是在寫代碼而是在調(diào)參數(shù)、試模型、處理各種邊界情況。如果你也想搭一套類似的工具我的建議是先把最小可用版本跑起來哪怕檢索準(zhǔn)確率只有60%也比一開始就追求完美要強。跑通之后再逐步優(yōu)化每一步的改進(jìn)都能立刻看到效果這種正反饋會讓整個過程輕松很多。最后分享一個我在實際使用中總結(jié)的小技巧建索引的時候把圖片的文件名也一起存進(jìn)元數(shù)據(jù)里。因為有時候檢索出來的圖光看內(nèi)容想不起來是哪個項目的但一看文件名就立刻明白了。這個細(xì)節(jié)雖小但在實際使用中能省不少事。