딥러닝용 컨테이너 세팅

김민준·2025년 5월 13일

ros 제외한것들이 설치 된 이미지 사용

sudo docker run -it \
  --name dlearning_container \
  -v /home/hi/container_train_videos:/workspace/container_train_videos \
  --net=host \
  --privileged \
  --env="DISPLAY=$DISPLAY" \
  --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" \
  --device=/dev/video0 \
  --gpus all \
  heimerdinger123/kimminjun_cudapytorchcudnn22.04:1.0 \
  bash


# 컨테이너 안에서
cd /workspace

apt update && apt install -y \
  libsm6 libice6 libx11-xcb1 libxcb1 libxcb-shape0 libxcb-shm0 \
  libxcb-xfixes0 libxcb-randr0 libxcb-image0 libxcb-keysyms1 \
  libxcb-icccm4 libxcb-sync1 libxcb-xinerama0 libxkbcommon-x11-0 \
  libgl1-mesa-glx libglib2.0-0 libxext6 libxrender1 libx11-dev



echo 'export QT_QPA_PLATFORM_PLUGIN_PATH=/usr/local/lib/python3.10/dist-packages/cv2/qt/plugins/platforms' >> ~/.bashrc

source ~/.bashrc


# 감시 스크립트 복사 후 테스트
from ultralytics import YOLO
# 사전 학습된 모델 로드
model = YOLO("yolo11n.pt")
# CUDA 0번 GPU에서 학습
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=0)
# 여러 GPU 사용 예시 (0번, 1번 GPU 사용)
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])

또는 CLI에서는 다음과 같이 사용할 수 있습니다.
yolo detect train data=coco8.yaml model=yolo11n.pt epochs=100 imgsz=640 device=0
yolo detect train data=coco8.yaml model=yolo11n.pt epochs=100 imgsz=640 device=0,1

추론(예측) 예제는 아래와 같이 CUDA 장치를 지정하여 실행할 수 있습니다.


from ultralytics import YOLO
model = YOLO("yolo11n.pt")
results = model.predict(source="이미지경로.jpg", device=0)

from ultralytics import YOLO

model = YOLO('yolo11n.pt')
model.predict(source='/workspace/container_train_videos/vid_20250513_103219.mp4', show=True, device=0)


# train_watch_folder.py - YOLOv10 학습 자동화 (mp4 → 이미지 → 학습)

import os
import time
import glob
import subprocess
from datetime import datetime
import shutil

WATCH_DIR = "/workspace/container_train_videos"
BASE_DATASET_DIR = "/workspace/datasets"
YOLO_DIR = "/workspace/yolov10"
LOG_FILE = "/workspace/logs/training_log.txt"
PROCESSED_DIR = "/workspace/processed_videos"



def log(msg):
    os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
    with open(LOG_FILE, "a") as f:
        f.write(f"[{datetime.now()}] {msg}\n")
    print(msg)


def extract_frames(video_path, output_dir, fps=5):
    os.makedirs(output_dir, exist_ok=True)
    cmd = [
        "ffmpeg", "-i", video_path,
        "-qscale:v", "2",
        "-vf", f"fps={fps}",
        os.path.join(output_dir, "frame_%04d.jpg")
    ]
    subprocess.run(cmd, check=True)

#언제 처리된 파일인지 구분하기 위해서 base_dir인자로 년도 날짜 시각 등을 받아와서
#이걸 앞에 붙여서 언제 데이터기반한 데이터셋인지 구별할 수 있도록  
def create_dataset_structure(base_dir):
    os.makedirs(os.path.join(base_dir, "train/images"), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "train/labels"), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "valid/images"), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "valid/labels"), exist_ok=True)

#이것도 dataset_dir을 붙여서 언제 처리된 데이터셋인지 명시하기 위해 년도및 날짜/train/images 이렇게 yaml파일에서 지정
def generate_custom_yaml(dataset_dir, output_path, class_names):
    content = f"""
train: {os.path.join(dataset_dir, 'train/images')}
val: {os.path.join(dataset_dir, 'valid/images')}

nc: {len(class_names)}
names: {class_names}
"""
    with open(output_path, "w") as f:
        f.write(content.strip())


def train_yolov10(yaml_path, run_name):
    cmd = [
        "python3", f"{YOLO_DIR}/train.py",
        "--data", yaml_path,
        "--img", "640",
        "--batch", "8",
        "--epochs", "5",
        "--name", run_name,
        "--weights", "yolov10s.pt"
    ]
    log(f"Training started: {run_name}")
    try:
        subprocess.run(cmd, check=True)
        log(f"Training completed: {run_name}")
    except subprocess.CalledProcessError as e:
        log(f"Training failed: {run_name}: {e}")


def process_video(video_path):
    basename = os.path.splitext(os.path.basename(video_path))[0]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    dataset_name = f"train_{basename}_{timestamp}
    dataset_dir = os.path.join(BASE_DATASET_DIR, dataset_name)
    # 결과 /workspace/datasets/영상처리한시각/dataset/ 

    create_dataset_structure(dataset_dir)
    # 결과 /workspace/datasets/영상처리한시각/dataset/train/images 등으로 yaml에 기록될 위한 폴더 형성


    extract_frames(video_path, os.path.join(dataset_dir, "train/images"))


    #3번째 인자 클래스 이름. 2번째 인자 /workspace/datasets/영상처리한시각/dataset/custom.yaml
    generate_custom_yaml(dataset_dir, os.path.join(dataset_dir, "custom.yaml"), ["target"])
    train_yolov10(os.path.join(dataset_dir, "custom.yaml"), dataset_name)


def monitor_and_train():
    log("Monitoring started...")
    os.makedirs(WATCH_DIR, exist_ok=True)
    #"/workspace/datasets 의 디렉토리가 없을경우 생성하려 했으나 create_dataset_structure 에서 영상처리한 시간을 고려하여 디렉토리 생성
    # ex) /workspace/datasets/영상처리한시각/dataset/ 
    #os.makedirs(BASE_DATASET_DIR, exist_ok=True) 
    #한번 영상처리 된것들 지우지않고 백업용 디렉토리
    os.makedirs(PROCESSED_DIR, exist_ok=True)
    os.makedirs(YOLO_DIR, exist_ok=True) #만약을 위한 것, 사실 미리 여기에 train.py를 만들어 놔야한다 
    
    while True:
        videos = glob.glob(os.path.join(WATCH_DIR, "*.mp4"))
        #videos에는 전체경로가 반환 된다.
        for video in videos:
            process_video(video)
            backup_path = os.path.join(PROCESSED_DIR, os.path.basename(video))
            shutil.move(video, backup_path)
        time.sleep(30)


if __name__ == "__main__":
    monitor_and_train()



/workspace/
├── container_train_videos/   # 원본 .mp4 감시 폴더
├── processed_videos/         # 처리된 영상 백업
├── datasets/
│   └── train_vid001_20250514_133000/
│       ├── train/images/
│       ├── train/labels/
│       ├── valid/images/
│       ├── valid/labels/
│       └── custom.yaml
├── yolov10/
│   └── train.py  # 없으면 자동 생성됨
├── logs/
│   └── training_log.txt

데이터셋 구조
my_data
├── train
│   ├── images
│   └── labels
├── valid
│   ├── images
│   └── labels
└── test (optional)
    ├── images
    └── labels
profile
지금까지 해온 여러 활동들을 간략하게라도 정리해보고자 합니다.

0개의 댓글