Update app.py

app.py

@@ -3,33 +3,88 @@ import numpy as np
 import tensorflow as tf
 from PIL import Image
 import cv2
+from scipy.ndimage import binary_fill_holes
 
-# Certifique-se de que o nome do ficheiro do seu modelo é o mesmo que o que fez upload
+CLASS_COLORS = {
+    1: (0, 255, 0),       # diecast (verde)
+    2: (0, 0, 255),       # large_packaging (azul)
+    3: (255, 0, 0),       # packaging (vermelho)
+}
+CLASS_NAMES = {
+    1: "diecast",
+    2: "large_packaging",
+    3: "packaging",
+}
+
+# Caminho para o modelo
 MODEL_PATH = "segmentation_model.h5"
 model = tf.keras.models.load_model(MODEL_PATH)
 
 def predict_image(input_image):
-    # Converte a imagem de PIL para NumPy
-    img_np = np.array(input_image)
+    # Converte a imagem de PIL para NumPy (formato BGR para OpenCV)
+    original_img_np = np.array(input_image.convert('RGB'))
+    original_img_cv2 = cv2.cvtColor(original_img_np, cv2.COLOR_RGB2BGR)
 
-    # Redimensiona a imagem para o tamanho de entrada do modelo (256, 256)
-    img_resized = tf.image.resize(img_np, (256, 256))
+    # Redimensiona para o tamanho do modelo (256, 256)
+    img_resized = tf.image.resize(original_img_np, (256, 256))
     img_input = np.expand_dims(img_resized, axis=0)
 
-    # Faz a previsão
-    prediction = model.predict(img_input)
+    # Faz a previsão do modelo
+    prediction = model.predict(img_input, verbose=0)
+    
+    # Obtém a máscara e a confiança para cada pixel
     mask_predicted = np.argmax(prediction[0], axis=-1)
+    confidences = np.max(prediction[0], axis=-1)
+
+    # Redimensiona a máscara e a confiança para o tamanho da imagem original
+    original_size = original_img_np.shape[:2]
+    mask_final = cv2.resize(mask_predicted.astype(np.uint8), (original_size[1], original_size[0]), interpolation=cv2.INTER_NEAREST)
+    confidences_final = cv2.resize(confidences, (original_size[1], original_size[0]), interpolation=cv2.INTER_LINEAR)
+
+    final_img = original_img_cv2.copy()
+    confidence_threshold = 0.7  # Limiar de 80%
+
+    for class_id in np.unique(mask_final):
+        if class_id == 0:
+            continue
+        
+        class_name = CLASS_NAMES.get(class_id, f"Classe {class_id}")
+        binary_mask = (mask_final == class_id).astype(np.uint8)
+        binary_mask = binary_fill_holes(binary_mask)
+        
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        for contour in contours:
+            x_min, y_min, w, h = cv2.boundingRect(contour)
+            x_max, y_max = x_min + w, y_min + h
+            
+            region_mask = (mask_final[y_min:y_max, x_min:x_max] == class_id)
+            region_confidences = confidences_final[y_min:y_max, x_min:x_max][region_mask]
+            
+            if region_confidences.size > 0:
+                avg_confidence = np.mean(region_confidences)
+            else:
+                avg_confidence = 0
+            
+            if avg_confidence > confidence_threshold:
+                label_text = f"{class_name}: {avg_confidence:.2f}%"
+                color_tuple = CLASS_COLORS.get(class_id, (255, 255, 255))
+                
+                cv2.rectangle(final_img, (x_min, y_min), (x_max, y_max), color_tuple, 2)
+                
+                (text_width, text_height), baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+                cv2.rectangle(final_img, (x_min, y_min - text_height - 10), (x_min + text_width, y_min), color_tuple, -1)
+                cv2.putText(final_img, label_text, (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
+    
 
-    # Converte a máscara de volta para uma imagem e redimensiona
-    mask_image = Image.fromarray(mask_predicted.astype(np.uint8) * 50) # *50 para visualizar melhor
+    final_img_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(final_img_rgb)
 
-    return mask_image
 
-# Define a interface Gradio
 gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs="image",
-    title="Ferramenta de Segmentação de Objetos",
-    description="Carregue uma imagem e o modelo irá segmentar os objetos detetados."
+    title="Ferramenta de Segmentação e Detecção de Objetos",
+    description="Carregue uma imagem e o modelo irá detetar objetos com caixas e confiança."
 ).launch()