raratdit/krsbi-cnn-classification

KRSBI CNN Classification

Model CNN kustom untuk klasifikasi gambar KRSBI-B.

Jumlah kelas: 3

Cara load model:

import torch
import torch.nn as nn
import json

class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((4,4))
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256*4*4, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        return self.classifier(self.avgpool(self.features(x)))

with open("config.json") as f:
    cfg = json.load(f)
with open("class_to_idx.json") as f:
    class_to_idx = json.load(f)

model = SimpleCNN(num_classes=cfg["num_classes"])
state = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(state["model_state_dict"])
model.eval()