Pointcept:Examples:Weld:LVS

import os
import glob
import numpy as np
from torch.utils.data import Dataset
from plyfile import PlyData
from .builder import DATASETS
from .transform import Compose

def read_ply(path):
    plydata = PlyData.read(path)
    x = np.array(plydata['vertex']['x'])
    y = np.array(plydata['vertex']['y'])
    z = np.array(plydata['vertex']['z'])
    coords = np.stack([x, y, z], axis=1)
    return coords

@DATASETS.register_module()
class WeldDataset(Dataset):
    def __init__(self,
                 split='train',
                 data_root='data/weld',
                 transform=None,
                 test_mode=False,
                 loop=1):
        super(WeldDataset, self).__init__()
        self.split = split
        self.data_root = data_root
        self.transform = Compose(transform)
        self.loop = loop
        self.test_mode = test_mode

        # 클래스 매핑 정의
        self.class_map = {
            "face1": 0,
            "face2": 1,
            "bead": 2,
        }

        self.data_list = self.get_data_list()

    def get_data_list(self):
        split_dir = os.path.join(self.data_root, self.split)
        if not os.path.exists(split_dir):
            return []
        return sorted(glob.glob(os.path.join(split_dir, "*")))

    def __getitem__(self, idx):
        scene_path = self.data_list[idx % len(self.data_list)]

        all_coords = []
        all_labels = []

        for file_name, label in self.class_map.items():
            file_path = os.path.join(scene_path, f"{file_name}.ply")
            if os.path.exists(file_path):
                coords = read_ply(file_path)
                all_coords.append(coords)
                all_labels.append(np.full(len(coords), label, dtype=np.int64))

        if len(all_coords) == 0:
            # 빈 데이터 처리 (에러 방지)
            coords = np.zeros((0, 3), dtype=np.float32)
            labels = np.zeros((0,), dtype=np.int64)
        else:
            coords = np.concatenate(all_coords, axis=0).astype(np.float32)
            labels = np.concatenate(all_labels, axis=0).astype(np.int64)

        # Pointcept standard keys: coord, segment (instead of label)
        data_dict = dict(coord=coords, segment=labels, name=os.path.basename(scene_path))

        if self.transform is not None:
            data_dict = self.transform(data_dict)

        return data_dict

    def __len__(self):
        return len(self.data_list) * self.loop

# pointcept/datasets/__init__.py 에 추가
from .weld import WeldDataset

dataset_type = "WeldDataset"
data_root = "data/weld"

data = dict(
    train=dict(
        type=dataset_type,
        split="train",
        data_root=data_root,
        transform=[
            dict(type="CenterShift", apply_z=True),
            dict(type="RandomRotate", angle=[-1, 1], axis="z", center=[0, 0, 0], p=0.5),
            dict(type="RandomScale", scale=[0.9, 1.1]),
            dict(type="DefaultPredictDoNothing"), # placeholder
            dict(type="Collect", keys=("coord", "label"), feat_keys=("coord")),
            dict(type="ToTensor"),
        ],
        loop=10, # 데이터가 적을 경우 반복 학습
    ),
    val=dict(dataset_type,
        split="val",
        data_root=data_root,
        transform=[
            dict(type="CenterShift", apply_z=True),
            dict(type="Collect", keys=("coord", "label"), feat_keys=("coord")),
            dict(type="ToTensor"),
        ],
    ),
)

_base_ = ["../_base_/default_runtime.py"]

# misc custom setting
batch_size = 1  # bs: total bs in all gpus
num_worker = 4
mix_prob = 0.0
empty_cache = False
enable_amp = True

