[차량속성모델 만들기 2] 차량의 종류와 색상 멀티태스크 모델 만들기, RESNET50, Yolo Segmentation 적용

만들고자 하는 최종 목표 모델 : 차량의 종류, 색상을 뽑는 모델 만들기
목표 정확도 : 99% 이상

 

이전에 RESNET50으로 차량의 색상 모델을 만들었다.

2025.05.19 - [Project Portfolio] - [차량속성모델 만들기 1] 차량의 색상 모델 만들기, RESNET50, Yolo Segmentation 적용

[차량속성모델 만들기 1] 차량의 색상 모델 만들기, RESNET50, Yolo Segmentation 적용

 

이번에는 차량의 종류와 색상을 한번에 뽑는 멀티태스크 모델을 설계하여 최종 모델을 만들어보고자 한다.

멀티태스크 모델?
하나의 모델이 동시에 여러 가지 다른 작업을 학습하여 수행하는 모델

 

입력 데이터 (예: 차량 이미지)
        │
하나의 공통 특징 추출기 (예: RESNET50)
        │
    ┌───┴───┐
    │       │
 작업1    작업2
(색상)   (차종)

 

색상 모델을 만들 때와 마찬가지로, 

도로와 창문 등 차량과 관련 없는 요소들을 제거하기 위해 Yolo Segmentation을 적용했다.

그리고 데이터 정규화, 데이터 마스킹 기법을 포함한 이미지 전처리, 손실함수에 클래스별 가중치 적용, 얼리스타핑을 적용했다.

 

1. Dataset

먼저 차량과 창문을 딴 라벨링 데이터가 필요하다. 라벨링은 labelme라는 툴에서 하면 된다.

라벨링을 하게 되면, JSON 파일에 라벨링 결과가 저장되어있다.

(Yolo Segmentation 모델을 만드는 방식에 대해서는 다른 글에서 설명하겠다.)

2025.05.19 - [오늘의 공부일기/머신러닝, 딥러닝 공부일기] - Yolo Segmentation 커스텀모델 Training, JsonToYolo 포함

 

이미지 당 하나의 JSON 파일이 있어야 한다.

그럼 모든 준비가 끝났다.

 

2. 데이터 정규화 및 훈련, 검증 데이터 나누기

보통 표준으로 많이 사용하는 ImageNet 데이터의 RGB 평균과 표준편차 대신

현재의 데이터셋에서 평균과 표준편차를 직접 구해 적용했다.

또한 훈련과 검증 데이터를 8대 2로 나누었다.

# ImageNet 데이터셋의 RGB 채널 평균과 표준편차
transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                     std=[0.229, 0.224, 0.225]) 

# 직접 구해 적용
temp_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])
temp_dataset = VehicleDataset(train_folder, temp_transform)
temp_loader = DataLoader(temp_dataset, batch_size=16, shuffle=False)

mean = torch.zeros(3)
std = torch.zeros(3)
for imgs, _, _ in temp_loader:
    batch_size = imgs.size(0)
    for i in range(3):
        mean[i] += imgs[:, i, :, :].mean() * batch_size
        std[i] += imgs[:, i, :, :].std() * batch_size
mean.div_(len(temp_dataset))
std.div_(len(temp_dataset))
print("Total_length", len(temp_dataset))
print(f"Dataset Mean: [{mean[0]:.4f}, {mean[1]:.4f}, {mean[2]:.4f}], "
      f"STD: [{std[0]:.4f}, {std[1]:.4f}, {std[2]:.4f}]")
      
# 훈련, 검증 데이터 8대 2로 나누기
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_subset, val_subset = torch.utils.data.random_split(dataset, [train_size, val_size], generator=torch.Generator().manual_seed(42))

train_loader = DataLoader(train_subset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=16, shuffle=False)
print(f"Train 이미지 수: {len(train_subset)}")
print(f"Val 이미지 수: {len(val_subset)}")
print(f"Train 배치 수: {len(train_loader)}, Val 배치 수: {len(val_loader)}")

데이터셋 평균, 표준편차, 이미지셋 통계 출력

 

3. 이미지 전처리, 마스킹 적용

도로와 창문과 같이 정확한 차량의 색상을 뽑기에 장애가 되는 요소들을 제거하기 위해서 이미지 마스킹을 적용했다.

또한, 밝기를 조절하여 너무 어두운 영상, 너무 밝은 영상을 조절했다.

def preprocess_image_with_mask(img_path: str, json_path: str):
    image = Image.open(img_path).convert("RGB")
    img_np = np.array(image)

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

    mask = np.zeros((img_np.shape[0], img_np.shape[1]), dtype=np.uint8)

    for shape in metadata['shapes']:
        points = np.array(shape['points'], dtype=np.int32)
        if shape['label'] in type_classes:
            cv2.fillPoly(mask, [points], color=1)
        elif shape['label'] == 'window':
            cv2.fillPoly(mask, [points], color=0)

    # 마스크 적용: 창문을 제외한 차량 외곽만
    masked_img = cv2.bitwise_and(img_np, img_np, mask=mask)

    # HSV 변환 및 밝기 조절
    img_hsv = cv2.cvtColor(masked_img, cv2.COLOR_RGB2HSV)
    h, s, v = cv2.split(img_hsv)
    mean_brightness = np.mean(v[mask == 1])

    if mean_brightness < 90:
        v = np.clip(v.astype(np.float32) * 1.3, 0, 255).astype(np.uint8)
    elif mean_brightness > 160:
        v = np.clip(v.astype(np.float32) * 0.8, 0, 255).astype(np.uint8)

    v = cv2.equalizeHist(v)
    img_hsv = cv2.merge((h, s, v))
    img_rgb = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)

    img = Image.fromarray(img_rgb)
    return img

