File size: 1,969 Bytes
ebc51c9 3ba327b ebc51c9 3ba327b ebc51c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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) |