Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import cv2 | |
| import math | |
| import numpy as np | |
| import os | |
| import numpy as np | |
| import mediapipe as mp | |
| from mediapipe.tasks import python | |
| from mediapipe.tasks.python import vision | |
| print("hello world") | |
| # Height and width that will be used by the model | |
| DESIRED_HEIGHT = 480 | |
| DESIRED_WIDTH = 480 | |
| # Performs resizing and showing the image | |
| def resize_and_show(image): | |
| h, w = image.shape[:2] | |
| if h < w: | |
| img = cv2.resize(image, (DESIRED_WIDTH, math.floor(h/(w/DESIRED_WIDTH)))) | |
| else: | |
| img = cv2.resize(image, (math.floor(w/(h/DESIRED_HEIGHT)), DESIRED_HEIGHT)) | |
| cv2.imshow('color', img) | |
| cv2.waitKey(1000) | |
| cv2.destroyAllWindows() | |
| def segmentate(filepath): | |
| BG_COLOR = (192, 192, 192) # gray | |
| MASK_COLOR = (255, 255, 255) # white | |
| # Create the options that will be used for ImageSegmenter | |
| base_options = python.BaseOptions(model_asset_path='./hair_segmenter.tflite') | |
| options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True) | |
| # Create the image segmenter | |
| with vision.ImageSegmenter.create_from_options(options) as segmenter: | |
| # Loop through demo image(s) | |
| # Create the MediaPipe image file that will be segmented | |
| image = mp.Image.create_from_file(filepath) | |
| # Retrieve the masks for the segmented image | |
| segmentation_result = segmenter.segment(image) | |
| category_mask = segmentation_result.category_mask | |
| # 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_mask.numpy_view(),) * 3, axis=-1) > 0.2 | |
| output_image = np.where(condition, fg_image, bg_image) | |
| # print(f'Segmentation mask of {name}:') | |
| # resize_and_show(output_image) | |
| return output_image | |
| # GUI | |
| title = 'mediapipe hair segmentation' | |
| description = 'hair segmentation using mediapipe' | |
| examples = [[f'examples/{name}', 3] for name in sorted(os.listdir('examples'))] | |
| iface = gr.Interface( | |
| fn=segmentate, | |
| inputs=[ | |
| gr.Image(type='filepath', label='Input Image') | |
| ], | |
| outputs=[ | |
| gr.Image(label='image segmentated') | |
| ], | |
| examples=examples, | |
| allow_flagging='never', | |
| cache_examples=False, | |
| title=title, | |
| description=description | |
| ) | |
| iface.launch() |