onlywonhand/SDD-Seg_Structural_Damage_Dataset_for_Segmentation

import cv2
import torch
import pandas as pd
from torch.utils.data import Dataset
import numpy as np

class StructureDamageDataset(Dataset):
    def __init__(self, csv_path):
        self.df = pd.read_csv(csv_path)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        
        # 이미지 로드
        image = cv2.imread(row['image_path'])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 마스크 로드 (PNG 파일은 0~7의 손상 값을 가짐)
        mask_png = cv2.imread(row['mask_path'], cv2.IMREAD_UNCHANGED)
        
        # 14채널 통합 마스크 생성 (H, W, 14)
        H, W = image.shape[:2]
        target = np.zeros((H, W, 14), dtype=np.float32)
        
        # [Ch 0~7] 손상 정보 One-hot Encoding
        for cls_id in range(1, 8):
            target[:, :, cls_id] = (mask_png == cls_id).astype(np.float32)
            
        # [Ch 8~13] 구조물 정보 (Weakly Supervised - 전체 채우기)
        struct_id = int(row['structure_class_id']) # 0~5
        target[:, :, 8 + struct_id] = 1.0 
        
        return image, target