fixed bug in segmentation_utils.py

segmentation_utils.py

@@ -2,7 +2,6 @@ import requests
 from pycocotools import mask
 import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw, ImageOps, ImageFont
-from dotenv import find_dotenv, load_dotenv
 import os
 import base64
 import io
@@ -11,13 +10,94 @@ import numpy as np
 import cv2
 from image_utils import print_text_on_image_centered, create_background_image
 from icecream import ic
+import traceback
+from pprint import pprint
 
 
-load_dotenv(find_dotenv())
-HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-API_URL = "https://api-inference.huggingface.co/models/facebook/mask2former-swin-tiny-coco-panoptic"
-headers = {"Authorization": f"Bearer {HUGGINGFACEHUB_API_TOKEN}"}
+
+# Función para transformar la entrada en un array de numpy
+# Si la entrada es una URL, descarga la imagen y la convierte en un array de numpy
+# Si la entrada es una ruta de archivo, carga la imagen y la convierte en un array de numpy
+# Si la entrada ya es un array de numpy, devuélvela tal cual
+# Si la entrada no es ninguna de las anteriores, lanza un ValueError
+
+def transform_image_to_numpy_array(input):
+    if isinstance(input, np.ndarray):
+        # Si la entrada es un array de numpy, devuélvela tal cual
+        h, w = input.shape[:2]
+        new_height = int(h * (500 / w))
+        return cv2.resize(input, (500, new_height))
+    elif isinstance(input, str):
+        # Si la entrada es una cadena, podría ser una URL o una ruta de archivo
+        if input.startswith('http://') or input.startswith('https://'):
+            # Si la entrada es una URL, descarga la imagen y conviértela en un array de numpy
+            # se necesita un header para evitar el error 403
+            headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36"}
+            response = requests.get(input, headers=headers)
+            ic(response.status_code)
+            image_array = np.frombuffer(response.content, dtype=np.uint8)
+            image = cv2.imdecode(image_array, -1)
+            
+            # Si la imagen tiene 3 canales (es decir, es una imagen en color),
+            # convertirla de BGR a RGB
+            if image.ndim == 3:
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            image = Image.fromarray(image).convert("RGBA") 
+            image = np.array(image)
+        else:
+            # Si la entrada es una ruta de archivo, carga la imagen y conviértela en un array de numpy
+            image = cv2.imread(input)
+       
+        h, w = image.shape[:2]
+        new_height = int(h * (500 / w))
+        return cv2.resize(image, (500, new_height))
+    else:
+        raise ValueError("La entrada no es un array de numpy, una URL ni una ruta de archivo.")
+
+def transform_image_to_numpy_array2(input):
+    if isinstance(input, np.ndarray):
+        # Si la entrada es un array de numpy, devuélvela tal cual
+        return cv2.resize(input, (500, 500))
+    elif isinstance(input, str):
+        # Si la entrada es una cadena, podría ser una URL o una ruta de archivo
+        if input.startswith('http://') or input.startswith('https://'):
+            # Si la entrada es una URL, descarga la imagen y conviértela en un array de numpy
+            # se necesita un header para evitar el error 403
+            headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36"}
+            response = requests.get(input, headers=headers)
+            ic(response.status_code)
+            image_array = np.frombuffer(response.content, dtype=np.uint8)
+            image = cv2.imdecode(image_array, -1)
+            
+             # Si la imagen tiene 3 canales (es decir, es una imagen en color),
+             # convertirla de BGR a RGB
+            if image.ndim == 3:
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            image = Image.fromarray(image).convert("RGBA") 
+            image = np.array(image)
+        else:
+            # Si la entrada es una ruta de archivo, carga la imagen y conviértela en un array de numpy
+            image = cv2.imread(input)
+       
+        return cv2.resize(image, (500, 500))
+    else:
+        raise ValueError("La entrada no es un array de numpy, una URL ni una ruta de archivo.")
+    
+def segment_image_from_numpy(image_array, api_token, model):
+    
+    #API_URL = "https://api-inference.huggingface.co/models/facebook/mask2former-swin-tiny-coco-panoptic"
+    API_URL = f"https://api-inference.huggingface.co/models/{model}"
+    headers = {"Authorization": f"Bearer {api_token}"}
+    ic(API_URL)
+    ic(headers)
+    # Convert the image to bytes
+    is_success, im_buf_arr = cv2.imencode(".jpg", image_array)
+    data = im_buf_arr.tobytes()
+    response = requests.post(API_URL, headers=headers, data=data)
+    pprint(response.json())
+    return response.json()  
+
 
 def segment_image_from_path(image_path):
     with open(image_path, "rb") as f:
@@ -39,13 +119,90 @@ def decode_mask(mask_str, size):
     mask_image = mask_image.resize(size).convert("L")
     return mask_image
 
