[CV] VGG19 fine-tuning - Image Classifiaction for chart, diagram, table

◾프로젝트 소개

BITAMin 12기 마지막 프로젝트로 Auto Readme Generator 프로젝트를 진행하고 있다.
발표용 pdf 자료를 입력하면 자동으로 Github Repository용 readme 파일을 생성해주는 서비스를 개발하는 것이 목표이다.

이를 위해 여러 작업을 수행해야 하는데, 크게 이미지 처리 작업과 텍스트 처리 작업으로 나뉜다.

그 중 이미지 처리에서는 pdf에서 이미지를 추출하여 해당 이미지가 readme 파일에 포함되어야 할 중요한 이미지인지 아닌지를 구분하는 작업이 필요하다.

해당 작업을 위해 우리는 일반적인 머신러닝 및 딥러닝 프로젝트에서 chart(plot, graph), diagram, table 이미지가 중요하게 사용된다는 점에 착안하여 이미지 분류를 통해 주요 이미지를 선택하자는 아이디어를 떠올렸다.

기존에 chart나 diagram이나 table 이미지 분류에 대한 선행 프로젝트들을 조사하다가 높은 정확도를 보이는 chart classification 프로젝트를 발견했다.
Chart-Classification-Using-CNN--Keras 깃헙에서 관련 내용을 볼 수 있다.
기존에 CV에 대해 공부를 많이 한 상태가 아니었기에 가장 기본적인 CNN 기반 모델을 사용한다는 점에서 난이도도 쉬워 보였고, 성능도 해당 깃헙에서 잘 나와서 참고하기 좋았다.
해당 깃헙에서는 VGG 모델을 fine tuning하여 정확도 100%를 달성했다.
해당 선행 프로젝트의 단점이라면 데이터셋이 너무 작고 이미지에 다양성이 부족해 성능에 대한 신뢰도가 떨어진다는 점이었다.

◾VGG

🔻모델 간단 소개

VGG 모델은 3 by 3 filter의 Convolutional Layer를 겹겹이 쌓은 모델이다. VGG19는 19겹, VGG16은 16겹을 쌓았다.
convolutional filter의 크기를 3x3으로 한 이유는 이것이 이미지의 좌우/상하/중앙 등의 특징을 파악하는 데 용이한 최소 단위이기 때문이라고 한다.
VGG 논문에 따르면 classifiaction task의 accuracy는 이미지의 representation depth와 연관이 크다고 한다.

🔻모델 fine tuning

기본적인 구조는 Chart-Classification-Using-CNN--Keras 깃헙의 FineTune.py 코드를 참고하였다.
기존 코드가 Keras 기반으로 작성되어서 GPT에 torch 코드로 수정해주라고 했고, 오류가 나는 부분과 이미지 로드 부분 등을 수정하여 사용했다.

🔸데이터셋, 이미지 로드 및 전처리

RoboFlow를 사용해 auto-training을 시킬 땐 이미지 증강을 사용했을 때 성능이 0.1~0.2%p 증가했다.
그러나 lcoal 환경에서는 오히려 이미지 증강 기법이 덜 효과적이었다. 따라서 차라리 이미지 증강을 사용하기보단 데이터셋을 더 다양하고 많이 구성하는 것이 효과적이었다.
현재 우리팀이 수집한 데이터셋은 최대한 다양한 chart, table, diagram으로 구성하려고 노력했다.
diagram의 경우, 데이터셋을 구성할 때 그 기준이 애매했다. 어디까지 diagram으로 인정해야 하는지 결정하기 어려웠다. 그래서 일단 ML/DL 분야의 pdf라는 가정 하에 최대한 화살표로 어떤 flow를 설명하는 이미지를 diagram 데이터셋에 포함시켰다. 또 각종 모델 아키텍쳐, use case diagram, as-is/to-be diagram, class diagram 등 일반적인 다이어그램도 직접 검색하여 수집했다.
수집한 데이터셋별 출처는 다음과 같다.

