app.py

import streamlit as st
from PIL import Image
from io import BytesIO
import requests
import os
import time
import torch
import torch.nn.functional as F
from torchvision.transforms import (Compose,
                                    Normalize,
                                    CenterCrop,
                                    Resize,
                                    ToTensor)

# Variabel global untuk menyimpan token API
TOKEN = st.secrets["TOKEN"]

# Header yang digunakan untuk mengautentikasi ke API dengan token
headers = {"Authorization": f"Bearer {TOKEN}"}

# Perangkat yang digunakan (misal: CPU) untuk mengeksekusi model PyTorch
device = torch.device('cpu')

# Transformasi data untuk memproses gambar sebelum penggunaan pada model
pred_transforms = Compose(
    [
        Resize(224),
        CenterCrop(224),
        ToTensor(),
        Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
    ]
)

# Kamus untuk memetakan indeks kelas menjadi label kematangan pisang
id2label = {0: 'matang', 1: 'mentah', 2: 'setengah-matang', 3: 'terlalu-matang'}


# Fungsi untuk mengirimkan permintaan prediksi gambar ke API menggunakan model dari Hugging Face
def query_hf(filename):
    # Membaca data gambar dari file yang diunggah
    with open(filename, "rb") as f:
        data = f.read()

    # Mengirimkan permintaan prediksi gambar ke API menggunakan HTTP POST request dengan data gambar
    response = requests.post(API_URL, headers=headers, data=data)
    return response.json()


# Fungsi untuk memproses gambar menggunakan model PyTorch
def query_pytorch(image, model):
    # Mengubah gambar menjadi tensor dan memprosesnya dengan model PyTorch
    input_tensor = pred_transforms(image.convert("RGB")).to(device)
    encoding = input_tensor.unsqueeze(0)
    with torch.no_grad():
        # Mengubah gambar menjadi tensor dan memprosesnya dengan model PyTorch
        input_tensor = pred_transforms(image.convert("RGB"))
        encoding = input_tensor.unsqueeze(0)
        
        # Mengirimkan gambar yang telah diubah menjadi tensor melalui model PyTorch untuk mendapatkan output prediksi
        outputs = model(encoding)

        # Menghitung probabilitas kelas menggunakan fungsi softmax
        probabilities = F.softmax(outputs, dim=1)

    return probabilities



# Fungsi untuk menampilkan hasil prediksi menggunakan model dari Hugging Face
def hunggingface_pred(upload):
    # Membuka gambar yang diunggah
    image = Image.open(upload) 
    
    with st.spinner("memprediksi..."):
        # Mengirimkan permintaan prediksi gambar ke API menggunakan model dari Hugging Face
        response = query_hf(upload)

        # Jika terdapat kesalahan pada response (misalnya model sedang dimuat), tampilkan indikator progres
        while 'error' in response:
            estimated_time = round(response['estimated_time'])
            progress_text = "Model sedang dimuat. Mohon tunggu..."
            my_bar = st.progress(0, text=progress_text)

            start_time = time.time()
            # Menunggu hingga model selesai dimuat berdasarkan estimated_time
            while time.time() - start_time < estimated_time:
                time.sleep(0.1)
                elapsed_time = time.time() - start_time
                progress_percent = min(
                    100, int((elapsed_time / estimated_time) * 100))
                my_bar.progress(progress_percent,
                                text=f"{progress_text} {progress_percent}%")

            # Mengirimkan permintaan kembali jika terdapat kesalahan pada response
            response = query_hf(upload)
            
            # Menghilangkan progress bar saat sudah mencapai 100%
            if my_bar is not None:
                my_bar.empty()

        # Menampilkan gambar yang diunggah dan hasil prediksi dari model Vision Transformer
        img1.image(image, width=300)
        for i, item in enumerate(response, start=1):
            label = item["label"]
            score = item["score"]
            score_percent = score * 100
            labels[i - 1].text(f"{i}. {label:<15}: {score_percent:.1f}%")


# Fungsi untuk menampilkan hasil prediksi menggunakan model dari PyTorch
def pytorch_pred(upload, PATH):
    # Membuka gambar yang diunggah
    image = Image.open(upload)

    with st.spinner("memprediksi..."):
        # Memuat model PyTorch yang telah dilatih
        model = torch.load(PATH, map_location=device)
        model = model.eval()

        # Memproses gambar dengan model PyTorch dan menghitung probabilitas kelas menggunakan fungsi softmax
        response = query_pytorch(image, model)

        # Menampilkan gambar yang diunggah dan hasil prediksi dari model Convolutional Neural Network
        img1.image(image, width=300)

        # Mengurutkan hasil prediksi berdasarkan probabilitas tertinggi
        sorted_indices = torch.argsort(response, descending=True)
        for i, idx in enumerate(sorted_indices[0]):
            class_name = id2label[idx.item()]
            score_percent = response[0, idx].item() * 100
            labels[i].text(f"{i+1}. {class_name:<15}: {score_percent:.1f}%")


# Fungsi untuk mendapatkan ekstensi file dari nama file yang diunggah
def get_extension(file_name):
    _, ext = os.path.splitext(file_name)
    return ext.lower()

# Fungsi untuk mengubah gambar menjadi format PNG
def convert_image(img):
    buf = BytesIO()
    img.save(buf, format="PNG")
    byte_im = buf.getvalue()
    return byte_im



# Mengatur konfigurasi halaman aplikasi
st.set_page_config(
    layout="wide", page_title="KlaKePi", page_icon=":banana:"
)

# Menampilkan judul dan deskripsi aplikasi
st.write("## Klasifikasi Tingkat Kematangan Pisang (KlaKePi)")
st.write("Unggah gambar pisang untuk memprediksi tingkat kematangannya.")