# model settings
model = dict(
    type="DefaultSegmentorV2",
    num_classes=3,
    backbone_out_channels=64,
    backbone=dict(
        type="PT-v3m1",
        in_channels=3,
        order=("z", "z-trans", "hilbert", "hilbert-trans"),
        stride=(2, 2, 2, 2),
        enc_depths=(2, 2, 2, 6, 2),
        enc_channels=(32, 64, 128, 256, 512),
        enc_num_head=(2, 4, 8, 16, 32),
        enc_patch_size=(1024, 1024, 1024, 1024, 1024),
        dec_depths=(2, 2, 2, 2),
        dec_channels=(64, 64, 128, 256),
        dec_num_head=(4, 4, 8, 16),
        dec_patch_size=(1024, 1024, 1024, 1024),
        mlp_ratio=4,
        qkv_bias=True,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
        drop_path=0.3,
        shuffle_orders=True,
        pre_norm=True,
        enable_rpe=False,
        enable_flash=False,
        upcast_attention=True,
        upcast_softmax=True,
        enc_mode=False,
        pdnorm_bn=False,
        pdnorm_ln=False,
        pdnorm_decouple=True,
        pdnorm_adaptive=False,
        pdnorm_affine=True,
        pdnorm_conditions=("ScanNet", "S3DIS", "Structured3D"),
    ),
    criteria=[
        dict(type="CrossEntropyLoss", loss_weight=1.0, ignore_index=-1),
    ],
)

# scheduler settings
epoch = 100
optimizer = dict(type="AdamW", lr=0.006, weight_decay=0.05)
scheduler = dict(
    type="OneCycleLR",
    max_lr=0.006,
    pct_start=0.05,
    anneal_strategy="cos",
    div_factor=10.0,
    final_div_factor=1000.0,
)

# dataset settings
dataset_type = "WeldDataset"
data_root = "data/weld"

data = dict(
    num_classes=3,
    ignore_index=-1,
    names=["face1", "face2", "bead"],
    train=dict(
        type=dataset_type,
        split="train",
        data_root=data_root,
        transform=[
            dict(type="CenterShift", apply_z=True),
            dict(type="RandomRotate", angle=[-1, 1], axis="z", center=[0, 0, 0], p=0.5),
            dict(type="RandomScale", scale=[0.9, 1.1]),
            dict(
                type="GridSample",
                grid_size=0.02,
                hash_type="fnv",
                mode="train",
                return_grid_coord=True,
            ),
            dict(type="SphereCrop", point_max=102400, mode="random"),
            dict(type="ToTensor"),
            dict(
                type="Collect",
                keys=("coord", "grid_coord", "segment"),
                feat_keys=["coord"],
            ),
        ],
        loop=10,
    ),
    val=dict(
        type=dataset_type,
        split="train",  # Only have one scene for now
        data_root=data_root,
        transform=[
            dict(type="CenterShift", apply_z=True),
            dict(type="Copy", keys_dict={"segment": "origin_segment"}),
            dict(
                type="GridSample",
                grid_size=0.02,
                hash_type="fnv",
                mode="train",
                return_grid_coord=True,
                return_inverse=True,
            ),
            dict(type="ToTensor"),
            dict(
                type="Collect",
                keys=("coord", "grid_coord", "segment", "origin_segment", "inverse"),
                feat_keys=["coord"],
            ),
        ],
    ),
)

upcast_attention=False,  # fp32 upcast 해제
upcast_softmax=False,

enc_patch_size=(512, 512, 512, 512, 512),  # 1024 → 512
dec_patch_size=(512, 512, 512, 512),

# SphereCrop point_max=102400 → 51200

# grid_size 0.02 → 0.025

# 메모리 주범: attention upcast 해제 + patch 축소
upcast_attention=False,   # True → False (line 38)
upcast_softmax=False,     # True → False (line 39)
enc_patch_size=(256, 256, 256, 256, 256),   # 1024 → 256
dec_patch_size=(256, 256, 256, 256),

...

# 그리고 transform에서:
dict(type="GridSample", grid_size=0.03, ...),   # 0.02 → 0.03
dict(type="SphereCrop", point_max=40960, ...),  # 102400 → 40960

enc_depths=(1, 1, 1, 3, 1),
enc_channels=(16, 32, 64, 128, 256),
enc_num_head=(1, 2, 4, 8, 16),
dec_depths=(1, 1, 1, 1),
dec_channels=(32, 32, 64, 128),
dec_num_head=(2, 2, 4, 8),
backbone_out_channels=32,