Spaces:
Sleeping
Sleeping
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 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) | |