Update app.py

app.py

@@ -1,50 +1,27 @@
 import gradio as gr
 import tensorflow as tf
-import numpy as np
-import cv2
-
-# 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")
+import cv2
+import numpy as np
 
 # Load the TensorFlow model
-tf_model_path = 'modelo_treinado.h5'
+tf_model_path = 'modelo_treinado.h5'  # Update with the path to your TensorFlow model
 tf_model = tf.keras.models.load_model(tf_model_path)
 
 class_labels = ["Normal", "Cataract"]
 
 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
+    # 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
+    img_array = cv2.resize(img_array, (224, 224))
+    img_array = img_array / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
 
     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
+    return class_labels[label_index], confidence_tf
 
 demo = gr.Interface(
     fn=predict,