Update app.py

app.py

@@ -2,40 +2,54 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 import cv2
-import datetime
 
-# Carregue o modelo treinado
-model_path = 'modelo_treinado.h5'
-model = tf.keras.models.load_model(model_path)
+# Load the pre-trained ResNet-18 model from PyTorch
+import torch
+import requests
+from torchvision import transforms
+
+resnet_model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()
+response = requests.get("https://git.io/JJkYN")
+labels = response.text.split("\n")
+
+# Load the TensorFlow model
+tf_model_path = 'modelo_treinado.h5'
+tf_model = tf.keras.models.load_model(tf_model_path)
 
 class_labels = ["Normal", "Cataract"]
 
-def classify_image(input_image):
-    # Redimensione a imagem de entrada para o tamanho esperado pelo modelo
-    input_image = tf.image.resize(input_image, (192, 256))
-    input_image = (input_image / 255.0)
-    input_image = np.expand_dims(input_image, axis=0)
-
-    # Faça a previsão
-    prediction = model.predict(input_image)
-    class_index = np.argmax(prediction)
-    predicted_class = class_labels[class_index]
-
-    # Crie uma cópia da imagem de entrada para exibir
-    output_image = input_image[0].copy()
-    output_image = (output_image * 255).astype('uint8')
-
-    # Desenhe uma caixa de detecção na imagem
-    cv2.putText(output_image, predicted_class, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-    cv2.putText(output_image, timestamp, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
-
-    return output_image
-
-# Configurar a interface Gradio
-gr.Interface(
-    fn=classify_image,
-    inputs="image",
-    outputs="image",
-    live=True
-).launch()
+def predict(inp):
+    # First, use the ResNet-18 model to predict labels
+    inp_resized = transforms.ToTensor()(inp).unsqueeze(0)
+    with torch.no_grad():
+        prediction_resnet = torch.nn.functional.softmax(resnet_model(inp_resized)[0], dim=0)
+        confidences_resnet = {labels[i]: float(prediction_resnet[i]) for i in range(1000)}
+
+    # Then, use the TensorFlow model to predict Normal or Cataract
+    img_array = cv2.cvtColor(np.array(inp), cv2.COLOR_RGB2BGR)
+    img_array = cv2.resize(img_array, (224, 224))  # Resize to match the input size expected by the TF model
+    img_array = img_array / 255.0  # Normalize pixel values
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+
+    prediction_tf = tf_model.predict(img_array)
+    label_index = np.argmax(prediction_tf)
+    confidence_tf = float(prediction_tf[0, label_index])
+
+    # Combine the ResNet-18 and TensorFlow predictions
+    resnet_label = max(confidences_resnet, key=confidences_resnet.get)
+    if confidence_tf >= 0.5:
+        final_label = class_labels[label_index]
+        confidence = confidence_tf
+    else:
+        final_label = resnet_label
+        confidence = confidences_resnet[resnet_label]
+
+    return final_label, confidence
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=["label", "number"]
+)
+
+demo.launch()