Contents Creator

MMdetection을 사용하기 위한 초기 환경설정

1. 가상환경 구축

conda create -n mmdetection python=3.11 # ( 3.7 이상 버전이면 ok)

2. Pytorch 설치

conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

pip install mmcv==2.0.0 -f https://download.openmmlab.com/mmcv/dist/cu117/torch2.0/index.html

or 

git clone https://github.com/open-mmlab/mmdetection.git

pip install -r requirements/build.txt

pip install -v -e .

pip mmdet # 설치가 잘 됐는지 확인사살

keypoint가 라벨링된 json 파일(train)을 이용하여 이미지에 적용하여 학습하여

나머지 이미지(test)로 keypoint 찾기

# keypointrcnn_resnet50_fpn 모델 이용한 Keypoint 실습

PyTorch 제공하는 Object detection reference training scripts 다운로드

 - 다운로드 사이트  https://github.com/pytorch/vision/tree/main/references/detection 

Customdataset.py

import torch
import json
import cv2
import numpy as np
import os
from torch.utils.data import Dataset
from torchvision.transforms import functional as F


class KeypointDataset(Dataset):
    def __init__(self, root, transform=None, demo=False):
        self.demo = demo
        # Visualize를 통해 시각화로 확인하고자 할때 비교를 위한 변수
        self.root = root
        # 현재 데이터셋은 root 인자에 dataset 하위 폴더의 train, test 폴더를 택일하도록 되어 있음
        self.imgs_files = sorted(os.listdir(os.path.join(self.root, "images")))
        # 이미지 파일 리스트. train 또는 test 폴더 하위에 있는 images 폴더를 지정하고, 해당 폴더의 내용물을 받아옴
        # 이미지를 이름 정렬순으로 불러오도록 sorted를 붙임
        self.annotations_files = sorted(os.listdir(os.path.join(self.root, "annotations")))
        # 라벨링 JSON 파일 리스트. 상기 images 폴더를 받아온 것과 동일
        self.transform = transform

    def __getitem__(self, idx):
        img_path = os.path.join(self.root, "images", self.imgs_files[idx])
        # 이번에 호출되는 idx번째 이미지 파일의 절대경로
        annotations_path = os.path.join(self.root, "annotations", self.annotations_files[idx])
        # 이번에 호출되는 idx번째 이미지 파일의 라벨 JSON 파일 경로

        img_original = cv2.imread(img_path)
        img_original = cv2.cvtColor(img_original, cv2.COLOR_BGR2RGB)
        # 이미지를 읽은 후, BGR 순서를 RGB 형태로 바꿈

        with open(annotations_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            # 라벨 JSON 파일을 JSON 모듈로 받아옴
            bboxes_original = data["bboxes"]
            # JSON 하위 "bboxes" 키로 bbox 정보들이 담겨있음
            keypoints_original = data["keypoints"]
            # "keypoints" 키로 키포인트 정보가 담겨있음
            bboxes_labels_original = ['Glue tube' for _ in bboxes_original]
            # 현재 데이터셋은 모든 객체가 접착제 튜브이므로,
            # 모든 bbox에 대해 일관적으로 'Glue tube'라는 라벨을 붙여줌

        if self.transform:  # if self.transform is not None:
            # "keypoints": [
            #   [[1019, 487, 1], [1432, 404, 1]], [[861, 534, 1], [392, 666, 1]]
            # ]
            keypoints_original_flattened = [el[0:2] for kp in keypoints_original for el in kp]
            # kp : [[1019, 487, 1]]
            # el : [1019, 487, 1]
            # el[0:2] : [1019, 487] (평면 이미지에서 3번째 축 요소는 필요없기 때문에 제거)
            # keypoints_original_flattened = [[1019, 487], [1432, 404], [861, 534], [392, 666]]
            # albumentation transform은 평면 이미지에 적용되므로 2차원 좌표만이 필요함

            # albumentation 적용
            transformed = self.transform(image=img_original, bboxes=bboxes_original,
                                         bboxes_labels=bboxes_labels_original,
                                         keypoints=keypoints_original_flattened)

            img = transformed["image"]  # albumentation transform이 적용된 image
            bboxes = transformed["bboxes"]

            keypoints_transformed_unflattened = np.reshape(np.array(transformed["keypoints"]), (-1, 2, 2)).tolist()
            # transformed["keypoints"] : [1019, 487, 1432, 404, 861, 534, 392, 666]
            # keypoints_transformed_unflattened : [[[1019, 487], [1432, 404]], [[861, 534], [392, 666]]]

            keypoints = []
            for o_idx, obj in enumerate(keypoints_transformed_unflattened):
                obj_keypoints = []
                # o_idx : 현재 순회중인 요소의 순번 (index)
                # obj : 현재 순회중인 요소, ex) [[1019, 487], [1432, 404]]
                for k_idx, kp in enumerate(obj):
                    # k_idx : 현재 순회중인 하위 요소의 순번 (index)
                    # kp : 현재 순회중인 요소, ex) [1019, 487]
                    obj_keypoints.append(kp + [keypoints_original[o_idx][k_idx][2]])
                    # torch.Tensor에서 벡터곱을 하는 과정에서 필요할 3번째 축 요소를 덧붙임 ex) [1019, 487, 1]
                keypoints.append(obj_keypoints)
                # Tensor 형태로 사용할 keypoints 리스트에 3번째 축 요소를 덧붙인 키포인트 좌표를 담음

        else:
            img, bboxes, keypoints = img_original, bboxes_original, keypoints_original
            # transform이 없는 경우에는 변수 이름만 바꿔줌

        # transform을 통과한 값들을 모두 tensor로 변경
        bboxes = torch.as_tensor(bboxes, dtype=torch.float32)
        # as_tensor 메서드가 list를 tensor로 변환할 때 속도 이점이 있음
        target = {}
        # keypoint 모델에 사용하기 위한 label이 dictionary 형태로 필요하므로, dict 형태로 꾸림
        target["boxes"] = bboxes
        target["labels"] = torch.as_tensor([1 for _ in bboxes], dtype=torch.int64)
        # 모든 객체는 동일하게 접착제 튜브이므로, 동일한 라벨 번호 삽입
        target["image_id"] = torch.tensor([idx])
        # image_id는 고유번호를 지칭하는 경우도 있는데, 그러한 경우에는 JSON 파일에 기입이 되어있어야 함
        # 이번 데이터셋은 JSON상에 기입되어있지 않으므로, 현재 파일의 순번을 넣어줌
        target["area"] = (bboxes[:, 3] - bboxes[:, 1]) * (bboxes[:, 2] - bboxes[:, 0])
        # 해당하는 bbox의 넓이.
        # bboxes[:, 3] - bboxes[:, 1] : bboxes 내부 요소들의 (y2 - y1), 즉 세로 길이
        # bboxes[:, 2] - bboxes[:, 0] : bboxes 내부 요소들의 (x2 - x1), 즉 가로 길이
        target["iscrowd"] = torch.zeros(len(bboxes), dtype=torch.int64)
        # 이미지상에 키포인트 또는 bbox가 가려져있는지를 묻는 요소
        target["keypoints"] = torch.as_tensor(keypoints, dtype=torch.float32)

        img = F.to_tensor(img)
        # image의 텐서 변환

        bboxes_original = torch.as_tensor(bboxes_original, dtype=torch.float32)
        target_original = {}
        target_original["boxes"] = bboxes_original
        target_original["labels"] = torch.as_tensor([1 for _ in bboxes_original],
                                                    dtype=torch.int64)  # all objects are glue tubes
        target_original["image_id"] = torch.tensor([idx])
        target_original["area"] = (bboxes_original[:, 3] - bboxes_original[:, 1]) * (
                bboxes_original[:, 2] - bboxes_original[:, 0])
        target_original["iscrowd"] = torch.zeros(len(bboxes_original), dtype=torch.int64)
        target_original["keypoints"] = torch.as_tensor(keypoints_original, dtype=torch.float32)
        img_original = F.to_tensor(img_original)
        # demo=True일 경우, 원본 이미지와 변환된 이미지를 비교하기 위해 원본 이미지를 반환하기 위한 블록

        if self.demo:
            return img, target, img_original, target_original
        else:
            return img, target

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


if __name__ == "__main__":
    root_path = "./keypoint_dataset"
    train_dataset = KeypointDataset(f"{root_path}/train")
    for item in train_dataset:
        print(item)

visualize.py

import cv2
import albumentations as A
from Customdataset import KeypointDataset
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from utils import collate_fn


train_transform = A.Compose([
                    A.Sequential([
                        A.RandomRotate90(p=1), # 랜덤 90도 회전
                        A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, brightness_by_max=True,
                                                   always_apply=False, p=1) # 랜덤 밝기 및 대비 조정
                    ], p=1)
                ], keypoint_params=A.KeypointParams(format='xy'),
                   bbox_params=A.BboxParams(format="pascal_voc", label_fields=['bboxes_labels']) # 키포인트의 형태를 x-y 순으로 지정
                )

