import torch
from PIL import Image, ImageOps, ImageFile
from transformers import AutoImageProcessor, AutoModelForImageClassification
from pathlib import Path
from streamlit_paste_button import paste_image_button as pbutton
import pandas as pd
# Allow loading of truncated images
ImageFile.LOAD_TRUNCATED_IMAGES = True
# --- CONFIGURATION ---
# Set the path to your saved model folder
# Make sure this path is accessible from where you run the Streamlit app
MODEL_PATH = "./material_classifier_model_final"
# -------------------- Load Model --------------------
@st.cache_resource
def load_model(model_path):
"""
Loads the fine-tuned model and processor from a local path.
Caches the model in memory for faster performance.
"""
try:
processor = AutoImageProcessor.from_pretrained(model_path)
model = AutoModelForImageClassification.from_pretrained(model_path)
return processor, model
except Exception as e:
st.error(f"Error loading the model …
