safety-equipment-detection/
├── data/
│   ├── images/
│   │   ├── train/         # 학습 이미지
│   │   └── val/           # 검증 이미지
│   ├── labels/
│   │   ├── train/         # 학습 라벨 (YOLO TXT)
│   │   └── val/           # 검증 라벨
│   └── dataset.yaml       # YOLO 데이터셋 설정
├── labelstudio/
│   └── export/            # Label Studio 내보내기
├── runs/                  # 학습 결과
├── scripts/
│   ├── convert_labels.py  # 라벨 변환 스크립트
│   └── analyze_data.py    # 데이터 분석
└── inference.py

# 클래스 정의
CLASSES = {
    0: 'helmet',        # 안전모
    1: 'vest',          # 안전 조끼
    2: 'no-helmet',     # 안전모 미착용
    3: 'no-vest',       # 조끼 미착용
}

# dataset.yaml 작성
dataset_yaml = """
path: ./data
train: images/train
val: images/val

nc: 4
names:
  0: helmet
  1: vest
  2: no-helmet
  3: no-vest
"""

with open('data/dataset.yaml', 'w') as f:
    f.write(dataset_yaml)

# Label Studio 설정 (Docker 실행)
# docker run -it -p 8080:8080 heartexlabs/label-studio:latest

# 라벨링 인터페이스 XML 템플릿
labeling_config = """
<View>
  <Image name="image" value="$image"/>
  <RectangleLabels name="label" toName="image">
    <Label value="helmet" background="green"/>
    <Label value="vest" background="blue"/>
    <Label value="no-helmet" background="red"/>
    <Label value="no-vest" background="orange"/>
  </RectangleLabels>
</View>
"""

# Label Studio COCO Export → YOLO 변환
import json
import os

def coco_to_yolo(coco_json_path, output_dir, image_width, image_height):
    """COCO 포맷 라벨을 YOLO 포맷으로 변환합니다."""
    with open(coco_json_path) as f:
        coco = json.load(f)

    # 카테고리 매핑 (COCO ID → YOLO 클래스 인덱스)
    cat_map = {}
    for cat in coco['categories']:
        cat_map[cat['id']] = cat['id'] - 1  # 0-indexed

    # 이미지별 어노테이션 그룹화
    img_annotations = {}
    for ann in coco['annotations']:
        img_id = ann['image_id']
        if img_id not in img_annotations:
            img_annotations[img_id] = []
        img_annotations[img_id].append(ann)

    os.makedirs(output_dir, exist_ok=True)

    for img_info in coco['images']:
        img_id = img_info['id']
        img_w = img_info.get('width', image_width)
        img_h = img_info.get('height', image_height)
        filename = os.path.splitext(img_info['file_name'])[0] + '.txt'

        lines = []
        for ann in img_annotations.get(img_id, []):
            cls_id = cat_map[ann['category_id']]
            x, y, w, h = ann['bbox']  # COCO: (x_min, y_min, w, h)

            # YOLO: (cx, cy, w, h) 정규화
            cx = (x + w / 2) / img_w
            cy = (y + h / 2) / img_h
            nw = w / img_w
            nh = h / img_h
            lines.append(f"{cls_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}")

        with open(os.path.join(output_dir, filename), 'w') as f:
            f.write('\n'.join(lines))

    print(f"변환 완료: {len(coco['images'])}개 이미지")

# 변환 실행
coco_to_yolo(
    'labelstudio/export/result.json',
    'data/labels/train',
    image_width=1920,
    image_height=1080,
)

import os
import numpy as np
import matplotlib.pyplot as plt
from collections import Counter

def analyze_dataset(label_dir, class_names):
    """데이터셋의 클래스 분포와 바운딩 박스 통계를 분석합니다."""
    class_counts = Counter()
    bbox_sizes = []
    objects_per_image = []

    for label_file in os.listdir(label_dir):
        if not label_file.endswith('.txt'):
            continue

        with open(os.path.join(label_dir, label_file)) as f:
            lines = f.readlines()

        objects_per_image.append(len(lines))

        for line in lines:
            parts = line.strip().split()
            cls_id = int(parts[0])
            w, h = float(parts[3]), float(parts[4])

            class_counts[cls_id] += 1
            bbox_sizes.append(w * h)  # 정규화된 면적

    # 클래스 분포 시각화
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))

    # 클래스별 개수
    classes = [class_names[i] for i in sorted(class_counts.keys())]
    counts = [class_counts[i] for i in sorted(class_counts.keys())]
    axes[0].bar(classes, counts, color=['green', 'blue', 'red', 'orange'])
    axes[0].set_title('클래스별 객체 수')
    axes[0].set_ylabel('개수')

    # 이미지당 객체 수
    axes[1].hist(objects_per_image, bins=20, edgecolor='black')
    axes[1].set_title('이미지당 객체 수 분포')
    axes[1].set_xlabel('객체 수')

    # 바운딩 박스 크기 분포
    axes[2].hist(bbox_sizes, bins=50, edgecolor='black')
    axes[2].set_title('바운딩 박스 크기 분포 (정규화 면적)')
    axes[2].set_xlabel('면적')

    plt.tight_layout()
    plt.savefig('outputs/data_analysis.png', dpi=150)
    print(f"총 이미지: {len(objects_per_image)}장")
    print(f"총 객체: {sum(counts)}개")
    for cls, cnt in zip(classes, counts):
        print(f"  {cls}: {cnt}개 ({cnt/sum(counts)*100:.1f}%)")

analyze_dataset('data/labels/train', CLASSES)

from ultralytics import YOLO

# 모델 선택
model = YOLO('yolo11m.pt')  # Medium 모델 (균형 잡힌 성능)

