from f_segment_img import *
from f_measurents import *
import gradio as gr
import dotenv
import ast
dotenv.load_dotenv()
# 
def create_sam():
    sam_checkpoint = "sam_vit_h_4b8939.pth"
    model_type = "vit_h"; device = "cuda"
    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    return sam

def plt2arr(fig, draw=True):
    if draw: fig.canvas.draw()
    rgba_buf = fig.canvas.buffer_rgba()
    (w,h) = fig.canvas.get_width_height()
    rgba_arr = np.frombuffer(rgba_buf, dtype=np.uint8).reshape((h,w,4))
    return rgba_arr

def frame_size_width_mm(dropdown_label):
    if dropdown_label == 'Small (142 mm)': frame_width_px = 142
    elif dropdown_label == 'Medium (xx mm)': frame_width_px = 150
    elif dropdown_label == 'Large (xx mm)': frame_width_px = 155
    return frame_width_px

# 
def ipd_app(image,dropdown_label):
    # Measure image
    landmarks = ast.literal_eval(os.environ['landmarks'])
    frame_processed, measurements = measure_landmarks_img(image, landmarks, plot_landmarks_on_img = True, plot_data_on_img = True)
    # Segment Frame
    image, img_cropped, masks_selection, objects_segmented = segment_frame_from_img(image, landmarks, create_sam())
    # Calibrate measurements
    frame_width_px = get_frame_width(masks_selection)
    frame_width_mm = frame_size_width_mm(dropdown_label)
    ipd_mm = ipd_calibration(measurements['ipd_px'], frame_width_px, frame_width_mm)
    text_ipd = 'IPD: ' + str(round(ipd_mm,2)) + ' mm'
    # Check
    sam_check = plot_sam_check_segmentation_frame(image, img_cropped, objects_segmented)
    sam_check_numpy = plt2arr(sam_check, draw = True)
    # 
    return text_ipd, frame_processed, str(measurements), sam_check_numpy


dropdown = gr.Dropdown(["Small (142 mm)", "Medium (xx mm)", "Large (xx mm)"], label="Refractives Frame Size", info="For calibration")
demo = gr.Interface(fn=ipd_app, inputs=["image",dropdown], outputs=["text", "image", "text", "image"])
demo.launch(debug=True)