# === 3. 데이터셋 정의 (train) ===
class VehicleDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_paths = []
        for ext in valid_extensions:
            self.img_paths.extend(glob.glob(os.path.join(img_dir, f"*{ext}")))
        self.transform = transform

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

    def __getitem__(self, idx):
        path = self.img_paths[idx]

        # JSON 파일에서 color와 label 정보 읽기
        json_path = os.path.splitext(path)[0] + ".json"
        with open(json_path, 'r') as f:
            metadata = json.load(f)
        
        for shape in metadata['shapes']:
            if shape['label'] == 'window':
                continue
            else:
                color = shape['color']
                vtype = shape['label']
        color_label = color_classes.index(color)
        type_label = type_classes.index(vtype)

        # image = Image.open(path).convert("RGB")
        # image = preprocess_image(image)  # 이미지 전처리
        image = preprocess_image_with_mask(path, json_path)  # 이미지 전처리
        if self.transform:
            image = self.transform(image)

        return image, torch.tensor(color_label), torch.tensor(type_label)

마스킹, 밝기 조절 등 이미지 전처리 결과

 

4. 손실함수, 클래스별 가중치 적용

클래스별 가중치를 계산하여 적용했을 때의 장점 : 

1. 클래스 불균형 문제 해결

2. 정확성 및 일반성 성능 향상

3. 모델의 편향 감소

4. 데이터 수가 없는 클래스에 대한 안정적 처리

 

클래스 불균형 문제를 해결하고, 안정적인 학습을 위해서 클래스별 가중치를 계산하여 적용했다.

"역빈도 가중치" 방식을 사용했다.

 

• 데이터가 적은 클래스에 높은 가중치를 부여
• 데이터가 많은 클래스에 상대적으로 낮은 가중치를 부여
• 데이터가 없는 클래스는 평균 가중치로 대체

# === 6. 클래스별 weight 계산 ===
from collections import Counter

def compute_class_weights(label_list, num_classes):
    counter = Counter(label_list)
    total = sum(counter.values())
    raw_weights = []

    for i in range(num_classes):
        count = counter.get(i, 0)
        if count > 0:
            weight = total / (count * num_classes)
        else:
            weight = None
        raw_weights.append(weight)

    # 평균 가중치로 대체
    mean_weight = np.mean([w for w in raw_weights if w is not None])
    weights = [w if w is not None else mean_weight for w in raw_weights]

    return torch.tensor(weights, dtype=torch.float)

클래스별 가중치 출력

 

클래스 불균형을 고려하기 위해 색상과 종류 각각의 클래스 가중치를 계산하여 손실함수에 적용했다.

훈련을 해본 결과, 색상 정확도가 낮게 나와서 전체손실을 계산할 때, 색상에 더 큰 비중을 부여하였다.

loss = 0.7 * criterion_color(out_color, color_lbl) + 0.3 * criterion_type(out_type, type_lbl)

이렇게 색상을 좀 더 집중적으로 학습하도록 했다.

 

5. 얼리스타핑 적용

과적합을 방지하기 의해서 얼리스타핑을 적용했다.

# Early stopping
if val_losses[-1] < best_valid_loss:
    best_valid_loss = val_losses[-1]
    early_stop_counter = 0
    torch.save(model.state_dict(), os.path.join(save_dir,"vehicle_colorNtype_resnet50_best.pt"))
    print("✅ Best model saved!")
else:
    early_stop_counter += 1

 

6. 모델 설계

RESNET50을 백본으로 사용하고, 차량 이미지의 색상과 종류를 분류하는 모델을 설계했다.

RESNET50의 마지막 완전연결층을 제거하여 클래스분류기능없이 오로지 특징 추출만 수행하도록 했다.

백본에서 추출된 특징을 받아 색상과 종류 클래스를 예측했다.

# === 5. 모델 정의 ===
class VehicleClassifier(nn.Module):
    def __init__(self, num_colors, num_types):
        super(VehicleClassifier, self).__init__()
        base = resnet50(weights=ResNet50_Weights.DEFAULT)
        self.backbone = nn.Sequential(*list(base.children())[:-1])  # remove FC
        self.fc_color = nn.Linear(2048, num_colors)
        self.fc_type = nn.Linear(2048, num_types)

    def forward(self, x):
        x = self.backbone(x)          # (B, 2048, 1, 1)
        x = x.view(x.size(0), -1)     # (B, 2048)
        return self.fc_color(x), self.fc_type(x)

 

7. 훈련 결과

훈련 과정에서의 Loss, Acc, 얼리스타핑, 테스트결과 정확도 출력

모델의 테스트 결과, Confusion Matrix 표 (정규화)

Confusion Matrix란?
분류 모델의 성능을 평가하기 위한 도구로, 실제 데이터의 클래스와 모델이 예측한 클래스 간의 관계를 시각적으로 나타낸 표.

 

색상과 종류 모두 잘 뽑아낸 것을 확인할 수 있다.

 

모든 추론 결과를 json에 저장 및 업데이트하도록 하고, 이를 시각화해보면, 아래와 같다.