Spaces:
Running
Running
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| examples = [ | |
| '176788.jpg', | |
| '182862.jpg' | |
| ] | |
| masks = [ | |
| 'mask1.jpg', | |
| 'mask2.jpg' | |
| ] | |
| def edge_detection(mode_selector, upload_input, image_selector, image_selected, threshold1, threshold2): | |
| if mode_selector == "Upload": | |
| image = upload_input | |
| else: | |
| image = image_selected | |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| edges = cv2.Canny(gray, threshold1, threshold2) | |
| return edges | |
| def mode_selector_changed(mode_selector, image_selector): | |
| if mode_selector == "Upload": | |
| return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False) | |
| else: | |
| return gr.update(visible=False), gr.update(visible=True), gr.update(visible=image_selector is not None, value=image_selector) | |
| def mode_selector_changed_inpainting(mode_selector, image_selector, mask_selector): | |
| if mode_selector == "Upload": | |
| return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) | |
| else: | |
| return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=image_selector is not None, value=image_selector), gr.update(visible=True), gr.update(visible=image_selector is not None, value=mask_selector) | |
| def close_selected_image(image_selected): | |
| if image_selected is None: | |
| return gr.update(visible=True, value=None), gr.update(visible=False, value=None) | |
| else: | |
| return gr.update(visible=True), gr.update(visible=True) | |
| def close_selected_mask(mask_selected): | |
| if mask_selected is None: | |
| return gr.update(visible=True, value=None), gr.update(visible=False, value=None) | |
| else: | |
| return gr.update(visible=True), gr.update(visible=True) | |
| def image_segmentation(mode_selector, upload_input, image_selector, image_selected): | |
| if mode_selector == "Upload": | |
| image = upload_input | |
| else: | |
| image = image_selected | |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| _, segmented = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| return segmented | |
| def image_inpainting(mode_selector, upload_input, upload_mask, image_selector, image_selected, mask_selector, mask_selected): | |
| if mode_selector == "Upload": | |
| image = upload_input | |
| mask = upload_mask | |
| else: | |
| image = image_selected | |
| mask = mask_selected | |
| if mask.shape[:2] != image.shape[:2]: | |
| mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST) | |
| if len(mask.shape) == 3: | |
| mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY) | |
| mask = (mask > 0).astype(np.uint8) * 255 | |
| inpainted = cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA) | |
| return inpainted | |
| with gr.Blocks() as demo: | |
| with gr.Tab("Edge Detection"): | |
| mode_selector = gr.Radio(["Upload", "Example"], value="Upload", label="Upload an image or using example image") | |
| # Mode 1 | |
| upload_input = gr.Image(type="numpy", label="Upload Image", visible=True) | |
| # Mode 2 | |
| image_selector = gr.Radio(examples, label="Select an example image", visible=False) | |
| image_selected = gr.Image(None, label="Selected Image", visible=False, sources=[]) | |
| threshold1 = gr.Slider(0, 255, value=50, step=1, label="Threshold 1") | |
| threshold2 = gr.Slider(0, 255, value=100, step=1, label="Threshold 2") | |
| output = gr.Image(type="numpy", label="Edge Image") | |
| inputs = [ | |
| mode_selector, | |
| upload_input, | |
| image_selector, | |
| image_selected, | |
| threshold1, | |
| threshold2 | |
| ] | |
| mode_selector.change( | |
| mode_selector_changed, | |
| inputs=[mode_selector, image_selector], | |
| outputs=[upload_input, image_selector, image_selected] | |
| ) | |
| image_selector.change( | |
| mode_selector_changed, | |
| inputs=[mode_selector, image_selector], | |
| outputs=[upload_input, image_selector, image_selected] | |
| ) | |
| image_selected.change( | |
| close_selected_image, | |
| inputs=[image_selected], | |
| outputs=[image_selector, image_selected] | |
| ) | |
| gr.Interface( | |
| fn=edge_detection, | |
| inputs=inputs, | |
| outputs=output, | |
| description="Upload an image and adjust the thresholds to perform edge detection", | |
| flagging_mode='never' | |
| ) | |
| with gr.Tab("Image Segmentation"): | |
| mode_selector = gr.Radio(["Upload", "Example"], value="Upload", label="Upload an image or using example image") | |
| # Mode 1 | |
| upload_input = gr.Image(type="numpy", label="Upload Image", visible=True) | |
| # Mode 2 | |
| image_selector = gr.Radio(examples, label="Select an example image", visible=False) | |
| image_selected = gr.Image(None, label="Selected Image", visible=False, sources=[]) | |
| output = gr.Image(type="numpy", label="Edge Image") | |
| inputs = [ | |
| mode_selector, | |
| upload_input, | |
| image_selector, | |
| image_selected, | |
| ] | |
| mode_selector.change( | |
| mode_selector_changed, | |
| inputs=[mode_selector, image_selector], | |
| outputs=[upload_input, image_selector, image_selected] | |
| ) | |
| image_selector.change( | |
| mode_selector_changed, | |
| inputs=[mode_selector, image_selector], | |
| outputs=[upload_input, image_selector, image_selected] | |
| ) | |
| image_selected.change( | |
| close_selected_image, | |
| inputs=[image_selected], | |
| outputs=[image_selector, image_selected] | |
| ) | |
| gr.Interface( | |
| fn=image_segmentation, | |
| inputs=inputs, | |
| outputs=output, | |
| description="Upload an image and adjust the thresholds to perform edge detection", | |
| flagging_mode='never' | |
| ) | |
| with gr.Tab("Image Inpainting"): | |
| mode_selector = gr.Radio(["Upload", "Example"], value="Upload", label="Upload an image or using example image") | |
| # Mode 1 | |
| upload_input = gr.Image(type="numpy", label="Upload Image", visible=True) | |
| upload_mask = gr.Image(type="numpy", label="Upload Mask", visible=True) | |
| # Mode 2 | |
| image_selector = gr.Radio(examples, label="Select an example image", visible=False) | |
| image_selected = gr.Image(None, label="Selected Image", visible=False, sources=[]) | |
| mask_selector = gr.Radio(masks, label="Select an example mask", visible=False) | |
| mask_selected = gr.Image(None, label="Selected Mask", visible=False, sources=[]) | |
| output = gr.Image(type="numpy", label="Edge Image") | |
| inputs = [ | |
| mode_selector, | |
| upload_input, | |
| upload_mask, | |
| image_selector, | |
| image_selected, | |
| mask_selector, | |
| mask_selected | |
| ] | |
| mode_selector.change( | |
| mode_selector_changed_inpainting, | |
| inputs=[mode_selector, image_selector, mask_selector], | |
| outputs=[upload_input, upload_mask, image_selector, image_selected, mask_selector, mask_selected] | |
| ) | |
| image_selector.change( | |
| mode_selector_changed, | |
| inputs=[mode_selector, image_selector], | |
| outputs=[upload_input, image_selector, image_selected] | |
| ) | |
| mask_selector.change( | |
| mode_selector_changed, | |
| inputs=[mode_selector, mask_selector], | |
| outputs=[upload_mask, mask_selector, mask_selected] | |
| ) | |
| image_selected.change( | |
| close_selected_image, | |
| inputs=[image_selected], | |
| outputs=[image_selector, image_selected] | |
| ) | |
| mask_selected.change( | |
| close_selected_mask, | |
| inputs=[mask_selected], | |
| outputs=[mask_selector, mask_selected] | |
| ) | |
| gr.Interface( | |
| fn=image_inpainting, | |
| inputs=inputs, | |
| outputs=output, | |
| description="Upload an image and adjust the thresholds to perform edge detection", | |
| flagging_mode='never' | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(allowed_paths=["/home/teng_aicv/Desktop/"]) |