root_path = "./keypoint_dataset/"
dataset = KeypointDataset(f"{root_path}/train/", transform=train_transform, demo=True)
# demo=True 인자를 먹으면 Dataset이 변환된 이미지에 더해 transform 이전 이미지까지 반환하도록 지정됨
data_loader = DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=collate_fn)

iterator = iter(data_loader)
# for문으로 돌릴 수 있는 복합 자료형들은 iterable (반복 가능) 속성을 갖고 있음
# iter()로 그러한 자료형을 감싸면 iterator (반복자) 가 되고,
# next(iterator)를 호출하면서 for문을 돌리듯이 내부 값들을 순회할 수 있게 됨
batch = next(iterator)
# iterator에 대해 next로 감싸서 호출을 하게 되면, 
# for item in iterator의 예시에서 item에 해당하는 단일 항목을 반환함
# 아래 4줄에 해당하는 코드와 같은 의미
# batch_ = None
# for item in data_loader:
#     batch_ = item
#     break

keypoints_classes_ids2names = {0: "Head", 1: "Tail"}
# bbox 클래스는 모두 접착제 튜브 (Glue tube) 로 동일하지만, keypoint 클래스는 위의 dict를 따름


def visualize(image, bboxes, keypoints, image_original=None, bboxes_original=None, keypoints_original=None):
    # pyplot을 통해서 bbox와 키포인트가 포함된
    # 원본 이미지와 변환된 이미지를 차트에 띄워서 대조할 수 있는 편의함수
    fontsize = 18
    # cv2.putText에 사용될 글씨 크기 변수

    for bbox in bboxes:
        # bbox = xyxy
        start_point = (bbox[0], bbox[1])
        # 사각형의 좌측 상단
        end_point = (bbox[2], bbox[3])
        # 사각형의 우측 하단
        image = cv2.rectangle(image.copy(), start_point, end_point, (0, 255, 0), 2)
        # 이미지에 bbox 좌표에 해당하는 사각형을 그림

    for kpts in keypoints:
        # keypoints : JSON 파일에 있는 keypoints 키의 값 (즉, keypoints 최상위 리스트)
        # kpts : keypoints 내부의 각 리스트 (즉, 각 bbox의 키포인트 리스트)
        for idx, kp in enumerate(kpts):
            # kp : kpts 내부의 각 리스트 (즉, 키포인트 리스트 내부의 xy좌표쌍, 키포인트 점)
            image = cv2.circle(image.copy(), tuple(kp), 5, (255, 0, 0), 10)
            # 현재 키포인트에 점을 찍음
            image = cv2.putText(image.copy(), f" {keypoints_classes_ids2names[idx]}", tuple(kp),
                                cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 3, cv2.LINE_AA)
            # 현재 키포인트가 Head인지 Tail인지 위에 선언한 dict에 해당하는 문자를 집어넣음

    # 변환된 이미지만을 확인할 경우, 원본 이미지가 없을 것이므로 그대로 이미지 처리 끝냄
    if image_original is None and keypoints_original is None:
        plt.figure(figsize=(40, 40))
        # 이미지를 그릴 차트 선언
        plt.imshow(image)
        # 위에서 bbox와 키포인트를 그린 이미지를 출력

    else:
        for bbox in bboxes_original:
        # bbox = xyxy
            start_point = (bbox[0], bbox[1])
            # 사각형의 좌측 상단
            end_point = (bbox[2], bbox[3])
            # 사각형의 우측 하단
            image_original = cv2.rectangle(image_original.copy(), start_point, end_point, (0, 255, 0), 2)
            # 이미지에 bbox 좌표에 해당하는 사각형을 그림

        for kpts in keypoints_original:
            # keypoints : JSON 파일에 있는 keypoints 키의 값 (즉, keypoints 최상위 리스트)
            # kpts : keypoints 내부의 각 리스트 (즉, 각 bbox의 키포인트 리스트)
            for idx, kp in enumerate(kpts):
                # kp : kpts 내부의 각 리스트 (즉, 키포인트 리스트 내부의 xy좌표쌍, 키포인트 점)
                image_original = cv2.circle(image_original.copy(), tuple(kp), 5, (255, 0, 0), 10)
                # 현재 키포인트에 점을 찍음
                image_original = cv2.putText(image_original.copy(), f" {keypoints_classes_ids2names[idx]}", tuple(kp),
                                    cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 3, cv2.LINE_AA)
                # 현재 키포인트가 Head인지 Tail인지 위에 선언한 dict에 해당하는 문자를 집어넣음
        
        f, ax = plt.subplots(1, 2, figsize=(40, 20))
        # 두 장의 이미지를 1행 2열, 즉 가로로 길게 보여주는 subplots 생성

        ax[0].imshow(image_original)
        # 첫번째 subplot에는 원본 이미지를 출력
        ax[0].set_title("Original Image", fontsize=fontsize)
        # 이미지 제목

        ax[1].imshow(image)
        # 두번째 subplot에는 변환이 완료된 이미지를 출력
        ax[1].set_title("Transformed Image", fontsize=fontsize)

        plt.show()


