YoloX를 사용하여 커스텀 데이터셋을 학습하는 과정을 정리했습니다.
전체 과정은 아래 유튜브 영상에서도 보실 수 있습니다. 라벨링 하는 과정은 유튜브 영상을 참고하세요.
Yolo가 여러 버전이 있지만 라이센스 문제로 상용으로 자유롭게 사용가능한 것은 YoloX입니다.
2025. 5. 20 최초작성
우분투 환경에서 진행해야 필요한 패키지 설치시 문제가 발생하지 않았습니다. 윈도우에서 시도시 문제가 발생했었습니다.
다음 포스트를 참고하여 파이썬 개발환경을 구축하여 사용합니다.
Visual Studio Code와 Miniconda를 사용한 Python 개발 환경 만들기( Windows, Ubuntu, WSL2)
1. Python 가상 환경을 생성하여 진행합니다.
Python 3.9를 사용해야 requirements.txt를 사용한 패키지 설치시 문제없이 진행가능했습니다.
$ conda create -n yolo python=3.9
$ conda activate yolo
2. 깃허브 저장소를 가져옵니다.
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
3. 해당 YOLOX 디렉토리로 이동하여 필요한 패키지를 설치합니다.
cd YOLOX
pip install -r requirements.txt
일부 패키지가 설치안되는 경우 추가로 설치하면 됩니다.
4. PyTorch에서 CUDA를 지원하지 않도록 설치되어 있어서 Fasle가 출력되었습니다.
python -c "import torch; print(torch.cuda.is_available())"
False
False가 나왔다면 CUDA가 지원되도록 Pytorch를 다시 설치합니다. CUDA 11.8로는 동작에 문제가 있어 CUDA 12.4를 사용하도록 설치했습니다.
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
다음처럼 True가 나와야 PyTorch에서 GPU가속이 가능합니다.
python -c "import torch; print(torch.cuda.is_available())"
True
5. 이제 YOLOX 설치를 진행합니다.
pip3 install -v -e .
6. 미리 학습된 YOLOX weights를 다운로드합니다.
wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth
6. 이제 예제를 실행해봅니다.
python tools/demo.py image -f ./exps/default/yolox_s.py -c ./yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu
실행 결과 마지막에 검출 결과가 저장된 이미지의 위치가 출력됩니다.
Saving detection result in ./YOLOX_outputs\yolox_m\vis_res\2025_04_06_22_12_19\dog.jpg
디폴트로 설치된 이미지 뷰어인 eog를 사용하여 결과 이미지를 터미널에서 확인합니다.
eog ./YOLOX_outputs/yolox_m/vis_res/2025_04_04_15_18_44/dog.jpg
7. 이제 커스텀 데이터를 사용해봅니다.
Labelme을 사용하여 이미지에 라벨을 부여합니다. 여기에선 scissors(가위) 라벨을 부여합니다.
자세한 과정은 유튜브 영상을 참고하세요.
영상에선 다루지 않지만 라벨링을 하지 않은 이미지를 추가해야 오탐지를 줄일 수 있다고합니다.
네거티브 샘플(Negative Samples) 추가:
가장 효과적인 방법 중 하나입니다. 탐지하려는 객체가 없는 배경 이미지나, 오탐지를 일으키는 객체들만 포함된 이미지를 학습 데이터셋에 추가합니다.
이때, 해당 이미지에는 어떤 바운딩 박스(Bounding Box) 정보도 포함하지 않아야 합니다.
이렇게 하면 모델은 '이런 이미지는 탐지 대상이 아니다'라는 것을 명시적으로 학습하게 되어 오탐지를 줄일 수 있습니다.
Labelme를 설치합니다.
pip install labelme
실행합니다.
labelme
다음처럼 에러나는경우
(yolo) webnautes@webnautes-Legion-5-15ACH6H:~/YOLOX$ labelme
2025-04-09 12:24:04.781 | INFO | labelme.config:get_config:66 - Loading config file from: /home/webnautes/.labelmerc
QObject::moveToThread: Current thread (0x20a23740) is not the object's thread (0x21ab6f80).
Cannot move to target thread (0x20a23740)
qt.qpa.plugin: Could not load the Qt platform plugin "xcb" in "/home/webnautes/miniconda3/envs/yolo/lib/python3.9/site-packages/cv2/qt/plugins" even though it was found.
This application failed to start because no Qt platform plugin could be initialized. Reinstalling the application may fix this problem.
Available platform plugins are: xcb, eglfs, linuxfb, minimal, minimalegl, offscreen, vnc, wayland-egl, wayland, wayland-xcomposite-egl, wayland-xcomposite-glx, webgl.
중지됨 (코어 덤프됨)
다음 주소를 로그에서 찾아서 해당 디렉토리를 삭제하고 다시 labelme를 실행합니다. .
(yolo) webnautes@webnautes-Legion-5-15ACH6H:~/YOLOX$ rm -rf home/webnautes/miniconda3/envs/yolo/lib/python3.9/site-packages/cv2/qt/plugins
라벨링된 폴더와 테스트용 이미지 파일을 복사해옵니다. labelme_data 폴더를 YOLOX 폴더로 가져오면 됩니다.
다음 링크에 있는 코드를 labelme_to_coco.py로 저장하여 실행합니다.
$ python labelme_to_coco.py
변환 완료! 총 23개 이미지가 COCO 형식으로 변환되었습니다.
- 훈련 세트: 18개
- 검증 세트: 5개
COCO 구조가 datasets/coco128에 생성되었습니다.
다음과 같은 디렉토리 구조를 갖게됩니다.
$ ls datasets/coco128/
annotations/ train2017/ val2017/
다음 파일을 수정합니다.
./exps/example/custom/yolox_s.py
수정내용
클래스 개수
self.num_classes = 1
사양이 낮아서 줄임
self.data_num_workers = 1
디스크 공간 부족으로 평가하는 횟수를 줄여야 중간 모델 저장횟수가 줄어듭니다.
self.eval_interval = 50
yolox/data/datasets/coco_classes.py 파일을 수정하여 다음처럼 저장합니다. labelme에서 지정한 이름을 사용하세요.
COCO_CLASSES = (
"scissors",
)
학습을 진행해봅니다.
특별한 경우를 제외하고는 항상 모델을 초기화할 때 COCO 사전 학습된 가중치를 사용하는 것이 좋습니다.
Exp 파일과 제공된 COCO 사전 학습 가중치를 받으면 다음 명령어를 사용하여 모델을 학습시킬 수 있습니다:
python tools/train.py -f ./exps/example/custom/yolox_s.py -d 1 -b 8 --fp16 -o -c yolox_s.pth
이제 학습이 완료된 모델을 테스트 해봅니다. 테스트시 conf 값을 높이면 더 신뢰있는 물체만 감지합니다.
python tools/demo.py image -f ./exps/example/custom/yolox_s.py -c ./YOLOX_outputs/yolox_s/best_ckpt.pth --path 1.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu
python tools/demo.py image -f ./exps/example/custom/yolox_s.py -c ./YOLOX_outputs/yolox_s/best_ckpt.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu
에러가 발생합니다.
체크포인트 파일의 출처를 신뢰할 수 있다면 (직접 학습했거나 공식적인 경로로 받았다면), torch.load 함수 호출 시 weights_only=False 옵션을 명시적으로 추가해줘야 합니다.
tools/demo.py 파일을 엽니다.
에러가 발생한 283번째 줄 근처의 코드를 찾습니다.
ckpt = torch.load(ckpt_file, map_location="cpu")
이 부분을 다음과 같이 수정합니다.
ckpt = torch.load(ckpt_file, map_location="cpu", weights_only=False)
다시 위 명령을 실행하면 잘 동작합니다.
참고한 곳.
https://yolox.readthedocs.io/en/latest/quick_run.html#installation
https://github.com/Megvii-BaseDetection/YOLOX/blob/main/README.md
https://github.com/Megvii-BaseDetection/YOLOX/blob/main/docs/train_custom_data.md
Member discussion