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까지 저장