st.markdown("---")
# Menampilkan bagian pengaturan arsitektur model
st.write("## Pengaturan :wrench:")
arsitektur = st.radio(
    "Arsitektur Model",
    ('Vision Transformer (ViT)', 'Convolutional Neural Network (CNN)')
)

# Jika dipilih arsitektur Vision Transformer (ViT), maka tampilkan pilihan varian model
if arsitektur == 'Vision Transformer (ViT)':
    option = st.selectbox(
        "Varian Model",
        ("ViT-B/16-in1k", "ViT-B/32-in1k", "ViT-L/16-in1k", "ViT-L/32-in1k", "ViT-H/14-in1k",
         "ViT-B/16-in21k", "ViT-B/32-in21k", "ViT-L/16-in21k", "ViT-L/32-in21k", "ViT-H/14-in21k"),
        index=7, help="in1k = Model pre-trained pada ImageNet-1k dan fine-tuned pada citra pisang primer\n\nin21k = Model pre-trained pada ImageNet-21k dan fine-tuned pada citra pisang primer"
    )

# Jika dipilih arsitektur Convolutional Neural Network (CNN), maka tampilkan pilihan varian model
elif arsitektur == 'Convolutional Neural Network (CNN)':
    option = st.selectbox(
        "Varian Model",
        ("Alexnet-in1k", "Densenet121-in1k", "Densenet201-in1k", "MobileNet V2-in1k", "MobileNet V3-in1k",
         "Resnet152-in1k", "Resnet50-in1k", "Squeezenet-in1k", "VGG16-in1k", "VGG19-in1k"),
        index=0, help="in1k = Model pre-trained pada ImageNet-1k dan fine-tuned pada citra pisang primer"
    )

st.markdown("---")
# Menampilkan bagian untuk mengunggah gambar pisang
st.write("## Unggah Gambar :arrow_down:")
my_upload = st.file_uploader(
    "Unggah gambar pisang", type=["png", "jpg", "jpeg"]
)
if my_upload is not None:
    # Menyimpan gambar yang diunggah ke dalam folder 'uploaded_images'
    ext = get_extension(my_upload.name)
    output_filename_with_ext = f"pisangup{ext}"
    if not os.path.exists("uploaded_images"):
        os.makedirs("uploaded_images")
    with open(os.path.join("uploaded_images", output_filename_with_ext), "wb") as f:
        f.write(my_upload.read())

st.markdown("---")
# Menampilkan bagian output hasil prediksi
st.write("## Output :arrow_down:")

# Membagi halaman menjadi dua kolom
col1, col2 = st.columns(2)

# Menampilkan gambar pada kolom pertama
col1.write("### Gambar :camera:")
img1 = col1.empty()

# Menampilkan teks header untuk hasil prediksi pada kolom kedua
col2.write("### Label dan Skor Terdeteksi:")
labels = [col2.empty() for _ in range(4)]

# Daftar path model yang tersedia
model_paths = {
    "Alexnet-in1k": "./models/alexnet.pt",
    "Densenet121-in1k": "./models/densenet121.pt",
    "Densenet201-in1k": "./models/densenet201.pt",
    "MobileNet V2-in1k": "./models/mobilenet_V2.pt",
    "MobileNet V3-in1k": "./models/mobilenet_V3.pt",
    "Resnet50-in1k": "./models/resnet50.pt",
    "Resnet152-in1k": "./models/resnet152.pt",
    "Squeezenet-in1k": "./models/squeezenet.pt",
    "VGG16-in1k": "./models/vgg16.pt",
    "VGG19-in1k": "./models/vgg19.pt",
    "ViT-B/16-in1k": "./models/vit_b_16.pt",
    "ViT-B/32-in1k": "./models/vit_b_32.pt",
    "ViT-L/16-in1k": "./models/vit_l_16.pt",
    "ViT-L/32-in1k": "./models/vit_l_32.pt",
    "ViT-H/14-in1k": "./models/vit_h_14.pt",
}

# Daftar API URL untuk model Vision Transformer
api_urls = {
    "ViT-B/16-in21k": "https://api-inference.huggingface.co/models/aryap2/kematangan-pisang-vit-b-16-100eph-224-v4.4",
    "ViT-B/32-in21k": "https://api-inference.huggingface.co/models/aryap2/kematangan-pisang-vit-b-32-100eph-224-v4.4",
    "ViT-L/16-in21k": "https://api-inference.huggingface.co/models/aryap2/kematangan-pisang-vit-l-16-100eph-224-v4.4",
    "ViT-L/32-in21k": "https://api-inference.huggingface.co/models/aryap2/kematangan-pisang-vit-l-32-100eph-224-v4.4",
    "ViT-H/14-in21k": "https://api-inference.huggingface.co/models/aryap2/kematangan-pisang-vit-h-14-100eph-224-telyu",
}

# Jika model yang dipilih ada dalam daftar path model, maka gunakan model tersebut dengan fungsi pytorch_pred
if option in model_paths:
    PATH = model_paths[option]
    if my_upload is not None:
        pytorch_pred(f"./uploaded_images/{output_filename_with_ext}", PATH)
    else:
        pytorch_pred("./pisang.jpg", PATH)

# Jika model yang dipilih ada dalam daftar API URL, maka gunakan model tersebut dengan fungsi hunggingface_pred
elif option in api_urls:
    API_URL = api_urls[option]
    if my_upload is not None:
        hunggingface_pred(f"./uploaded_images/{output_filename_with_ext}")
    else:
        hunggingface_pred("./pisang.jpg")