# 학습 실행
results = model.train(
    data='data/dataset.yaml',
    epochs=100,
    imgsz=640,
    batch=16,
    patience=15,           # Early Stopping
    optimizer='AdamW',
    lr0=0.001,
    lrf=0.01,              # 최종 학습률 비율
    warmup_epochs=3,
    mosaic=1.0,            # Mosaic 증강
    mixup=0.1,             # MixUp 증강
    degrees=10,            # 회전 증강
    flipud=0.5,            # 상하 반전 (안전장비는 방향 중요)
    project='runs/detect',
    name='safety-equipment',
    save=True,
    plots=True,
)

# CLI로 학습 (동일한 설정)
# yolo detect train data=data/dataset.yaml model=yolo11m.pt \
#     epochs=100 imgsz=640 batch=16 patience=15 \
#     project=runs/detect name=safety-equipment

# 학습 결과 확인
import os

result_dir = 'runs/detect/safety-equipment'

# 결과 디렉터리 구조
# ├── weights/
# │   ├── best.pt          # 최고 성능 모델
# │   └── last.pt          # 마지막 에포크 모델
# ├── results.csv          # 에포크별 메트릭
# ├── confusion_matrix.png # 혼동 행렬
# ├── F1_curve.png         # F1-Confidence 곡선
# ├── P_curve.png          # Precision-Confidence 곡선
# ├── R_curve.png          # Recall-Confidence 곡선
# └── PR_curve.png         # Precision-Recall 곡선

# 결과 로드
import pandas as pd
results_df = pd.read_csv(os.path.join(result_dir, 'results.csv'))
results_df.columns = results_df.columns.str.strip()

# 주요 메트릭 출력
best_epoch = results_df['metrics/mAP50(B)'].idxmax()
print(f"최고 성능 에포크: {best_epoch + 1}")
print(f"mAP@50: {results_df.loc[best_epoch, 'metrics/mAP50(B)']:.4f}")
print(f"mAP@50:95: {results_df.loc[best_epoch, 'metrics/mAP50-95(B)']:.4f}")
print(f"Precision: {results_df.loc[best_epoch, 'metrics/precision(B)']:.4f}")
print(f"Recall: {results_df.loc[best_epoch, 'metrics/recall(B)']:.4f}")

from ultralytics import YOLO
import supervision as sv
import cv2

# 최고 모델 로드
model = YOLO('runs/detect/safety-equipment/weights/best.pt')

# 검증 세트 평가
metrics = model.val(data='data/dataset.yaml')
print(f"mAP@50: {metrics.box.map50:.4f}")
print(f"mAP@50:95: {metrics.box.map:.4f}")

# 클래스별 AP
for i, cls_name in CLASSES.items():
    print(f"  {cls_name}: AP@50 = {metrics.box.ap50[i]:.4f}")

# 이미지 추론 + supervision 시각화
def detect_and_visualize(image_path, model, confidence=0.5):
    """이미지에서 안전장비를 탐지하고 시각화합니다."""
    image = cv2.imread(image_path)
    results = model(image, conf=confidence)[0]
    detections = sv.Detections.from_ultralytics(results)

    # 어노테이터 설정
    box_annotator = sv.BoxAnnotator(thickness=2)
    label_annotator = sv.LabelAnnotator(text_scale=0.5)

    labels = [
        f"{results.names[int(cls)]} {conf:.2f}"
        for cls, conf in zip(detections.class_id, detections.confidence)
    ]

    # 시각화
    annotated = box_annotator.annotate(scene=image.copy(), detections=detections)
    annotated = label_annotator.annotate(scene=annotated, detections=detections, labels=labels)

    return annotated, detections

# 단일 이미지 추론
annotated, detections = detect_and_visualize('data/images/val/test.jpg', model)
cv2.imwrite('outputs/detection_result.jpg', annotated)

# 안전 위반 판정
def check_safety_violations(detections, class_names):
    """안전장비 미착용 위반을 판정합니다."""
    violations = []
    for cls_id, conf in zip(detections.class_id, detections.confidence):
        name = class_names[int(cls_id)]
        if name.startswith('no-'):
            violations.append(f"{name} (신뢰도: {conf:.2f})")

    if violations:
        print(f"⚠ 안전 위반 {len(violations)}건 탐지:")
        for v in violations:
            print(f"  - {v}")
    else:
        print("✓ 안전장비 정상 착용")

    return violations

violations = check_safety_violations(detections, CLASSES)

import cv2
import supervision as sv
from ultralytics import YOLO

model = YOLO('runs/detect/safety-equipment/weights/best.pt')

box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_scale=0.5)

def process_frame(frame):
    """프레임에서 안전장비를 탐지합니다."""
    results = model(frame, conf=0.5, verbose=False)[0]
    detections = sv.Detections.from_ultralytics(results)

    labels = [
        f"{results.names[int(cls)]} {conf:.2f}"
        for cls, conf in zip(detections.class_id, detections.confidence)
    ]

    annotated = box_annotator.annotate(scene=frame.copy(), detections=detections)
    annotated = label_annotator.annotate(scene=annotated, detections=detections, labels=labels)
    return annotated

# 영상 파일 처리
sv.process_video(
    source_path='input_video.mp4',
    target_path='output_video.mp4',
    callback=process_frame,
)

# ONNX 내보내기
model = YOLO('runs/detect/safety-equipment/weights/best.pt')
model.export(format='onnx', imgsz=640, simplify=True)

# TensorRT 내보내기 (NVIDIA GPU 환경)
model.export(format='engine', imgsz=640, half=True)

# TFLite 내보내기 (모바일/엣지)
model.export(format='tflite', imgsz=320)