if __name__=="__main__":

    visualize_image_show = True
    visualize_targets_show = True

    image = (batch[0][0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    # CustomDataset에서 Tensor로 변환했기 때문에 다시 plt에 사용할 수 있도록 numpy 행렬로 변경
    # img, target, img_original, target_original = batch이므로, batch[0]는 img를 지칭
    # batch[0][0]에 실제 이미지 행렬에 해당하는 텐서가 있을것 (batch[0][1]에는 dtype 등의 다른 정보가 있음)
    bboxes = batch[1][0]['boxes'].detach().cpu().numpy().astype(np.int32).tolist()
    # target['boxes']에 bbox 정보가 저장되어있으므로, 해당 키로 접근하여 bbox 정보를 획득

    keypoints = []
    for kpts in batch[1][0]['keypoints'].detach().cpu().numpy().astype(np.int32).tolist():
        keypoints.append([kp[:2] for kp in kpts])
        # 이미지 평면상 점들이 필요하므로, 3번째 요소로 들어있을 1을 제거

    image_original = (batch[2][0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    # batch[2] : image_original
    bboxes_original = batch[3][0]['boxes'].detach().cpu().numpy().astype(np.int32).tolist()
    # batch[3] : target

    keypoints_original = []
    for kpts in batch[3][0]['keypoints'].detach().cpu().numpy().astype(np.int32).tolist():
        keypoints_original.append([kp[:2] for kp in kpts])

    if visualize_image_show:
        visualize(image, bboxes, keypoints, image_original, bboxes_original, keypoints_original)
    if visualize_targets_show and visualize_image_show == False:
        print("Original targets: \n", batch[3], "\n\n")
        # original targets: (줄바꿈) original targets dict 출력 (두줄 내림)
        print("Transformed targets: \n", batch[1])

main.py

import torch
import torchvision
import albumentations as A

from engine import train_one_epoch, evaluate
from utils import collate_fn

from torch.utils.data import DataLoader
from Customdataset import KeypointDataset

from torchvision.models.detection import keypointrcnn_resnet50_fpn
from torchvision.models.detection.rpn import AnchorGenerator

def get_model(num_keypoints, weights_path=None):
    # 필요한 모델에 대해 키포인트 개수를 정의하고, 기존 모델이 있는 경우 로드하는 편의함수
    anchor_generator = AnchorGenerator(sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 4.0))

    # 여러 input size에 대해 feature map으로 넘어갈 때 적절한 비율 변환율을 도와주는 객체

    model = keypointrcnn_resnet50_fpn(
                        pretrained=False, # 모델 자체는 pretrain된 것을 사용하지 않음
                        # 현재는 학습용 코드이기 때문에, pretrain 모델을 fine-tuning 하지않고 처음부터 가중치 업데이트를 하도록 설정

                        pretrained_backbone=True, # backbone만 pretrain된 것을 사용함
                        # backbone은 모델 설계상 이미 pretrain 됐을 것으로 상정했기 때문에
                        # 실제 가중치 없데이트가 주로 일어날 부분에 비해 pretrain 여부가 크게 상관있지 않음
                        num_classes=2, # 무조건 배경 클래스를 포함함
                        num_keypoints=num_keypoints,
                        rpn_anchor_generator=anchor_generator)
    if weights_path: # 기존 모델이 있는 경우
        state_dict = torch.load(weights_path)
        model.load_state_dict(state_dict)

    return model

train_transform = A.Compose([
                    A.Sequential([
                        A.RandomRotate90(p=1), # 랜덤 90도 회전
                        A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, brightness_by_max=True,
                                                    always_apply=False, p=1) # 랜덤 밝기 및 대비 조정
                    ], p=1)
                ], keypoint_params=A.KeypointParams(format='xy'),
                      bbox_params=A.BboxParams(format="pascal_voc", label_fields=['bboxes_labels']) # 키포인트의 형태를 x-y 순으로 지정
                )

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

KEYPOINTS_FOLDER_TRAIN = "./keypoint_dataset/train/"
# train dataset이 있는 경로
train_dataset = KeypointDataset(KEYPOINTS_FOLDER_TRAIN, transform=train_transform)
train_dataloader = DataLoader(train_dataset, batch_size=6, shuffle=True, collate_fn=collate_fn)

model = get_model(num_keypoints=2)
model.to(device)

optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)
# momentum : 이전에 가중치를 업데이트한 기울기를 얼마나 반영할 것인가?
# weight_decay : 가중치 감소율
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.3)
# stepLR : step_size epoch 마다 lr이 기준 lr * gamma 로 변경됨 -> 즉 5 epoch 마다 lr이 0.3배씩 줄어듬

num_epochs = 20
# 최대 학습 횟수

# 현재 engine에서 받아온 train_one_epoch 함수는 손실함수를 넣지 않아도 되도록 처리된 함수이므로,
# 손실함수의 정의 는 생략됨

for epoch in range(num_epochs):
    train_one_epoch(model, optimizer, train_dataloader, device, epoch, print_freq=1000)
    # 학습이 진행되는 함수
    lr_scheduler.step()
    # 학습 1 epoch가 끝날때마다 scheduler 역시 업데이트
    if epoch % 10 == 0:
        torch.save(model.state_dict(), f"./keypointsrcnn_weights_{epoch}.pth")
        # 10 epochs마다 가중치 저장

torch.save(model.state_dict(), "./keypointsrcnn_weights_last.pth")
# 모든 epoch가 끝나면 last.pth까지 저장