Update app.py

app.py

@@ -4,121 +4,104 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import tempfile
+from skimage import measure
 import trimesh
+import torch.nn.functional as F
+import torchvision.transforms as transforms
 
 # Check if CUDA is available, otherwise use CPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Import Point-E modules
-try:
-    print("Loading Point-E model...")
-    from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
-    from point_e.diffusion.sampler import PointCloudSampler
-    from point_e.models.configs import MODEL_CONFIGS, model_from_config
-    from point_e.models.download import load_checkpoint
-    from point_e.util.plotting import plot_point_cloud
-except ImportError:
-    print("Point-E modules not available. Please make sure Point-E is installed.")
-    raise
-
-# Create base model for image encoder
-base_name = 'base40M-textvec'
-base_model = model_from_config(MODEL_CONFIGS[base_name], device)
-base_model.eval()
-base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name])
-
-# Create upsampler model
-upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)
-upsampler_model.eval()
-upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])
-
-# Create image to point cloud model
-img2pc_name = 'base300M'
-img2pc_model = model_from_config(MODEL_CONFIGS[img2pc_name], device)
-img2pc_model.eval()
-img2pc_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[img2pc_name])
-
-# Load checkpoints
-print("Loading model checkpoints...")
-base_model.load_state_dict(load_checkpoint(base_name, device))
-upsampler_model.load_state_dict(load_checkpoint('upsample', device))
-img2pc_model.load_state_dict(load_checkpoint(img2pc_name, device))
-
-# Create samplers
-sampler = PointCloudSampler(
-    device=device,
-    models=[base_model, upsampler_model],
-    diffusions=[base_diffusion, upsampler_diffusion],
-    num_points=[1024, 4096],
-    aux_channels=['R', 'G', 'B'],
-    guidance_scale=[3.0, 0.0],
-)
+# Define a simple neural network to extract depth from images
+class SimpleDepthNet(torch.nn.Module):
+    def __init__(self):
+        super(SimpleDepthNet, self).__init__()
+        self.conv1 = torch.nn.Conv2d(3, 32, kernel_size=3, padding=1)
+        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.conv3 = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.conv4 = torch.nn.Conv2d(128, 1, kernel_size=3, padding=1)
+        self.pool = torch.nn.MaxPool2d(2, 2)
+        self.upsample = torch.nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        
+    def forward(self, x):
+        # Encoder
+        x = F.relu(self.conv1(x))
+        x = self.pool(x)
+        x = F.relu(self.conv2(x))
+        x = self.pool(x)
+        
+        # Decoder
+        x = self.upsample(x)
+        x = F.relu(self.conv3(x))
+        x = self.upsample(x)
+        x = torch.sigmoid(self.conv4(x))
+        return x
 
-img2pc_sampler = PointCloudSampler(
-    device=device,
-    models=[img2pc_model],
-    diffusions=[img2pc_diffusion],
-    num_points=[1024],
-    aux_channels=['R', 'G', 'B'],
-    guidance_scale=[3.0],
-)
+# Initialize the model
+model = SimpleDepthNet().to(device)
 
-def preprocess_image(image):
-    # Resize to match expected input size
-    image = image.resize((256, 256))
-    return image
+# Define transformation for input images
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+])
 
-def image_to_3d(image, num_steps=64):
+def image_to_3d(image):
     """
-    Convert a single image to a 3D model using Point-E
+    Convert a single image to a 3D model using a simple depth extraction approach
     """
     if image is None:
         return None, "No image provided"
     
     try:
         # Preprocess image
-        processed_image = preprocess_image(image)
+        img_tensor = transform(image).unsqueeze(0).to(device)
+        
+        # Generate depth map
+        with torch.no_grad():
+            depth = model(img_tensor)[0, 0].cpu().numpy()
+        
+        # Convert depth map to 3D points
+        h, w = depth.shape
+        y, x = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
         
-        # Generate samples
-        samples = None
-        for i, x in enumerate(img2pc_sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[processed_image]))):
-            samples = x
+        # Normalize coordinates
+        x = (x - w/2) / max(w, h)
+        y = (y - h/2) / max(w, h)
+        z = depth - 0.5  # Center around zero
         
-        # Extract point cloud
-        pc = samples[-1]['pred_pc']
-        colors = samples[-1]['pred_pc_aux']['R', 'G', 'B']
+        # Create point cloud
+        points = np.stack([x.flatten(), y.flatten(), z.flatten()], axis=1)
         
-        # Create colored point cloud
-        points = pc.cpu().numpy()[0]
-        colors_np = colors.cpu().numpy()[0]
+        # Get colors from original image
+        img_np = np.array(image.resize((w, h))) / 255.0
+        colors = img_np.reshape(-1, 3)
         
-        # Create a mesh from point cloud
-        point_cloud = trimesh.PointCloud(vertices=points, colors=colors_np)
+        # Create a mesh from the point cloud (using marching cubes on the depth map)
+        verts, faces, _, _ = measure.marching_cubes(depth, 0.5)
+        mesh = trimesh.Trimesh(vertices=verts, faces=faces)
         
         # Save as OBJ
         with tempfile.NamedTemporaryFile(suffix='.obj', delete=False) as obj_file:
             obj_path = obj_file.name
-            point_cloud.export(obj_path)
+            mesh.export(obj_path)
         
-        # Save as PLY for better Unity compatibility
+        # Also save as PLY for better compatibility with Unity
         with tempfile.NamedTemporaryFile(suffix='.ply', delete=False) as ply_file:
             ply_path = ply_file.name
-            point_cloud.export(ply_path)
+            mesh.export(ply_path)
         
         return [obj_path, ply_path], "3D model generated successfully!"
     except Exception as e:
         return None, f"Error: {str(e)}"
 
-def process_image(image, num_steps):
+def process_image(image):
     try:
         if image is None:
             return None, None, "Please upload an image first."
         
-        results, message = image_to_3d(
-            image, 
-            num_steps=num_steps
-        )
+        results, message = image_to_3d(image)
         
         if results:
             return results[0], results[1], message
@@ -128,14 +111,13 @@ def process_image(image, num_steps):
         return None, None, f"Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="Image to 3D Point Cloud Converter") as demo:
-    gr.Markdown("# Image to 3D Point Cloud Converter")
-    gr.Markdown("Upload an image to convert it to a 3D point cloud that you can use in Unity or other engines.")
+with gr.Blocks(title="Simple Image to 3D Converter") as demo:
+    gr.Markdown("# Simple Image to 3D Converter")
+    gr.Markdown("Upload an image to convert it to a simple 3D model that you can use in Unity or other engines.")
     
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(type="pil", label="Input Image")
-            num_steps = gr.Slider(minimum=16, maximum=128, value=64, step=8, label="Number of Inference Steps")
             submit_btn = gr.Button("Convert to 3D")
         
         with gr.Column(scale=1):
@@ -145,7 +127,7 @@ with gr.Blocks(title="Image to 3D Point Cloud Converter") as demo:
     
     submit_btn.click(
         fn=process_image,
-        inputs=[input_image, num_steps],
+        inputs=[input_image],
         outputs=[obj_file, ply_file, output_message]
     )