차량 파손 데이터셋

아래 차량 파손 데이터셋을 다운로드 받은 후 예제를 진행합니다.

https://www.aihub.or.kr/aihubdata/data/view.do?currMenu=115&topMenu=100&aihubDataSe=data&dataSetSn=581

 

%pwd

%cd /content/drive/MyDrive/본인의 파일 경로/resource/damage/data

 

!unzip -oq ./images/damage.zip -d ./images/
!unzip -oq ./annotations/damage.zip -d ./annotations/

 

!pip install -q ultralytics opencv-python

 

import os
import random
import shutil
import cv2
import glob
import json
import numpy as np
import matplotlib.pyplot as plt
import ultralytics
from tqdm import tqdm
from pycocotools import mask as maskUtils

 

file_root = '/content/drive/MyDrive/본인의 파일 경로'
data_root = '/content/drive/MyDrive/본인의 파일 경로ᅥᆫ/resource/damage'
train_root = f'{data_root}/train'
valid_root = f'{data_root}/valid'
test_root = f'{data_root}/test'
cls_list = ['Scratched', 'Crushed', 'Breakage', 'Separated']

 

for folder in [train_root, valid_root, test_root]:
    if not os.path.exists(folder):
        os.makedirs(folder)
    for s in ['images', 'labels']:
        s_folder = f'{folder}/{s}'
        if not os.path.exists(s_folder):
            os.makedirs(s_folder)

 

data_list = glob.glob(f'{data_root}/data/images/*.jpg')
len(data_list)

 

data_list[0]

 

def draw_polygon_by_index(idx):
    image_path = data_list[idx]
    image = cv2.imread(image_path)
    json_path = image_path.replace('images', 'annotations').replace('.jpg', '.json')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    with open(json_path, 'r') as f:
        data = json.load(f)

    for ann in data['annotations']:
        if 'segmentation' in ann:
            segmentation = ann['segmentation']
            if isinstance(segmentation, list):
                for seg in segmentation:
                    if isinstance(seg, list):
                        points = np.array(seg).reshape(-1, 2).astype(np.int32)
                        cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=2)
    
    plt.figure(figsize=(10, 8))
    plt.imshow(image)
    plt.axis('off')
    plt.title(f'Polygon Overlay: {os.path.basename(image_path)}')
    plt.show()

 

draw_polygon_by_index(10)

 

draw_polygon_by_index(15)

 

draw_polygon_by_index(20)

 

damage_classes = {
    "Scratched": 0,
    "Crushed": 1,
    "Breakage": 2,
    "Separated": 3
}

 

def convert_json_to_yolo(json_path, save_dir):
    with open(json_path, 'r') as f:
        data = json.load(f)

    image_info = data['images']
    img_width, img_height = image_info['width'], image_info['height']

    yolo_annotations = []

    for ann in data['annotations']:
        if 'segmentation' in ann and 'damage' in ann:
            segmentation = ann['segmentation'][0][0]

            if isinstance(segmentation[0], list):
                segmentation = [coord for sublist in segmentation for coord in sublist]

            normalized_coords = [segmentation[i] / img_width if i % 2 == 0 else segmentation[i] / img_height for i in range(len(segmentation))]

            class_label = damage_classes.get(ann['damage'], -1)
            yolo_annotations.append(f'{class_label} ' + ' '.join(map(str, normalized_coords)))
    
    yolo_filename = os.path.basename(json_path).replace('.json', '.txt')
    save_path = os.path.join(save_dir, yolo_filename)

    with open(save_path, 'w') as f:
        f.write('\n'.join(yolo_annotations))

 

for json_path in tqdm(glob.glob(os.path.join(f'{data_root}/data/annotations/', '*.json'))):
    convert_json_to_yolo(json_path, f'{data_root}/data/labels')
print(f'YOLO 변환 완료!!')

 

random.seed(2025)
file_list = glob.glob(f'{data_root}/data/labels/*.txt')

random.shuffle(file_list)
test_ratio = 0.1
num_file = len(file_list)

test_list = file_list[:int(num_file*test_ratio)]
valid_list = file_list[int(num_file*test_ratio):int(num_file*test_ratio)*2]
train_list = file_list[int(num_file*test_ratio)*2:]

 

for i in tqdm(test_list):
    label_name = i.split('/')[-1]
    shutil.copyfile(i, f'{test_root}/labels/{label_name}')
    img_name = i.split('/')[-1].replace('txt', 'jpg')
    img_path = f'{data_root}/data/images/{img_name}'
    shutil.copyfile(img_path, f'{test_root}/images/{img_name}')

for i in tqdm(valid_list):
    label_name = i.split('/')[-1]
    shutil.copyfile(i, f'{valid_root}/labels/{label_name}')
    img_name = i.split('/')[-1].replace('txt', 'jpg')
    img_path = f'{data_root}/data/images/{img_name}'
    shutil.copyfile(img_path, f'{valid_root}/images/{img_name}')

for i in tqdm(train_list):
    label_name = i.split('/')[-1]
    shutil.copyfile(i, f'{train_root}/labels/{label_name}')
    img_name = i.split('/')[-1].replace('txt', 'jpg')
    img_path = f'{data_root}/data/images/{img_name}'
    shutil.copyfile(img_path, f'{train_root}/images/{img_name}')

 

%pwd

%cd /content/drive/MyDrive/본인의 파일 경로/resource/damage

 

import yaml

data = dict()

data['train'] = train_root
data['val'] = valid_root
data['test'] = test_root
data['nc'] = len(cls_list)
data['names'] = cls_list

with open(f'car_damage.yaml', 'w') as f:
    yaml.dump(data, f)

 

from ultralytics import YOLO

model = YOLO('yolov8s-seg.yaml')
results = model.train(
    data = 'car_damage.yaml',
    epochs = 100,
    imgsz = 224,
    batch = 16,
    device = 0,
    name = 'damage_s'
)