Spaces:
Sleeping
Sleeping
File size: 1,464 Bytes
d246850 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import streamlit as st
from utils import load_model, load_image, preprocess_image, predict
from ui import show_header, show_image
import os
# ========================================
# 🔧 Configuration
# ========================================
MODEL_DIR = "models"
MODEL_PATH = os.path.join(MODEL_DIR, "efficientnet_b3_full_ai_image_classifier.pt")
# ========================================
# 🚀 Streamlit App
# ========================================
def main():
st.set_page_config(page_title="AI Image Detector", page_icon="🧠", layout="centered")
show_header()
# Load model once and cache
@st.cache_resource
def get_model():
return load_model(MODEL_PATH)
model = get_model()
# User options
option = st.radio("Choose Input Type:", ("Upload Image", "From URL"))
img = None
if option == "Upload Image":
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_file:
img = load_image(uploaded_file)
else:
url = st.text_input("Enter Image URL")
if url:
img = load_image(url)
# Predict
if img is not None:
img_tensor = preprocess_image(img)
label, prob = predict(model, img_tensor)
show_image(img, label, prob)
else:
st.info("👆 Upload an image or enter a URL to start.")
if __name__ == "__main__":
main()
|