3d / app.py
nithinm19's picture
Update app.py
1038ffe verified
raw
history blame
2.83 kB
import gradio as gr
from transformers import AutoModel, AutoProcessor
from PIL import Image
import torch
import requests
from io import BytesIO
import trimesh
import plotly.graph_objects as go
# Load model and processor from Hugging Face
def load_model_and_processor():
try:
model = AutoModel.from_pretrained("zxhezexin/openlrm-mix-large-1.1")
processor = AutoProcessor.from_pretrained("zxhezexin/openlrm-mix-large-1.1")
return model, processor
except Exception as e:
print(f"Error loading model or processor: {e}")
return None, None
model, processor = load_model_and_processor()
# Example image URL (replace this with a suitable example)
example_image_url = "https://huggingface.co/datasets/nateraw/image-folder/resolve/main/example_1.png"
# Function to load example image from URL
def load_example_image():
try:
response = requests.get(example_image_url)
image = Image.open(BytesIO(response.content))
return image
except Exception as e:
print(f"Error loading example image: {e}")
return None
# Define function to generate 3D output from 2D image
def image_to_3d(image):
if processor is None or model is None:
return "Model or processor not loaded."
try:
# Preprocess input image
inputs = processor(images=image, return_tensors="pt")
# Run inference
with torch.no_grad():
outputs = model(**inputs)
# Convert outputs to a 3D mesh (replace with actual logic based on model output)
# Assuming 'vertices' and 'faces' are returned by the model (adjust as needed)
vertices = outputs['vertices'].numpy() # Placeholder for vertex output
faces = outputs['faces'].numpy() # Placeholder for face output
# Create a mesh using trimesh
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
# Visualize the mesh using Plotly
fig = go.Figure(data=[go.Mesh3d(x=vertices[:,0], y=vertices[:,1], z=vertices[:,2],
i=faces[:,0], j=faces[:,1], k=faces[:,2])])
return fig # return the figure for display
except Exception as e:
return f"Error during inference: {str(e)}"
# Load the example image for the Gradio interface
example_image = load_example_image()
# Gradio interface setup
interface = gr.Interface(
fn=image_to_3d,
inputs=gr.Image(type="pil", label="Upload an Image"),
outputs=gr.Plot(label="3D Model"),
title="OpenLRM Mix-Large 1.1 - Image to 3D",
description="Upload an image to generate a 3D model using OpenLRM Mix-Large 1.1.",
examples=[[example_image]] if example_image else None,
theme="compact"
)
# Display a suggestion below the upload widget
interface.launch()