+
 def overlay_masks_on_image(image, segments, transparency=0.4):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    original_image = image
+    if original_image.mode != 'RGBA':
+        original_image = original_image.convert('RGBA')
+    
+    overlay = Image.new("RGBA", original_image.size, (255, 255, 255, 0))
+    text_layer = Image.new("RGBA", original_image.size, (255, 255, 255, 0))  
+    
+    for segment in segments:
+        mask_str = segment['mask']
+        mask_image = decode_mask(mask_str, original_image.size)
+        color = generate_random_color()
+        
+        color_mask = ImageOps.colorize(mask_image, black="black", white=color)
+        color_mask.putalpha(mask_image)
+        
+        overlay = Image.alpha_composite(overlay, color_mask)
+        
+        # Calcula el centroide de la mascara
+        x, y = np.where(np.array(mask_image) > 0)
+        centroid_x = x.mean()
+        centroid_y = y.mean()
+        
+        # Imprime la etiqueta y la puntuación en la capa de texto
+        font_size = 30
+        draw = ImageDraw.Draw(text_layer)
+        font = ImageFont.load_default().font_variant(size=font_size)
+        label = segment['label']
+        score = segment['score']
+        text =f"{label}: {score}"
+        
+        # Calcula el tamaño del texto
+        text_bbox = draw.textbbox((0, 0), text, font=font)
+        text_width = text_bbox[2] - text_bbox[0]
+        text_height = text_bbox[3] - text_bbox[1]
+        
+        # Asegúrate de que las coordenadas del texto están dentro de los límites de la imagen
+        text_x = max(0, min(centroid_x - text_width / 2, original_image.size[0] - text_width))
+        text_y = max(0, min(centroid_y - text_height / 2, original_image.size[1] - text_height))
+
+        draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font)
+    
+    # Ajusta la transparencia de la capa de superposición
+    overlay = Image.blend(original_image, overlay, transparency)
+
+    # Combina la capa de superposición con la capa de texto
+    final_image = Image.alpha_composite(overlay, text_layer)
+    
+    return final_image
+
+
+
+
+
+
+
+
+def overlay_masks_on_image2(image, segments, transparency=0.4):
     # Convert numpy array to PIL Image
-    original_image = Image.fromarray(image).convert("RGBA")
- 
+    #original_image = Image.fromarray(image).convert("RGBA")
+    #original_image = image
     #original_image = Image.open(image).convert("RGBA")
-    overlay = Image.new("RGBA", original_image.size, (255, 255, 255, 0))
+    # para file es str
+    # para url es numpy.ndarray
+    # para cv.imread es numpy.ndarray
+    
+    # Convertir el array de numpy a una imagen PIL si es necesario
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    print(type(image))
+    print(image)
+    original_image = image
+    
+    if original_image.mode != 'RGBA':
+        original_image = original_image.convert('RGBA')
     
+    print(original_image.size)
+    overlay = Image.new("RGBA", original_image.size, (255, 255, 255, 0))
+    print(overlay.size)
     # Nueva capa para el texto
 
     text_layer = Image.new("RGBA", original_image.size, (255, 255, 255, 0))  
@@ -56,6 +213,27 @@ def overlay_masks_on_image(image, segments, transparency=0.4):
         print(segment['label'] + " " + str(segment['score']))
         mask_str = segment['mask']
         mask_image = decode_mask(mask_str, original_image.size)
+        
+        
+        
+        # Convierte la imagen de la máscara a un array de numpy
+        mask_array = np.array(mask_image)
+
+        # Encuentra los píxeles blancos
+        y, x = np.where(mask_array > 0)
+
+        # Calcula el cuadro delimitador de los píxeles blancos
+        x_min, y_min, width, height = cv2.boundingRect(np.array(list(zip(x, y))))
+
+    
+        # Crea un objeto ImageDraw para dibujar en la imagen original
+        draw = ImageDraw.Draw(original_image)   
+        
+                
+        # Dibuja el cuadro delimitador en la imagen original
+        draw.rectangle([(x_min, y_min), (x_min + width, y_min + height)], outline=(0, 255, 0), width=2)
+                        
+        
         color = generate_random_color()
         
         color_mask = ImageOps.colorize(mask_image, black="black", white=color)
@@ -63,6 +241,7 @@ def overlay_masks_on_image(image, segments, transparency=0.4):
         
         overlay = Image.alpha_composite(overlay, color_mask)
         
+        
         # Calcula el centroide de la mascara
         
         x, y = np.where(np.array(mask_image) > 0)
@@ -83,11 +262,33 @@ def overlay_masks_on_image(image, segments, transparency=0.4):
        
         text_width = 500
         text_height = 100
