Agora vai!

app.py

@@ -6,10 +6,9 @@ import pickle
 from PIL import Image
 import zipfile
 import os
-import os as _os
 
 # Pega o token do Hugging Face da variável de ambiente
-hf_token = _os.getenv("HUGGINGFACE_HUB_TOKEN")
+hf_token = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
 # Repositório no Hugging Face
 repo_id = "davebraga/wrdbTI6"
@@ -17,24 +16,28 @@ repo_id = "davebraga/wrdbTI6"
 # Baixar o modelo compactado com autenticação
 model_zip_path = hf_hub_download(repo_id, "saved_model.zip", token=hf_token)
 
+# Diretório onde o modelo será extraído
+extract_dir = "saved_model"
+
 # Descompactar o modelo se ainda não existir
-if not os.path.exists("saved_model"):
+if not os.path.exists(extract_dir):
     with zipfile.ZipFile(model_zip_path, 'r') as zip_ref:
-        zip_ref.extractall("saved_model")
+        zip_ref.extractall(extract_dir)
 
 # Baixar os encoders com autenticação
 category_encoder_path = hf_hub_download(repo_id, "category_encoder.pkl", token=hf_token)
 color_encoder_path = hf_hub_download(repo_id, "color_encoder.pkl", token=hf_token)
 
 # Carregar modelo e encoders
-model = load_model("saved_model")
+model = load_model(extract_dir)
+
 with open(category_encoder_path, "rb") as f:
     category_encoder = pickle.load(f)
 with open(color_encoder_path, "rb") as f:
     color_encoder = pickle.load(f)
 
 # Função de previsão
-def predict(image):
+def predict(image: Image.Image):
     image = image.resize((160, 160))
     image_array = np.array(image) / 255.0
     image_array = np.expand_dims(image_array, axis=0)
@@ -45,7 +48,6 @@ def predict(image):
 
     return f"Categoria: {category}", f"Cor: {color}"
 
-
 # Interface Gradio
 iface = gr.Interface(
     fn=predict,