Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| from transformers import AutoModelForImageClassification, AutoFeatureExtractor | |
| # Load the Marigold model | |
| feature_extractor = AutoFeatureExtractor.from_pretrained("prs-eth/marigold-depth-lcm-v1-0") | |
| model = AutoModelForImageClassification.from_pretrained("prs-eth/marigold-depth-lcm-v1-0") | |
| def load_image(image): | |
| """ Convert uploaded image to grayscale. """ | |
| return image.convert('L') | |
| def compute_gradients(image): | |
| """ Compute horizontal and vertical gradients of the image. """ | |
| image_array = np.asarray(image, dtype=float) | |
| x_gradient = np.gradient(image_array, axis=1) | |
| y_gradient = np.gradient(image_array, axis=0) | |
| return x_gradient, y_gradient | |
| def create_normal_map(image): | |
| """ Generate a normal map from an image. """ | |
| image = load_image(image) | |
| x_grad, y_grad = compute_gradients(image) | |
| # Normalize gradients | |
| max_grad = max(np.max(np.abs(x_grad)), np.max(np.abs(y_grad))) | |
| x_grad /= max_grad | |
| y_grad /= max_grad | |
| # Calculate z component of the normal (assumed perpendicular to the surface) | |
| z = np.sqrt(1 - (x_grad ** 2) - (y_grad ** 2)) | |
| # Normalize to 0-255 and format as uint8 | |
| normal_map = np.dstack(((x_grad * 0.5 + 0.5) * 255, | |
| (y_grad * 0.5 + 0.5) * 255, | |
| (z * 1.0) * 255)).astype(np.uint8) | |
| return Image.fromarray(normal_map, 'RGB') | |
| def estimate_depth(image): | |
| """ Estimate depth using the Marigold model. """ | |
| inputs = feature_extractor(images=image, return_tensors="pt") | |
| outputs = model(**inputs) | |
| depth_map = outputs.logits # Adjust if the output tensor is named differently | |
| return depth_map | |
| def interface(image): | |
| normal_map = create_normal_map(image) | |
| grayscale_image = load_image(normal_map) | |
| depth_map = estimate_depth(image) | |
| return normal_map, grayscale_image, depth_map | |
| # Set up the Gradio interface | |
| iface = gr.Interface( | |
| fn=interface, | |
| inputs=gr.Image(type="pil", label="Upload Image"), | |
| outputs=[ | |
| gr.Image(type="pil", label="Normal Map"), | |
| gr.Image(type="pil", label="Grayscale Image"), | |
| gr.Image(type="pil", label="Depth Map") # Adjust the output type if needed | |
| ], | |
| title="Normal Map, Grayscale, and Depth Map Generator", | |
| description="Upload an image to generate its normal map, a grayscale version, and estimate its depth." | |
| ) | |
| iface.launch() | |