import cv2
import numpy as np
import gradio as gr

def process_image(image, calibration_factor):
    """
    Processes the input image to measure the hair diameter.
    
    Args:
        image (numpy array): Uploaded image in RGB format.
        calibration_factor (float): Conversion factor from pixels to micrometers.
    
    Returns:
        tuple: A result text and the annotated image.
    """
    if image is None:
        return "No image provided.", None

    # Convert image from RGB (Gradio default) to BGR (OpenCV default)
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Convert to grayscale
    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    
    # Apply Gaussian blur to reduce noise
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # Apply thresholding (using Otsu's method) to create a binary image
    ret, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # Find contours in the thresholded image
    contours, hierarchy = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        return "No hair contour found.", image

    # Assume the hair strand is the largest contour
    hair_contour = max(contours, key=cv2.contourArea)
    
    # Fit a minimum area rectangle to the contour
    rect = cv2.minAreaRect(hair_contour)
    (center, (width, height), angle) = rect
    
    # The smaller dimension is considered the hair's diameter in pixels
    diameter_pixels = min(width, height)
    diameter_real = diameter_pixels * calibration_factor  # Convert to micrometers
    
    # Draw the fitted rectangle on the image for visualization
    box = cv2.boxPoints(rect)
    box = np.int0(box)
    annotated = image_bgr.copy()
    cv2.drawContours(annotated, [box], 0, (0, 255, 0), 2)
    
    # Convert annotated image back to RGB for Gradio display
    annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
    
    result_text = f"Estimated Hair Diameter: {diameter_real:.2f} micrometers"
    return result_text, annotated_rgb

# Build the Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Hair Image"),
        gr.Number(value=0.05, label="Calibration Factor (micrometers per pixel)")
    ],
    outputs=[
        gr.Textbox(label="Result"),
        gr.Image(label="Annotated Image")
    ],
    title="Hair Diameter Measurement Tool",
    description=(
        "Upload a high-resolution image of a hair sample (captured with a dermascope) "
        "and provide a calibration factor based on a reference object included in the image. "
        "This tool processes the image to estimate the hair diameter using image processing techniques."
    )
)

if __name__ == '__main__':
    iface.launch()