• chart : 1,901
• Chart-Classification-Using-CNN--Keras(horizontal bar plot, scatter plot, line plot, vertical bar plot, pie plot) : 1,050
• Chart2TextImages(line plot, bar plot) : 200
• manually generated charts(confusion-matrix, heatmap, box plot, violoin plot) : 651
• diagram : 1,815
• DQA-Net : 1,589
• google search(model architectures, flow charts, as-is/to-be diagram, usecase diagram, ERD, etc) : 226
• table : 1,867
• TNCR_Dataset : 660
• PubTab Net : 1,207
• none : 1800
• DQA-Net 중 diagram이라고 구분하고 싶지 않은 유형의 이미지들
• google search randomly

Image Preprocessing & Loading

data_path = "<YOUR DATA DIR PATH>"
class ResizeWithPadding:
    def __init__(self, size=(244,244), bg_color=(255,255,255), fill=0):
        self.size = size
        self.bg_color = bg_color
        self.fill = fill

    def __call__(self, image):
        original_ratio = image.width / image.height
        output_ratio = self.size[0] / self.size[1]
        # Determine the new size that fits within the output size while maintaining the aspect ratio
        if original_ratio > output_ratio:
            # Fit to width
            new_width = self.size[0]
            new_height = int(new_width / original_ratio)
        else:
            # Fit to height
            new_height = self.size[1]
            new_width = int(new_height * original_ratio)
        
        # Resize the image
        resized_image = image.resize((new_width, new_height), Image.LANCZOS)

        # Create a new image with the desired size and background color
        new_image = Image.new("RGB", self.size, self.bg_color)
        
        # Paste the resized image onto the center of the new image
        paste_position = ((self.size[0] - new_width) // 2, (self.size[1] - new_height) // 2)
        new_image.paste(resized_image, paste_position)

        return new_image

image_datasets = datasets.ImageFolder(data_path, transform=transforms.Compose([
    ResizeWithPadding(size=(244,244)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]))

데이터 로더는 batch_size 20으로 구성했다. 참고했던 Chart-Classification-Using-CNN--Keras 깃헙에서는 데이터 개수가 적어 batch_size를 4로만 설정했는데 batch_size 20~60까지 다양하게 실험해본 결과 batch size가 작아도 성능이 잘 나와 20으로 설정했다. train-val-test 비율은 무난하게 70:20:10으로 했다.

DataLoader

num_samples = len(image_datasets)
combined_dataset = []
for image,label in image_datasets:
    combined_dataset.append((image,label))

### 📢데이터셋 분할
# 데이터를 70:20:10 비율로 분할
total_size = len(combined_dataset)
train_size = int(0.7 * total_size)
val_size = int(0.2 * total_size)
test_size = total_size - train_size - val_size

# 무작위로 분할
random.shuffle(combined_dataset)
train_dataset = combined_dataset[:train_size]
val_dataset = combined_dataset[train_size:train_size + val_size]
test_dataset = combined_dataset[train_size + val_size:]

# DataLoader 생성
def collate_fn(batch):
    images, labels = zip(*batch)
    return torch.stack(images, 0), torch.tensor(labels)

dataloaders = {
    'train' : DataLoader(train_dataset, batch_size=20, shuffle=True, 
                            collate_fn=collate_fn), #, num_workers=1) # Windows의 경우 num_workers를 설정하면 error가 날 수 있음.
    'val' : DataLoader(val_dataset, batch_size=20, shuffle=True, 
                        collate_fn=collate_fn), #, num_workers=1)
    'test': DataLoader(test_dataset, batch_size=20, shuffle=False, 
                            collate_fn=collate_fn) #, num_workers=1)
}

🔸Fine-tuning VGG19

ImageNet으로 pre-trained된 VGG19모델을 huggingface를 통해 import한다.
VGG19는 아래와 같이 구성되어 있는데 여기서 ClassifierHead만 다시 학습시키는 과정이 fine-tuning이다.

VGG19 Architecture

VGG(
(features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): ReLU(inplace=True)
    (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (24): ReLU(inplace=True)
    (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (26): ReLU(inplace=True)
    (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (31): ReLU(inplace=True)
    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (33): ReLU(inplace=True)
    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (35): ReLU(inplace=True)
    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(pre_logits): ConvMlp(
    (fc1): Conv2d(512, 4096, kernel_size=(7, 7), stride=(1, 1))
    (act1): ReLU(inplace=True)
    (drop): Dropout(p=0.0, inplace=False)
    (fc2): Conv2d(4096, 4096, kernel_size=(1, 1), stride=(1, 1))
    (act2): ReLU(inplace=True)
)
(head): ClassifierHead(
    (global_pool): SelectAdaptivePool2d(pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))
    (drop): Dropout(p=0.0, inplace=False)
    (fc): Linear(in_features=4096, out_features=1000, bias=True)
    (flatten): Identity()
)
)

모델의 ClassifierHead를 수정하는 코드는 아래와 같다.
우선 Classifier의 output dimension을 우리 프로젝트의 이미지 클래스 개수인 4로 맞춰주고, 가중치 벡터의 gradient를 classifier만 활성화시킨다.

Modifying model Classifier

vgg19_model.reset_classifier(num_classes=len(class_names))
# Freeze all layers except the last layer for fine-tune
for param in vgg19_model.parameters():
    param.requires_grad = False

for param in vgg19_model.get_classifier().parameters():
    param.requires_grad = True
    

🔸Train & Test

이후엔 일반적인 training 코드에 따라 내가 준비한 데이터셋으로 훈련시키고 그 결과를 확인하면 된다.
• Earyl Stopping 코드 : jeffheaton - app_deep_learning
• train_model, predict 함수 : Chart-Classification-Using-CNN--Keras의 함수를 torch문법에 맞게 수정

Training & Test

# Early Stopping
import copy
class EarlyStopping:
    def __init__(self, patience=10, min_delta=0, restore_best_weights=True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_model = None
        self.best_loss = None
        self.counter = 0
        self.status = ""

    def __call__(self, model, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.best_model = copy.deepcopy(model.state_dict())
        elif self.best_loss - val_loss >= self.min_delta:
            self.best_model = copy.deepcopy(model.state_dict())
            self.best_loss = val_loss
            self.counter = 0
            self.status = f"Improvement found, counter reset to {self.counter}"
        else:
            self.counter += 1
            self.status = f"No improvement in the last {self.counter} epochs"
            if self.counter >= self.patience:
                self.status = f"Early stopping triggered after {self.counter} epochs."
                if self.restore_best_weights:
                    model.load_state_dict(self.best_model)
                return True
        return False
# Training function
def train_model(model, criterion, optimizer, num_epochs=5, patience=10):
    history = {'train_accuracy': [], 'val_accuracy': [], 'train_loss': [], 'val_loss': []}
    es = EarlyStopping(patience=patience)
    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # Zero the parameter gradients
                optimizer.zero_grad()

                # Forward
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # Backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # Statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            dataset_size = train_size if phase=='train' else val_size if phase=='val' else test_size 
            epoch_loss = running_loss / dataset_size
            epoch_acc = running_corrects.double() / dataset_size

            history[f'{phase}_loss'].append(epoch_loss)
            history[f'{phase}_accuracy'].append(epoch_acc.item())
            
            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
        if es(model,val_loss=history['val_loss'][-1]):
            print("📢 Early Stopped when",epoch)
            break # Early Stopping

    return model, history

# Prediction function
def predict(model, dataloaders):
    model.eval()
    all_preds = []
    all_labels = []
    all_probs = []

    with torch.no_grad():
        for inputs, labels in dataloaders['test']:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)

            # prediction probabilities 구하기
            probs = nn.functional.softmax(outputs,dim=1)
            
            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
            all_probs.extend(probs.cpu().numpy())

    return np.array(all_labels), np.array(all_preds), np.array(all_probs)

### Train
vgg19_model, history = train_model(vgg19_model, criterion, optimizer, num_epochs=300, patience=15)
        

train accuracy 98.96%, validation accuracy 94.65%로 훈련이 종료되었다. 노트북 GPU 사용 시 52분 정도 걸렸다.
Chart-Classification-Using-CNN--Keras의 plotting 코드를 참고하여 결과를 확인해보면 chart와 table은 어느정도 고정적인 특징을 보여 구분을 잘하는 모습을 보이지만, 아무래도 diagram은 다양한 형태를 가져 성능이 조금 떨어지는 모습을 보인다.
사실 RoboFlow로 훈련시키면 정확도 99%나 나오지만, Auto-training을 시켜 정확히 어떤 모델을 사용하는지 알 수 없으므로 그냥 VGG19를 사용하기로 결정했다.

Plot Metrics

# Setting up for creating confusion metrics
cm = confusion_matrix(y_true=test_labels, y_pred=predictions)
cm_plot_labels = class_names
# Plotting Data
# Accuracy plot
fig = plt.figure()
plt.plot(history['train_accuracy'])
plt.plot(history['val_accuracy'])
plt.title('Model Accuracy')
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
plt.legend(['Train', 'Val'], loc='upper left')

# Loss plot
fig = plt.figure()
plt.plot(history['train_loss'])
plt.plot(history['val_loss'])
plt.title('Model Loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Val'], loc='upper left')

# Confusion matrix plot
fig = plt.figure()
plot_confusion_matrix(cm=cm, classes=cm_plot_labels, title='Confusion Matrix')

### Test Matrixs
print("Accuracy :",accuracy_score(test_labels,predictions))

plot_precision_recall_curve_multiclass(test_labels, probs, class_names)
plt.show()

🔻Save model checkpoints

VGG19 모델 전체의 가중치를 저장하니까 총 540MB 크기의 pth 파일이 나왔다...
깃헙에는 100MB 이상의 파일을 업로드하지 못하므로 fine-tuning한 ClassifierHead의 가중치만 저장해야 한다.

Save & Load Classifier parameter weights

# Save only Weights of Classifier
classifier_weights = vgg19_model.get_classifier().state_dict()
torch.save(classifier_weights, '<YOUR PATH>.pth')

# Load Classifier weights
clf_weights = torch.load('<YOUR PATH>.pth')
vgg19_model.get_classifier().load_state_dict(clf_weights)

◾사용해보기

bitamin_auto_readme_generator/code/assets에 fine-tuning한 classifier weights가 저장되어 있다.
이를 다운받고, hugging face에서 VGG19모델을 import한 후 classifier를 해당 weight으로 load해주면 된다.
새로운 이미지로 test할 때는 Resize만 하면 되고, 이미지 정규화는 하지 않는다.
(그런데 GPU 환경에서 학습한 모델의 checkpoints는 GPU 환경에서만 불러올 수 있는 것 같다.)

Test for Custom Image

transform_test_img = transforms.Compose([
    ResizeWithPadding(size=(244,244)) # 위에서 선언한 ResizeWithPadding 클래스가 필요하다.
    transforms.ToTensor()
])
test_img = Image.open('<YOUR IMAGE PATH>')
test_img = transform_test_img(test_img)

# Inference
vgg19_model.eval() # set as inference model
result = vgg19_model(test_img.view(1,3,244,244).to(device))

# See the result
def softmax(a):
	c = np.max(a)
	exp_a = np.exp(a - c)
	sum_exp_a = np.sum(exp_a)
	y = exp_a / sum_exp_a
	return y
def get_classname(result):
    result = result.cpu().detach().numpy().tolist()[0]
    mi = result.index(max(result))
    return [round(x,2) for x in softmax(result)], class_names[mi] # class_names도 위에서 이미 선언했다.
print(get_classname(result))
        

[test1.png, test2.png]

[test3.png, test4.png]

[test5.png]

위의 5개 이미지로 테스트를 진행해봤다.
test1,2는 diagram, test3,5는 chart, test4는 none class이다.

[0.0, 1.0, 0.0, 0.0] # test1.png softmax probability for each class
[0.0, 0.7, 0.3, 0.0]
[0.55, 0.36, 0.08, 0.01]
[0.0, 0.16, 0.84, 0.0]
[0.98, 0.01, 0.01, 0.0]
diagram diagram chart none chart 

각 클래스일 확률과 최대 확률을 보이는 클래스를 출력했을 때, 잘 예측하는 걸 볼 수 있다.
Fine-tuning 전체 코드는 bitamin_auto_readme_generator/code/image_classification/에서 볼 수 있다. 시간이 된다면, VGG보다 개선된 모델인 ResNet과 ViT pre-trained 모델로도 실험해보려 한다.