[Object Detection] 커스텀 모델로 Detection 하기

데이터셋은 다음 링크의 MNIST 데이터를 사용하였다.
https://www.kaggle.com/competitions/digit-recognizer/data

 

 

Dataset

데이터셋이 위와 같은 .csv 파일이므로 이미지 출력을 위해 2차원 array로 만들어 준다.

 

def show_img_from_df(df, index):
    plt.imshow(np.reshape(np.array(df.iloc[index, 1:]), (28, 28)), cmap="gray")
    
show_img_from_df(train_df, 17)

기본 이미지

 

Object Detection 하기 위한 새로운 이미지를 생성한다. 아이디어는 다음과 같다.

1. (90, 90) 사이즈의 새로운 빈 이미지를 생성한다.
2. 사이즈 범위 내의 랜덤한 (x, y) 좌표를 정한 뒤 원래의 이미지를 해당 위치에 넣는다.
3. 원래 이미지의 테두리를 bounding box로 설정한다.

new_size = 90

df = train_df
ind = randrange(1000)

img = np.reshape(np.array(df.iloc[ind, 1:]), (28, 28))

# create the new image
new_img = np.zeros((new_size, new_size))

# randomly select a bottom left corner to use for img
x_min = randrange(new_size - img.shape[0])
y_min = randrange(new_size - img.shape[0])

x_max = x_min + img.shape[0]
y_max = y_min + img.shape[0]

x_center = x_min + (x_max - x_min) / 2
y_center = y_min + (y_max - y_min) / 2

new_img[x_min:x_max, y_min:y_max] = img
new_img = cv2.rectangle(new_img, (y_max, x_min), (y_min, x_max), 255, 1)

plt.imshow(new_img, cmap="gray")
plt.show()

생성된 이미지

생성된 이미지로 새로운 dataset을 생성한다.

class CustomDataset(Dataset):
    def __init__(self, df):
        self.df = df

    def __getitem__(self, index):
        image = np.reshape(np.array(self.df.iloc[index, 1:]), (28, 28)) / 255.

        # create the new image
        new_img = np.zeros((90, 90))

        # randomly select a bottom left corner to use for img
        x_min = randrange(90 - image.shape[0])
        y_min = randrange(90 - image.shape[0])

        x_max = x_min + image.shape[0]
        y_max = y_min + image.shape[0]

        new_img[x_min:x_max, y_min:y_max] = image
        new_img = np.reshape(new_img, (1, 90, 90))

        label = [int(self.df.iloc[index, 0]), np.array([y_max, x_min, y_min, x_max]).astype("float32")]

        return new_img, label

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

 

Model

• 모델 끝단에class(digit) / bounding box를 예측하는 fully-connected-layer를 따로 둔다.
• class : 0 ~ 9 까지의 숫자 예측
• bounding box : matplotlib 기준 좌하단, 우상단 좌표 예측

class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()

        self.conv0 = nn.Conv2d(1, 16, 3, padding=2)
        self.pool0 = nn.MaxPool2d(2, 2)

        self.conv1 = nn.Conv2d(16, 16, 3, padding=3)
        self.pool1 = nn.MaxPool2d(2, 2)

        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(16*25*25, 256),
            nn.ReLU()
        )

        self.linear = nn.Linear(256, 10)  # class(digit) layer
        self.linear_bbox = nn.Linear(256, 4)  # bounding box layer

    def forward(self, x):
        x = self.conv0(x)
        x = self.pool0(x)
        x = F.relu(x)

        x = self.conv1(x)
        x = self.pool1(x)
        x = F.relu(x)

        x = self.flatten(x)
        x = self.linear_relu_stack(x)

        logits = self.linear(x)  # class(digit)
        bbox = self.linear_bbox(x)  # bounding box

        return logits, bbox

 

Training

파라미터 세팅은 다음과 같다.

• alpha : 100
• beta : 1
• class(digit) loss : CrossEntropy
• bounding box loss : MSE (Mean Squared Error)
• 전체 Loss : (alpha * class loss) + (beta * bounding box loss)
• optimizer : AdamW
• learning rate : 1e-3

loss_fn = nn.CrossEntropyLoss()
loss_mse = nn.MSELoss()

optimizer = optim.AdamW(model.parameters(), lr=1e-3)
alpha = 100
beta = 1

def train(model, dataloader, loss_fn, loss_mse, optimizer, alpha, beta):
    model.train()
    size = len(dataloader.dataset)

    loss_cls_list = []
    loss_point_list = []

    for batch, (inputs, labels) in enumerate(dataloader):
        inputs = inputs.to(device)
        label_class = labels[0].to(device)
        label_point = labels[1].to(device)

        pred_cls, pred_point = model(inputs.float())
		
        # 종합 loss
        loss = alpha*loss_fn(pred_cls, label_class) + beta*loss_mse(pred_point, label_point.float())
        
        # class(digit) loss
        loss_cls = loss_fn(pred_cls, label_class)
        
        # bounding box loss
        loss_point = loss_mse(pred_point, label_point)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
		
        # 100 batch 마다 출력
        if (batch != 0) and (batch % 100 == 0):
            loss = loss.item()
            current = batch*len(inputs)

            loss_cls = loss_cls.item()
            loss_point = loss_point.item()

            loss_cls_list.append(loss_cls)
            loss_point_list.append(loss_point)

            print(f"main loss: {loss:>7f} [{current:>5d}/{size:5d}]")
            print(f"Digit prediction loss: {loss_cls:>7f} [{current:>5d}/{size:>5d}]")
            print(f"Coordinate prediction loss: {loss_point:>7f} [{current:>5d}/{size:>5d}]")
            print("-"*20)

 

Test

• 기존의 CustomDataset과 거의 유사하지만 label이 존재하지 않는다.

class TestDataset(Dataset):
    def __init__(self, df):
        self.df = df

    def __getitem__(self, index):
        image = np.reshape(np.array(self.df.iloc[index, :]), (28, 28)) / 255.

        # create the new image
        new_img = np.zeros((90, 90))

        # randomly select a bottom left corner to use for img
        x_min = randrange(90 - image.shape[0])
        x_max = x_min + image.shape[0]
        y_min = randrange(90 - image.shape[0])
        y_max = y_min + image.shape[0]

        new_img[x_min:x_max, y_min:y_max] = image
        new_img = np.reshape(new_img, (1, 90, 90))

        return new_img

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

 

test dataset을 생성한 뒤 학습된 모델을 테스트 한다.

inputs = next(iter(test_dataloader))
index = randrange(63)

with torch.no_grad():
    pred_cls, pred_point = model(inputs.to(device).float())
    pred_cls = np.argmax(pred_cls[index].cpu().detach().numpy())

img = np.reshape(inputs[index].cpu().detach().numpy(), (90, 90))

y_max, x_min, y_min, x_max = int(pred_point[index][0]), int(pred_point[index][1]), int(pred_point[index][2]), int(pred_point[index][3])

new_img = img.copy()

# 다음 line을 실행하면 숫자는 안보이고 bounding box만 그려진다. *해결못함* 
# new_img = cv2e.rectangle(new_img, (y_max, x_min), (y_min, x_max), 255, 1)

plt.imshow(new_img, cmap="gray")
plt.plot(y_max, x_min, "or", markersize=10)
plt.plot(y_min, x_max, "og", markersize=10)
plt.show()

print("Digit:", int(pred_cls))
print(f"{y_max}, {x_min}, {y_min}, {x_max}")

 

어느정도 잘 동작하는 것을 볼 수 있다.

 

<전체 코드>

object_detection_custom_model.ipynb

0.08MB

---

📌 Reference

MNIST object detection