add

app.py

@@ -12,8 +12,9 @@ from einops import rearrange
 from shap_e.diffusion.sample import sample_latents
 from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
 from shap_e.models.download import load_model, load_config
-from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, create_custom_cameras
+from shap_e.util.notebooks import create_pan_cameras, decode_latent_images
 import spaces
+from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
 
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
@@ -25,6 +26,56 @@ from src.utils.camera_util import (
 from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground
 
+def create_custom_cameras(size: int, device: torch.device, azimuths: list, elevations: list, 
+                          fov_degrees: float,distance) -> DifferentiableCameraBatch:
+    # Object is in a 2x2x2 bounding box (-1 to 1 in each dimension)
+    object_diagonal =  distance # Correct diagonal calculation for the cube
+    
+    # Calculate radius based on object size and FOV
+    fov_radians = math.radians(fov_degrees)
+    radius = (object_diagonal / 2) / math.tan(fov_radians / 2)  # Correct radius calculation
+    # print(radius)
+    # exit(0)
+    origins = []
+    xs = []
+    ys = []
+    zs = []
+    
+    for azimuth, elevation in zip(azimuths, elevations):
+        azimuth_rad = np.radians(azimuth-90)
+        elevation_rad = np.radians(elevation)
+        
+        # Calculate camera position
+        x = radius * np.cos(elevation_rad) * np.cos(azimuth_rad)
+        y = radius * np.cos(elevation_rad) * np.sin(azimuth_rad)
+        z = radius * np.sin(elevation_rad)
+        origin = np.array([x, y, z])
+        
+        # Calculate camera orientation
+        z_axis = -origin / np.linalg.norm(origin)  # Point towards center
+        x_axis = np.array([-np.sin(azimuth_rad), np.cos(azimuth_rad), 0])
+        y_axis = np.cross(z_axis, x_axis)
+        
+        origins.append(origin)
+        zs.append(z_axis)
+        xs.append(x_axis)
+        ys.append(y_axis)
+
+    return DifferentiableCameraBatch(
+        shape=(1, len(origins)),
+        flat_camera=DifferentiableProjectiveCamera(
+            origin=torch.from_numpy(np.stack(origins, axis=0)).float().to(device),
+            x=torch.from_numpy(np.stack(xs, axis=0)).float().to(device),
+            y=torch.from_numpy(np.stack(ys, axis=0)).float().to(device),
+            z=torch.from_numpy(np.stack(zs, axis=0)).float().to(device),
+            width=size,
+            height=size,
+            x_fov=fov_radians,
+            y_fov=fov_radians,
+        ),
+    )
+
+
 @spaces.GPU(duration=60)
 def load_models():
     """Initialize and load all required models"""