IA_models / main.py
imamsaid's picture
Update main.py
4ab2b9e verified
import gradio as gr
from tensorflow.keras.models import load_model
import numpy as np
# Charger les modèles
model_TB = load_model("Tuberculosis_model.h5")
model_Pneumonia = load_model("Pneunomia_model.h5")
model_ImageTB = load_model("Image_TB_classifier.h5")
def predict(model_name, input_text):
# Transformer l'entrée texte en tableau numpy
# Exemple simple : input_text est des chiffres séparés par virgule
try:
input_array = np.array([[float(x) for x in input_text.split(",")]])
except:
return "Format d'entrée incorrect. Utilisez des nombres séparés par des virgules."
# Choix du modèle
if model_name == "Tuberculosis":
pred = model_TB.predict(input_array)
elif model_name == "Pneumonia":
pred = model_Pneumonia.predict(input_array)
elif model_name == "Image TB Classifier":
pred = model_ImageTB.predict(input_array)
else:
return "Modèle inconnu."
return str(pred[0])
# Interface Gradio
iface = gr.Interface(
fn=predict,
inputs=[
gr.Dropdown(["Tuberculosis", "Pneumonia", "Image TB Classifier"], label="Choisir le modèle"),
gr.Textbox(lines=1, placeholder="Entrée : valeurs séparées par virgule")
],
outputs="text",
title="Prédiction IA TeamAI",
description="Tester les modèles Tuberculosis, Pneumonia, Image TB Classifier"
)
if __name__ == "__main__":
iface.launch()