from f_measurents import * import cv2 import numpy as np import matplotlib.pyplot as plt import mediapipe as mp from mediapipe.tasks.python import vision import sys import os from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor # sys.path.insert( 1, '/Users/danielfiuzadosil/Documents/GitHub_Repo/Bryant_Medical/eCommerce/App_IPD [Master]/ipd_app/src/modules') # def remove_background_img(img): THRESHOLD = 0.12 # initialize mediapipe mp_selfie_segmentation = mp.solutions.selfie_segmentation selfie_segmentation = mp_selfie_segmentation.SelfieSegmentation( model_selection=1) # get the result results = selfie_segmentation.process(img) # extract segmented mask mask = np.stack((results.segmentation_mask,) * 3, axis=-1) > THRESHOLD mask_binary = mask.astype(int)*255 img_masked = img.copy() img_masked[mask_binary == 0] = 255 # return img_masked def remove_background_img_v2(image_path, MODEL_PATH='/Users/danielfiuzadosil/Documents/GitHub_Repo/Bryant_Medical/eCommerce/App_IPD [Master]/ipd_app/data/external/mediapipe_models/deeplabv3.tflite'): # THRESHOLD = 0.12 # BG_COLOR = (0, 0, 0) # black MASK_COLOR = (255, 255, 255) # white # BaseOptions = mp.tasks.BaseOptions OutputType = vision.ImageSegmenterOptions.OutputType # Create the options that will be used for ImageSegmenter base_options = BaseOptions(model_asset_path=MODEL_PATH) options = vision.ImageSegmenterOptions( base_options=base_options, output_type=OutputType.CATEGORY_MASK) # Create the MediaPipe image file that will be segmented image = mp.Image.create_from_file(image_path) with vision.ImageSegmenter.create_from_options(options) as segmenter: # Retrieve the masks for the segmented image category_masks = segmenter.segment(image) # Generate solid color images for showing the output segmentation mask. image_data = image.numpy_view() fg_image = np.zeros(image_data.shape, dtype=np.uint8) fg_image[:] = MASK_COLOR bg_image = np.zeros(image_data.shape, dtype=np.uint8) bg_image[:] = BG_COLOR # condition = np.stack( (category_masks[0].numpy_view(),) * 3, axis=-1) > THRESHOLD mask_binary = np.where(condition, fg_image, bg_image) # img_masked = image_data.copy() img_masked[mask_binary == 0] = 255 # return img_masked def segment_frame_from_img(image, landmarks, sam): # Read image image_0 = image.copy() # Generate facial landmarks mp_face_mesh = mp.solutions.face_mesh face_mesh = mp_face_mesh.FaceMesh( static_image_mode=True, max_num_faces=1, refine_landmarks=True, min_detection_confidence=0.5) # Calculate facial landmaks and other data result, mesh_points = face_mesh_points(image) df_mesh_points = mesh_points_to_df(mesh_points) width_face_px, height_face_px = get_face_dimensions_px( mesh_points, landmarks) # Calculate rectangle where the frame will likely be (based on a reference landmarks) squareBoxEyes = mesh_points[landmarks['squareBoxEyes']] # width_rectangle = squareBoxEyes[2][0] - squareBoxEyes[0][0] tolerance_x = 0.2*width_rectangle height_rectangle = squareBoxEyes[2][1] - squareBoxEyes[0][1] tolerance_y = height_rectangle*0.2 # x_start, y_start = [squareBoxEyes[0][0] - tolerance_x, squareBoxEyes[0][1] - tolerance_y] x_end, y_end = [squareBoxEyes[2][0] + tolerance_x, squareBoxEyes[2][1]] # Cropped image to region where the frames will be located img_cropped = image_0[int(y_start):int(y_end), int(x_start):int(x_end)] # Use Meta's Segment Anything Model (SAM) to segment the frame mask_generator = SamAutomaticMaskGenerator(sam) masks = mask_generator.generate(img_cropped) # Select the right object by defining the expected range of area occupied by the frame height, width, _ = img_cropped.shape area_photo = height*width area_frame_min = 0.2 area_frame_max = 0.6 # Iterate through the different masks and store the ones that fulfill the criteria masks_selection = [] objects_segmented = [] for i in range(len(masks)): mask = masks[i] area_object = mask['area'] if area_photo*area_frame_max > area_object > area_photo*area_frame_min: masks_selection.append(mask) # mask_binary_temp = mask['segmentation'].astype(int)*255 object_segmented = img_cropped.copy() object_segmented[mask_binary_temp == 0] = 255 objects_segmented.append(object_segmented) return image, img_cropped, masks_selection, objects_segmented def plot_sam_check_segmentation_frame(image, img_cropped, objects_segmented): # ax1 = plt.subplot(311) ax1.imshow(image) ax1.axis("Off") ax2 = plt.subplot(312) ax2.imshow(img_cropped) ax2.axis("Off") ax3 = plt.subplot(313) ax3.imshow(objects_segmented[0]) fig = plt.gcf() return fig def plot_sam_check_segmentation_frame_and_save(image, img_cropped, objects_segmented, output_folder, filepath): # ax1 = plt.subplot(311) ax1.imshow(image) ax1.axis("Off") ax2 = plt.subplot(312) ax2.imshow(img_cropped) ax2.axis("Off") ax3 = plt.subplot(313) ax3.imshow(objects_segmented[0]) # plt.savefig(output_folder + os.path.basename(filepath), transparent=True, bbox_inches='tight') # plt.show() def plot_segmented_object_with_bb(objects_segmented, masks_selection, image): image = cv2.rectangle(objects_segmented[0], masks_selection[0]['bbox'], [255,0,0], 4) plt.imshow(image) ax = plt.gca() # img_height = image.shape[0] img_width = image.shape[1] # major_ticks_height = np.arange(0, img_height, 100) minor_ticks_height = np.arange(0, img_height, 10) major_ticks_width = np.arange(0, img_width, 100) minor_ticks_width = np.arange(0, img_width, 10) # ax.set_yticks(major_ticks_height) ax.set_yticks(minor_ticks_height, minor=True) ax.set_xticks(major_ticks_width) ax.set_xticks(minor_ticks_width, minor=True) # ax.grid(which='both') ax.grid(which='minor', alpha=0.2) ax.grid(which='major', alpha=0.5) # # plt.show() def get_frame_width(masks_selection): frame_width = masks_selection[0]['bbox'][2] return frame_width def ipd_calibration(ipd_px, frame_width_px, frame_width_mm): calibration_factor = frame_width_px/frame_width_mm ipd_mm = ipd_px/calibration_factor return ipd_mm