sachbney/ai-detector

1import os2 3os.environ["STREAMLIT_HOME"] = "/app/.streamlit"4os.environ["XDG_CACHE_HOME"] = "/app/.cache"5os.environ["HF_HOME"] = "/tmp/huggingface"6os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"7 8import streamlit as st9import torch10import torch.nn.functional as F11from PIL import Image12from transformers import AutoImageProcessor, AutoModelForImageClassification13 14# ------------------------------15# Chemins relatifs robustes16# ------------------------------17IMAGE_PROCESSOR_PATH = os.path.join(os.getcwd(), "image-processor")18IMAGE_MODEL_PATH = os.path.join(os.getcwd(), "image-model")19 20# ------------------------------21# Chargement du modèle d’image22# ------------------------------23image_processor = AutoImageProcessor.from_pretrained(IMAGE_PROCESSOR_PATH, local_files_only=True)24image_model = AutoModelForImageClassification.from_pretrained(IMAGE_MODEL_PATH, local_files_only=True)25image_model.eval()26 27# Device28device = torch.device("cuda" if torch.cuda.is_available() else "cpu")29image_model.to(device)30 31# ------------------------------32# Fonction de prédiction33# ------------------------------34def predict_image(image):35    inputs = image_processor(images=image, return_tensors="pt").to(device)36    with torch.no_grad():37        outputs = image_model(**inputs)38        probs = F.softmax(outputs.logits, dim=-1)39        pred = torch.argmax(probs, dim=-1).item()40        confidence = probs[0][pred].item()41    label = image_model.config.id2label[pred]42    return label, round(confidence * 100, 2)43 44# ------------------------------45# Interface Streamlit46# ------------------------------47st.set_page_config(page_title="", layout="centered")48 49uploaded_file = st.file_uploader("Chargez une image", type=["png", "jpg", "jpeg"])50 51if uploaded_file:52    image = Image.open(uploaded_file).convert("RGB")53    label, proba = predict_image(image)54 55    # Choix de l'émoji56    emoji = "🤖" if label.lower() == "ia" else "🧠"57 58    # Afficher le résultat avec l'émoji59    st.success(f"{emoji} L'image est probablement : **{label}** ({proba} % de confiance)")60 61    # Puis afficher l'image62    st.image(image, caption="Image analysée", width=400)63