-        draw.text((centroid_x - text_width / 2, centroid_y - text_height / 2), text, fill=(255, 255, 255, 255), font=font)
         
+        
+        # Asegúrate de que las coordenadas del texto están dentro de los límites de la imagen
+        text_x = max(0, min(centroid_x - text_width / 2, original_image.size[0] - text_width))
+        text_y = max(0, min(centroid_y - text_height / 2, original_image.size[1] - text_height))
+# Asegúrate de que las coordenadas del texto están dentro de los límites de la imagen
+        text_x = max(0, min(centroid_x, original_image.size[0] - text_width))
+        text_y = max(0, min(centroid_y, original_image.size[1] - text_height))
+
+
+        # Calcula las coordenadas del texto
+        text_x = centroid_x - text_width / 2
+        text_y = centroid_y - text_height / 2
+        
+        
+        # Asegúrate de que las coordenadas del texto están dentro de los límites de la imagen
+        text_x = max(0, min(text_x, original_image.size[0] - text_width))
+        text_y = max(0, min(text_y, original_image.size[1] - text_height))
+        
+        
+        draw.text((centroid_x - text_width / 2, centroid_y - text_height / 2), text, fill=(255, 255, 255, 255), font=font)
+
+        #draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=font)
     
     # Ajusta la transparencia de la capa de superposición
-    
+    print(original_image.size)
+    print(overlay.size)
     overlay = Image.blend(original_image, overlay, transparency)
 
     # Combina la capa de superposición con la capa de texto
@@ -102,63 +303,71 @@ def generate_random_color():
     return (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
 
 
-def segment_and_overlay_results(image, api_token, model):
+def segment_and_overlay_results(image_path, api_token, model):
     #segments = segment_image_from_image(image)
     #final_image = overlay_masks_on_image(image, segments)
     #return final_image
     processed_image = None  # Initialize processed_image
     segments = []
+    #image_type = None
+    #if isinstance(image_path, str):
+    #    image_type = 'FILE'
+    #    image = cv2.imread('cats.jpg')
+    #elif isinstance(image_path, np.ndarray):
+    #    image_type = 'NUMPY ARRAY'
+    #else:
+    #    raise ValueError("The image is neither a Image nor a local file.")  
+          
+    #ic(image_type)
+    image = transform_image_to_numpy_array(image_path)
+    # imprime tres primeros pixeles
+    print(type(image))
+    ic(image[0, 0:3])
+    
+    
+    
+    
     try:
         #segments = segment_image_from_image(image)
         #processed_image = overlay_masks_on_image(image, segments) 
         
         # debug image contents
-        
-        ic(image)
-          
-        if image.startswith('http://') or image.startswith('https://'):
-            ic("image is a URL: " + image)
-            response = requests.get(image)
-            image = Image.open(BytesIO(response.content))
-        else:
-            # Check if image is a local file
-           
-          
-            if os.path.isfile(os.path.join(os.getcwd(), image)):
-                ic("image is a file: " + image + "OK")
-                image = Image.open(image)
-            else:
-                raise ValueError("The image is neither a URL nor a local file.")
-               
-        
-        
+
         #if os.path.isfile(image):
        #     ic ("--- image is a file ---")
        #     image = Image.open(image)
        # if image is None:
        #     ic("image is None")
        #     return None, []
-        print(image)
-        ic("--- calling segment_image_from_image ---")    
-        #segments = segment_image_from_image(image)
-        segments = segment_image_from_path('cats.jpg')
-        for segment in segments:
-            print("segmentation_utils.py segment_and_overlay_results")
-            print(segment['label'] + " " + str(segment['score']))
+        
+        ic("--- calling segment_image_from_path ---")    
+        segments = segment_image_from_numpy(image, api_token, model)
+        #if image_type == 'FILE':
+        #    segments = segment_image_from_path(image_path)
+        #if image_type == 'NUMPY ARRAY':
+        #    segments = segment_image_from_image(image_path)
+            
+        ic("--- printing segments ---")
+        for segment in segments:    
+            ic(segment['label'] ,segment['score'])
         processed_image = print_text_on_image_centered(
                     create_background_image(500, 500, "white"),
                     'SEGMENTING OK',
                     'green'
                 )
-        print("--- calling overlay_masks_on_image ---")
+        ic("--- calling overlay_masks_on_image ---")
         processed_image = overlay_masks_on_image(image, segments)
+        return processed_image, segments
     except Exception as e:
+        print("EXCEPTION")
         ic(e)
+        print(traceback.format_exc())
         processed_image = print_text_on_image_centered(
                     create_background_image(500, 500, "white"),
                     e,
                     'green'
                 )
         segments = []
-    finally:
-        return processed_image, segments
+        return processed_image, segments
+    #finally:
+    #return processed_image, segments