Spaces:
Running
Running
File size: 2,303 Bytes
0aa7f8a d36b367 b95c986 d36b367 0aa7f8a |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from segmentation import predict as segmentation_predict
from depth_estimation import predict as depth_estimation_predict
def predict(image, color_map):
# inference
mask_image = segmentation_predict(image)
segmented_image = Image.composite(
image,
Image.new("RGB", image.size, (0, 0, 0)),
mask_image.convert("L")
)
depth_image = depth_estimation_predict(segmented_image)
# apply matplotlib colormap (e.g., viridis)
depth_array = np.array(depth_image) # Convert PIL image to NumPy array
colormap = plt.get_cmap(color_map) # Choose a colormap
depth_colored = colormap(depth_array / 255.0) # Normalize and apply colormap
depth_colored = (depth_colored * 255).astype(np.uint8) # Convert to RGB (discard alpha)
depth_colored = Image.fromarray(depth_colored)
return depth_colored
color_maps = [
'viridis', 'plasma', 'inferno', 'magma', 'cividis',
'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn',
'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone',
'pink', 'spring', 'summer', 'autumn', 'winter', 'cool',
'Wistia', 'hot', 'afmhot', 'gist_heat', 'copper',
'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu', 'RdYlBu',
'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic',
'twilight', 'twilight_shifted', 'hsv',
'Pastel1', 'Pastel2', 'Paired', 'Accent', 'Dark2',
'Set1', 'Set2', 'Set3', 'tab10', 'tab20', 'tab20b', 'tab20c',
'flag', 'prism', 'ocean', 'gist_earth', 'terrain',
'gist_stern', 'gnuplot', 'gnuplot2', 'CMRmap',
'cubehelix', 'brg', 'gist_rainbow', 'rainbow', 'jet',
'turbo', 'nipy_spectral', 'gist_ncar',
]
examples = [
["assets/examples/myself.jpeg", "afmhot"],
["assets/examples/myself.jpeg", "inferno"],
]
interface = gr.Interface(
fn=predict,
inputs=[
gr.Image(type="pil"),
gr.Dropdown(choices=color_maps),
],
outputs=gr.Image(type="pil"),
title="DepthPro: Colorify",
description="Applies segmentation on the input image, then creates the depth map and finally colorizes it.",
examples=examples,
)
if __name__ == "__main__":
interface.launch()
|