import torch
import numpy as np
import cv2
from transformers import SamModel, SamProcessor
import gradio as gr

# Load the SAM model and processor from Hugging Face
model_id = "facebook/sam-vit-huge"
device = "cuda" if torch.cuda.is_available() else "cpu"

sam = SamModel.from_pretrained(model_id).to(device)
processor = SamProcessor.from_pretrained(model_id)

def segment_rocks(image):
    # Preprocess the image
    inputs = processor(image, return_tensors="pt").to(device)
    
    # Generate image embeddings
    with torch.no_grad():
        image_embeddings = sam.get_image_embeddings(inputs["pixel_values"])
    
    # Generate masks
    masks = []
    for i in range(3):  # Generate multiple masks
        inputs = processor(
            image,
            input_points=None,
            return_tensors="pt",
            input_boxes=[[[0, 0, image.shape[1], image.shape[0]]]],
        ).to(device)
        
        with torch.no_grad():
            outputs = sam(
                input_points=inputs["input_points"],
                input_boxes=inputs["input_boxes"],
                image_embeddings=image_embeddings,
                multimask_output=True,
            )
        
        masks.extend(outputs.pred_masks.squeeze().cpu().numpy())
    
    return masks

def compute_rock_properties(mask):
    # Find contours of the mask
    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    properties = []
    for contour in contours:
        # Compute area
        area = cv2.contourArea(contour)
        
        # Compute perimeter
        perimeter = cv2.arcLength(contour, True)
        
        # Compute circularity
        circularity = 4 * np.pi * area / (perimeter ** 2) if perimeter > 0 else 0
        
        # Fit an ellipse to get major and minor axes
        if len(contour) >= 5:
            ellipse = cv2.fitEllipse(contour)
            major_axis = max(ellipse[1])
            minor_axis = min(ellipse[1])
            aspect_ratio = major_axis / minor_axis if minor_axis > 0 else 0
        else:
            major_axis = minor_axis = aspect_ratio = 0
        
        properties.append({
            'area': area,
            'perimeter': perimeter,
            'circularity': circularity,
            'major_axis': major_axis,
            'minor_axis': minor_axis,
            'aspect_ratio': aspect_ratio
        })
    
    return properties

def process_image(input_image):
    # Convert to RGB if needed
    if input_image.shape[2] == 4:  # RGBA
        input_image = cv2.cvtColor(input_image, cv2.COLOR_RGBA2RGB)
    elif len(input_image.shape) == 2:  # Grayscale
        input_image = cv2.cvtColor(input_image, cv2.COLOR_GRAY2RGB)
    
    masks = segment_rocks(input_image)
    
    results = []
    for i, mask in enumerate(masks):
        properties = compute_rock_properties(mask)
        
        # Visualize the segmentation
        masked_image = input_image.copy()
        masked_image[mask] = (masked_image[mask] * 0.7 + np.array([255, 0, 0]) * 0.3).astype(np.uint8)
        
        results.append((masked_image, f"Rock {i+1} properties: {properties}"))
    
    return results

# Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="numpy"),
    outputs=[gr.Image(type="numpy"), gr.Textbox(label="Properties")] * 3,
    title="Rock Segmentation using SAM",
    description="Upload an image to segment rocks and compute their properties."
)

# Launch the interface
iface.launch()