Penn-Fudan으로 알아보는 객체 탐지(Object Detection), 분할(Segmentation) with FasterRCNN

# !wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/engine.py
# !wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/utils.py
# !wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_utils.py
# !wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/coco_eval.py
# !wget https://raw.githubusercontent.com/pytorch/vision/main/references/detection/transforms.py

# !wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip
# !unzip PennFudanPed.zip

# Deafult
import os

# Image
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.image as img
import matplotlib.patches as patches
import cv2
from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks

# Model
import torch
import torchvision
from torchvision.io import read_image
from torchvision import tv_tensors
from torchvision.ops.boxes import masks_to_boxes
from torchvision.transforms.v2 import functional as F
from torchvision.transforms import v2 as T
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator

# Custom Function
import utils
from engine import train_one_epoch, evaluate

# 버전 확인
# Torchvision의 버전이 0.15이하이면 앞으로 진핼할 아래의 코드는 실행이 되지 않음
print('Torch Version : ',torch.__version__)
print('Torchvision Version : ',torchvision.__version__)
print('Matplotlib Version : ',matplotlib.__version__)
print('Cv2 Version : ',cv2.__version__)

폴더 구조
PennFudanPed/
    PedMasks/
        FudanPed00001_mask.png
        FudanPed00002_mask.png
        FudanPed00003_mask.png
        FudanPed00004_mask.png
        ...
    PNGImages/
        FudanPed00001.png
        FudanPed00002.png
        FudanPed00003.png
        FudanPed00004.png
        ...
    Annotation/
        FudanPed00001.txt
        FudanPed00002.txt
        FudanPed00003.txt
        FudanPed00004.txt
        ...
        

# 샘플 데이터 확인
def draw_images(annotation_file_path):
    # 어노테이션 파일 읽기
    with open(annotation_file_path, 'r') as file:
        lines = file.readlines()

    # 바운딩 박스 정보와 마스크 이미지 경로 추출
    bounding_boxes = []
    mask_image_path = None
    for line in lines:
        if line.startswith('Bounding box for object'):
            coordinates = line.split(': ')[1].strip().replace('(', '').replace(')', '').split(' - ')
            xmin, ymin = map(int, coordinates[0].split(', '))
            xmax, ymax = map(int, coordinates[1].split(', '))
            bounding_boxes.append(((xmin, ymin), (xmax, ymax)))
        elif line.startswith('Pixel mask for object'):
            mask_image_path = line.split(': ')[1].strip().replace('"', '')

    # 이미지와 마스크 이미지 불러오기
    image_path = annotation_file_path.replace('Annotation', 'PNGImages').replace('.txt', '.png')
    png_image = cv2.imread(image_path)
    png_image_rgb = cv2.cvtColor(png_image, cv2.COLOR_BGR2RGB)
    mask_image = cv2.imread(mask_image_path, 0)  # Grayscale

    # 이미지에 바운딩 박스 그리기
    boundingbox_image = png_image_rgb.copy()
    for box in bounding_boxes:
        cv2.rectangle(boundingbox_image, box[0], box[1], (255, 0, 0), 2)  # Red bounding box

    # 이미지 출력
    fig, ax = plt.subplots(1, 3, figsize=(8, 9))
    ax[0].imshow(png_image_rgb)
    ax[0].axis('off')  # Hide axes
    ax[0].set_title('Basic Image')
    ax[1].imshow(boundingbox_image)
    ax[1].axis('off')  # Hide axes
    ax[1].set_title('Bounding Image')
    ax[2].imshow(mask_image)
    ax[2].axis('off')  # Hide axes
    ax[2].set_title('Segmentation Image')
    plt.show()

for i in range(1,6):
    annotation_file_path = f'PennFudanPed/Annotation/FudanPed0000{i}.txt'
    draw_images(annotation_file_path)

class PennFudanDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        # 모든 이미지 파일으 불러오고 정렬
        self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
        self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))

    def __getitem__(self, idx):
        # 이미지와 마스크 불러오기
        img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
        mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
        img = read_image(img_path)
        mask = read_image(mask_path)
        obj_ids = torch.unique(mask)
        # 첫번째 ID는 Background 이므로 제거 (CoCodataset 기준)
        obj_ids = obj_ids[1:]
        num_objs = len(obj_ids)

        # 색상으로 인코딩된 마스크를 세트로 분할
        masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8)

        # 각 마스크의 Bounding box 좌표
        boxes = masks_to_boxes(masks)

        # there is only one class
        labels = torch.ones((num_objs,), dtype=torch.int64)

        image_id = idx
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # 모든 인스턴스가 iscrod라고 가정
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # 샘플과 대상을 torchvision tv_tensors로 래핑합니다.
        img = tv_tensors.Image(img)

        target = {}
        target["boxes"] = tv_tensors.BoundingBoxes(boxes, format="XYXY", canvas_size=F.get_size(img))
        target["masks"] = tv_tensors.Mask(masks)
        target["labels"] = labels
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
        
        # 이미지 transform 설정
        if self.transforms is not None:
            img, target = self.transforms(img, target)

        return img, target

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

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor


def get_model_instance_segmentation(num_classes):
    # Coco dataset으로 사전 훈련된 resnet50 불러오기
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

    # classification 모델의 input feature 갯수 가져오기
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 모델의 head 변경
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # segmentation 모델의 input feature 갯수 가져오기
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    
    # 수정한 Layer들 model에 적용
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask,
        hidden_layer,
        num_classes
    )

    return model

# 이미지 변환 함수
def get_transform(train):
    transforms = []
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    transforms.append(T.ToDtype(torch.float, scale=True))
    transforms.append(T.ToPureTensor())
    return T.Compose(transforms)

# GPU or CPU 설정
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# 사람, Background class 설정
num_classes = 2

# 이미지 변환 및 Dataset 생성
dataset = PennFudanDataset('data/PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('data/PennFudanPed', get_transform(train=False))

# Train, Test Dataset 분할
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

# 데이터 로더 정의
data_loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    collate_fn=utils.collate_fn
)

data_loader_test = torch.utils.data.DataLoader(
    dataset_test,
    batch_size=1,
    shuffle=False,
    collate_fn=utils.collate_fn
)

# 모델 생성
model = get_model_instance_segmentation(num_classes)
model.to(device)

# Optimizer 구성
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
    params,
    lr=0.005,
    momentum=0.9,
    weight_decay=0.0005
)

# 스케쥴러 구성
lr_scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=3,
    gamma=0.1
)

# 5회 학습
num_epochs = 5

for epoch in range(num_epochs):
    # Train - 10회 마다 loss 출력
    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=100)
    # learning rate 업데이트
    lr_scheduler.step()
    # Test 데이터로 평가
    evaluate(model, data_loader_test, device=device)
print("종료")

image = read_image("./PennFudanPed/PNGImages/PennPed00040.png")
eval_transform = get_transform(train=False)

# 모델 평가 모드 후 예측
model.eval()
with torch.no_grad():
    x = eval_transform(image)
    x = x[:3, ...].to(device)
    predictions = model([x, ])
    pred = predictions[0]


# 이미지 변환
image = (255.0 * (image - image.min()) / (image.max() - image.min())).to(torch.uint8)
image = image[:3, ...]
pred_labels = [f"pedestrian: {score:.3f}" for label, score in zip(pred["labels"], pred["scores"])]
pred_boxes = pred["boxes"].long()
output_image = draw_bounding_boxes(image, pred_boxes, pred_labels, colors="red")

# Mask 생성
masks = (pred["masks"] > 0.7).squeeze(1)
output_image = draw_segmentation_masks(output_image, masks, alpha=0.5, colors="blue")

# 이미지 확인
plt.figure(figsize=(12, 12))
plt.imshow(output_image.permute(1, 2, 0))
plt.axis('off')
plt.show()