import gradio as gr
import spaces
from cellpose import models
import numpy as np
import cv2
import matplotlib.pyplot as plt
import tempfile
from PIL import Image
import io 
from huggingface_hub import hf_hub_download

HF_REPO_ID = "myang4218/cellposemodel"

MODEL_OPTIONS = {
    "Hemocytometer Model": "hemocytometermodel.npy",
    "General Model": "generalmodel.npy"
}

loaded_models = {}

@spaces.GPU
def segment_and_count(image, model_choice):
    
    model_filename = MODEL_OPTIONS[model_choice]
    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
    if model_filename in loaded_models:
        model = loaded_models[model_filename]
    else:
        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
        loaded_models[model_filename] = model
        
    # Convert PIL Image to numpy array
    image_np = np.array(image)
    
    # If grayscale, convert to RGB
    if len(image_np.shape) == 2:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
    elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
        # Handle RGBA images
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
    
    # Run Cellpose
    masks, flows, styles = model.eval(image_np, diameter=None, channels=[0, 0])
    
    # Count unique cells
    cell_count = len(np.unique(masks)) - 1  # subtract 1 for background
    
    # Create better overlay visualization
    overlay = image_np.copy().astype(np.float32)
    
    # Create colored mask overlay
    if masks.max() > 0:
        # Generate random colors for each cell
        np.random.seed(42)  # For reproducible colors
        colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
        colors[0] = [0, 0, 0]  # Background stays black
        
        # Create colored overlay
        colored_mask = colors[masks]
        
        # Blend with original image
        alpha = 0.4
        overlay = (1 - alpha) * overlay + alpha * colored_mask
    
    # Ensure values are in valid range and convert to uint8
    overlay = np.clip(overlay, 0, 255).astype(np.uint8)
    
    # Convert result to PIL Image for output
    overlay_image = Image.fromarray(overlay)
    
    return cell_count, overlay_image

# Gradio interface
demo = gr.Interface(
    fn=segment_and_count,
    inputs=[
        gr.Image(type="pil", label="Microscopy Image"),
        gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model", value="Hemocytometer Model")
    ],
    outputs=[
        gr.Number(label="Number of Cells"),
        gr.Image(type="pil", label="Segmented Overlay")
    ],
    title="Cell Counter with Cellpose",
    description="Upload a microscopy image and select a model to count the number of cells using Cellpose segmentation."
)

if __name__ == "__main__":
    demo.launch()