|
from doctest import Example |
|
import gradio as gr |
|
from transformers import DPTFeatureExtractor, DPTForDepthEstimation |
|
import torch |
|
import numpy as np |
|
from PIL import Image, ImageOps |
|
from pathlib import Path |
|
import glob |
|
from autostereogram.converter import StereogramConverter |
|
from datetime import datetime |
|
import time |
|
|
|
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large") |
|
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") |
|
|
|
stereo_converter = StereogramConverter() |
|
|
|
def process_image(image_path): |
|
print("\n\n\n") |
|
print("Processing image:", image_path) |
|
last_time = time.time() |
|
image_raw = Image.open(Path(image_path)) |
|
|
|
image = image_raw.resize( |
|
(1280, int(1280 * image_raw.size[1] / image_raw.size[0])), |
|
Image.Resampling.LANCZOS) |
|
|
|
|
|
encoding = feature_extractor(image, return_tensors="pt") |
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model(**encoding) |
|
predicted_depth = outputs.predicted_depth |
|
|
|
|
|
prediction = torch.nn.functional.interpolate( |
|
predicted_depth.unsqueeze(1), |
|
size=image.size[::-1], |
|
mode="bicubic", |
|
align_corners=False, |
|
).squeeze() |
|
output = prediction.cpu().numpy() |
|
depth_image = (output * 255 / np.max(output)).astype('uint8') |
|
depth_image_padded = np.array(ImageOps.pad( |
|
Image.fromarray(depth_image), (1280, 720))) |
|
|
|
stereo_image = stereo_converter.convert_depth_to_stereogram_with_thread_pool( |
|
depth_image_padded, False).astype(np.uint8) |
|
|
|
stereo_image_pil = Image.fromarray(stereo_image).convert('RGB') |
|
image_name = f'stereo_image_{datetime.now().strftime("%Y%m%d_%H%M%S")}.jpg' |
|
stereo_image_pil.save(image_name) |
|
print(time.time() - last_time) |
|
print("\n\n\n") |
|
return [depth_image_padded, stereo_image, image_name] |
|
|
|
|
|
examples_images = [[f] for f in sorted(glob.glob('examples/*.jpg'))] |
|
|
|
blocks = gr.Blocks() |
|
|
|
input_image = gr.Image(type="filepath", label="Input Image") |
|
predicted_depth = gr.Image(label="Predicted Depth", type="pil") |
|
autostereogram = gr.Image(label="Autostereogram", type="pil") |
|
file_download = gr.File(label="Download Image") |
|
|
|
|
|
def load_example(example_id): |
|
processed_examples = [ |
|
component.preprocess_example(sample) |
|
for component, sample in zip( |
|
[input_image], examples_images[example_id] |
|
) |
|
] |
|
if len(processed_examples) == 1: |
|
return processed_examples[0] |
|
else: |
|
return processed_examples |
|
|
|
|
|
with blocks: |
|
gr.Markdown(''' |
|
## Depth Image to Autostereogram (Magic Eye) |
|
This demo is a variation from the original [DPT Demo](https://huggingface.co/spaces/nielsr/dpt-depth-estimation). |
|
Zero-shot depth estimation from an image, then it uses [pystereogram](https://github.com/yxiao1996/pystereogram) |
|
to generate the autostereogram (Magic Eye) |
|
<base target="_blank"> |
|
|
|
''') |
|
|
|
with gr.Row(): |
|
examples_c = gr.components.Dataset( |
|
components=[input_image], |
|
samples=examples_images, |
|
type="index", |
|
) |
|
|
|
examples_c.click( |
|
load_example, |
|
inputs=[examples_c], |
|
outputs=[input_image], |
|
_postprocess=False, |
|
queue=False, |
|
) |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image.render() |
|
button = gr.Button("Predict") |
|
button.click(fn=process_image, inputs=[input_image], |
|
outputs=[predicted_depth, |
|
autostereogram, file_download], |
|
) |
|
|
|
with gr.Column(): |
|
predicted_depth.render() |
|
with gr.Row(): |
|
autostereogram.render() |
|
with gr.Row(): |
|
with gr.Column(): |
|
file_download.render() |
|
|
|
if __name__ == "__main__": |
|
blocks.launch(debug=True) |
|
|
