現在的圖標點選越來越多,如何一步解決,采用 YOLOX 目標檢測模式則可以輕松解決
要在 YOLOX 中使用按目錄分類的圖片數據集(每個目錄代表一個類別,目錄下是該類別的所有圖片),你需要進行以下配置步驟:
將需要檢測的圖標截取為獨立的圖片,放入到對應的目錄中,如 banner ,apple 等
1. 準備數據集結構
假設你的原始數據集結構如下:
dataset_root/class1/img1.jpgimg2.jpg...class2/img1.jpgimg2.jpg......
2. 轉換為 YOLOX 所需的格式
YOLOX 通常需要以下格式:
圖片文件對應的標注文件(每個圖片一個 .txt 文件)訓練/驗證集的劃分文件
創建 Python 轉換腳本
import os
import glob
from sklearn.model_selection import train_test_splitdef convert_directory_to_yolox(dataset_root, output_dir):# 獲取所有類別classes = [d for d in os.listdir(dataset_root) if os.path.isdir(os.path.join(dataset_root, d))]classes.sort()# 創建類別映射文件with open(os.path.join(output_dir, 'classes.txt'), 'w') as f:f.write('\n'.join(classes))# 收集所有圖片路徑和標簽image_paths = []labels = []for class_id, class_name in enumerate(classes):class_dir = os.path.join(dataset_root, class_name)for img_path in glob.glob(os.path.join(class_dir, '*')):if img_path.lower().endswith(('.png', '.jpg', '.jpeg')):image_paths.append(img_path)labels.append(class_id)# 劃分訓練集和驗證集train_img, val_img, train_lbl, val_lbl = train_test_split(image_paths, labels, test_size=0.2, random_state=42)# 創建標注文件和數據集列表def write_dataset(img_list, lbl_list, list_file):with open(list_file, 'w') as f_list:for img_path, label in zip(img_list, lbl_list):# 創建標注文件路徑base_name = os.path.splitext(os.path.basename(img_path))[0]txt_path = os.path.join(output_dir, 'labels', base_name + '.txt')# 寫入標注文件os.makedirs(os.path.dirname(txt_path), exist_ok=True)with open(txt_path, 'w') as f_txt:# YOLO格式: class_id x_center y_center width height# 對于全圖分類,bbox是整張圖片f_txt.write(f"{label} 0.5 0.5 1.0 1.0\n")# 寫入數據集列表f_list.write(f"{img_path}\n")os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)write_dataset(train_img, train_lbl, os.path.join(output_dir, 'train.txt'))write_dataset(val_img, val_lbl, os.path.join(output_dir, 'val.txt'))# 使用示例
convert_directory_to_yolox('path/to/your/dataset_root', 'path/to/yolox_dataset')
3. 配置 YOLOX
修改 exps 中的配置文件
找到或創建你的實驗配置文件(如 exps/example/yolox_s.py),修改以下部分:
class Exp(yolox_base.Exp):def __init__(self):super(Exp, self).__init__()self.num_classes = 你的類別數量self.data_dir = "轉換后的數據集路徑"self.train_ann = "train.txt"self.val_ann = "val.txt"self.input_size = (640, 640) # 根據你的需求調整self.test_size = (640, 640)
修改數據加載部分
確保你的數據加載器能夠處理這種格式。YOLOX 默認使用 MosaicDetection 數據集,你可能需要自定義:
from yolox.data import get_yolox_datadir
from yolox.data.datasets import DATASET_REGISTRY@DATASET_REGISTRY.register()
class YourDataset(Dataset):def __init__(self, ...):# 實現你的數據集加載邏輯pass
4. 訓練命令
使用修改后的配置運行訓練:
python tools/train.py -f exps/example/yolox_s.py -d 1 -b 64 --fp16 -o -c yolox_s.pth
注意事項
這種方法將整個圖像作為一個"邊界框"處理,適用于圖像分類任務如果你需要真正的目標檢測(圖像中有多個對象),需要為每個對象提供精確的邊界框標注考慮調整輸入尺寸以適應你的圖像比例對于分類任務,可能需要修改損失函數為更適合分類的損失(如交叉熵)
如果你需要真正的目標檢測而不是圖像分類,則需要為每個圖像提供包含對象位置和類別的詳細標注文件。