Pashtetuum's picture
Update app.py
d5dab40 verified
raw
history blame
No virus
2.12 kB
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import autokeras as ak
import pickle
import gradio as gr
import cv2
import io
import zipfile
#Загружаем модель
loaded_model = load_model('best_model', custom_objects=ak.CUSTOM_OBJECTS)
# Загружаем label_encoder
with open('label_encoder.pkl', 'rb') as le_file:
label_encoder = pickle.load(le_file)
def translator(Defect):
translation_dict = {
"Crazing": "Трещины",
"Inclusion": "Вкрапления",
"Patches": "Пятна",
"Pitted": "Рябь",
"Rolled": "Замятие",
"Scratches": "Царапины"
}
return translation_dict.get(Defect)
def predict(zip_bytes):
predictions = []
with zipfile.ZipFile(io.BytesIO(zip_bytes), 'r') as zip_ref:
for filename in zip_ref.namelist():
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
with zip_ref.open(filename) as file:
img_bytes = file.read()
img_np = np.frombuffer(img_bytes, np.uint8)
img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)
img_resized = cv2.resize(img, (200, 200))
img_array = np.expand_dims(img_resized, axis=0)
prediction = loaded_model.predict(img_array)
predicted_class = np.argmax(prediction)
decoded_class = label_encoder.inverse_transform([predicted_class])[0]
predicted_class_translation = translator(decoded_class)
predictions.append((filename, predicted_class_translation))
return predictions
interface = gr.Interface(
fn=predict,
inputs=gr.File(label="Выберите zip-файл (только с изображениями)"),
outputs=gr.Textbox(label="Предсказанный класс"),
title="Нейронная сеть для классификации брака на металлопрокате"
)
interface.launch(share=True)