import warnings
import numpy as np
import requests
from PIL import Image
import gradio as gr
from lang_sam import LangSAM
model = LangSAM()

text_prompt = """
Focus on the inner circular region captured by the endoscopic camera. Identify only small, raised, and rounded growths (polyps) with smooth and well-defined edges. These polyps typically appear brighter or redder than the surrounding tissue and have a soft, protruding texture. Ignore any flat areas, reflections, or irregular shapes that do not match these criteria. Avoid including empty or uniformly colored regions without distinct structures.
Multiple polyps may exist in one image.
"""


def highlight_mask_on_image(image, mask_np, color=(0, 200, 0), alpha=0.7):
    """Tô màu cho khu vực được xác định bởi mask lên hình ảnh gốc."""
    image_np = np.array(image)

    if mask_np.shape != image_np.shape[:2]:
        raise ValueError("Kích thước của mask không khớp với kích thước của hình ảnh")

    highlighted_image = image_np.copy()

    mask_indices = mask_np > 0  # Chỉ lấy các vùng có mask
    highlighted_image[mask_indices] = (1 - alpha) * image_np[mask_indices] + alpha * np.array(color)

    return Image.fromarray(highlighted_image.astype(np.uint8))
def display_masks_np(masks_np):
    """
    Hiển thị từng mask từ danh sách masks_np.

    Parameters:
        masks_np (list): Danh sách các mask dưới dạng numpy array.

    Returns:
        None: Hiển thị từng mask.
    """
    import matplotlib.pyplot as plt

    print(f"Number of masks to display: {len(masks_np)}")
    for i, mask_np in enumerate(masks_np):
        print(f"Mask {i + 1} statistics: min={np.min(mask_np)}, max={np.max(mask_np)}, sum={np.sum(mask_np)}")

        if np.sum(mask_np) > 0:  # Chỉ hiển thị mask nếu không rỗng
            plt.imshow(mask_np, cmap="gray")
            plt.title(f"Mask {i + 1}")
            plt.axis("off")
            plt.show()
        else:
            print(f"Mask {i + 1} is empty (all zeros).")

def main(image):
    # Convert image to RGB
    image_pil = image.convert("RGB")

    # Get prediction results from the model
    res = model.predict([image_pil], [text_prompt])

    # Check if the result has the expected structure
    if not isinstance(res, list) or len(res) == 0 or not isinstance(res[0], dict):
        raise ValueError("Unexpected result structure from model.predict")

    # Extract masks, boxes, phrases, and logits from the result
    masks = res[0].get("masks", [])
    boxes = res[0].get("boxes", [])
    phrases = res[0].get("phrases", [])
    logits = res[0].get("scores", [])

    # Handle the case where no masks are detected
    if len(masks) == 0:
        print(f"No objects of the '{text_prompt}' prompt detected in the image.")
        return image_pil  # Return the original image if no masks are detected

    # Convert masks to numpy arrays (if not already numpy arrays)
    masks_np = [mask.squeeze() if isinstance(mask, np.ndarray) else mask.cpu().numpy() for mask in masks]

    # Skip the first mask and process only masks from the second one onwards
    if len(masks_np) > 1:
        masks_np = masks_np[1:]  # Bỏ mask đầu tiên

    # Combine masks if there are multiple masks
    if len(masks_np) > 1:
        combined_mask = np.sum(masks_np, axis=0) > 0
    else:
        combined_mask = masks_np[0]

    # Highlight the combined mask on the original image
    highlighted_image = highlight_mask_on_image(image_pil, combined_mask)
    #display_masks_np(masks_np)  # Display the remaining masks
    return highlighted_image  # Return the highlighted image
    
def create_polyb():
  with gr.Blocks() as demo:
      gr.Markdown("Hãy tải ảnh lên và nhấn **Xử Lý** để khoanh vùng Polyp.")
      with gr.Row():
          inp = gr.Image(label= "Nhập Ảnh",type="pil",height=512, width=512,value="../anh/viemphoi.jpeg",interactive=True)
          out = gr.Image(label="Kết Quả Dự Đoán", type = 'pil', height=512, width=512)
      btn = gr.Button("Xử Lý")
      btn.click(fn=main, inputs=inp, outputs=out)
  return demo