Spaces:
Sleeping
Sleeping
import tensorflow as tf | |
import efficientnet.tfkeras as efn | |
from tensorflow.keras.layers import Input, GlobalAveragePooling2D, Dense | |
import numpy as np | |
import gradio as gr | |
from PIL import Image, ImageDraw, ImageFont | |
# Dimensões da imagem | |
IMG_HEIGHT = 224 | |
IMG_WIDTH = 224 | |
# Função para construir o modelo de detecção de objetos | |
def build_object_detection_model(img_height, img_width): | |
# Replace this with your object detection model architecture and weights | |
# For example, you can use a model from TensorFlow Hub or any other source | |
object_detection_model = None # Load your object detection model here | |
return object_detection_model | |
# Função para construir o modelo de classificação | |
def build_classification_model(img_height, img_width, n): | |
inp = Input(shape=(img_height, img_width, n)) | |
efnet = efn.EfficientNetB0( | |
input_shape=(img_height, img_width, n), | |
weights='imagenet', | |
include_top=False | |
) | |
x = efnet(inp) | |
x = GlobalAveragePooling2D()(x) | |
x = Dense(2, activation='softmax')(x) | |
model = tf.keras.Model(inputs=inp, outputs=x) | |
opt = tf.keras.optimizers.Adam(learning_rate=0.000003) | |
loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.01) | |
model.compile(optimizer=opt, loss=loss, metrics=['accuracy']) | |
return model | |
# Load the object detection and classification models | |
object_detection_model = build_object_detection_model(IMG_HEIGHT, IMG_WIDTH) | |
classification_model = build_classification_model(IMG_HEIGHT, IMG_WIDTH, 3) | |
classification_model.load_weights('modelo_treinado.h5') | |
# Function to preprocess the image for classification | |
def preprocess_image(input_image): | |
input_image = tf.image.resize(input_image, (IMG_HEIGHT, IMG_WIDTH)) | |
input_image = input_image / 255.0 | |
return input_image | |
# Function to perform object detection and classification | |
def predict_image(input_image): | |
# Realize o pré-processamento na imagem de entrada | |
input_image_classification = preprocess_image(input_image) | |
# Faça uma previsão usando o modelo de classificação carregado | |
input_image_classification = tf.expand_dims(input_image_classification, axis=0) | |
classification_prediction = classification_model.predict(input_image_classification) | |
# Perform object detection here using the object_detection_model | |
# Replace this with your object detection logic to get bounding box coordinates | |
# A saída será uma matriz de previsões (no caso de classificação de duas classes, será algo como [[probabilidade_classe_0, probabilidade_classe_1]]) | |
# Adicione lógica para interpretar o resultado e formatá-lo para exibição | |
class_names = ["Normal", "Cataract"] | |
predicted_class = class_names[np.argmax(classification_prediction)] | |
probability = classification_prediction[0][np.argmax(classification_prediction)] | |
# You can format the result with object detection bounding box and label here | |
# For example: | |
# formatted_text = f"Predicted Class: {predicted_class}\nProbability: {probability:.2%}\nObject Detection: {bounding_box_coordinates}" | |
# Create an output image with object detection | |
output_image = input_image # Replace this with your object detection visualization | |
# Convert the output image to bytes | |
output_image_bytes = Image.fromarray(np.uint8(output_image * 255)) | |
# Create an image with the label "Normal" or "Cataract" outside the image | |
draw = ImageDraw.Draw(output_image_bytes) | |
font = ImageFont.load_default() # You can customize the font and size here | |
label_text = f"Predicted Class: {predicted_class}" | |
label_size = draw.textsize(label_text, font=font) | |
label_position = (10, 10) # You can adjust the label position | |
draw.rectangle([label_position, (label_position[0] + label_size[0], label_position[1] + label_size[1])], fill="white") | |
draw.text(label_position, label_text, fill="black", font=font) | |
# Convert the image with the label to bytes | |
labeled_image_bytes = output_image_bytes.tobytes() | |
# Return both the image with object detection and the labeled image | |
return [labeled_image_bytes, f"Predicted Class: {predicted_class}"] | |
# Crie uma interface Gradio para fazer previsões | |
iface = gr.Interface( | |
fn=predict_image, | |
inputs=gr.inputs.Image(label="Upload an Image", type="file"), | |
outputs=[gr.outputs.Image(type="pil"), gr.outputs.Text(label="Prediction", type="markdown")], | |
interpretation="default" | |
) | |
# Execute a interface Gradio | |
iface.launch() | |