Update app_fullbody_pretrained.py

app_fullbody_pretrained.py

@@ -663,15 +663,16 @@ async def generate_3d_local(file: UploadFile = File(...)):
 # ============================================================
 # ✅ MONAI + PyVista 3D Full-Body Visualization (Final HF Safe)
 # ============================================================
+
 from fastapi import UploadFile, File
 from fastapi.responses import JSONResponse
-import io, os, numpy as np, torch, imageio
+import io, os, numpy as np, torch
 from PIL import Image
 from monai.networks.nets import UNet
 import torch.nn.functional as F
-import pyvista
-from pyvista import Plotter, start_xvfb
-import vtk
+import pyvista as pv
+from pyvista import start_xvfb
+import imageio.v3 as iio
 
 HF_SPACE_URL = os.getenv("HF_SPACE_URL", "https://abbhy123ghh-x-ray-analysis.hf.space").rstrip("/")
 
@@ -695,9 +696,7 @@ async def generate_3d_monai(file: UploadFile = File(...)):
         volume = np.expand_dims(volume, axis=0)  # [1, D, H, W]
 
         # --- Step 3: MONAI lightweight 3D UNet (shallow for CPU) ---
-        x = torch.tensor(volume, dtype=torch.float32).unsqueeze(0)  # [1,1,D,H,W]
-
-        # pad to multiples of 16 for shape-safe inference
+        x = torch.tensor(volume, dtype=torch.float32).unsqueeze(0)  # [1, 1, D, H, W]
         pad_d = (16 - x.shape[2] % 16) % 16
         pad_h = (16 - x.shape[3] % 16) % 16
         pad_w = (16 - x.shape[4] % 16) % 16
@@ -719,54 +718,39 @@ async def generate_3d_monai(file: UploadFile = File(...)):
         seg = (seg - seg.min()) / (seg.max() - seg.min() + 1e-8)
         seg = np.clip(seg * 255, 0, 255).astype(np.uint8)
 
-        # --- Step 4: Create volumetric grid safely (ImageData, no UniformGrid) ---
-        try:
-            grid = pyvista.ImageData()
-        except Exception:
-            grid = vtk.vtkImageData()
-
-        dims = np.array(seg.shape) + 1
-        try:
-            grid.dimensions = dims
-        except Exception:
-            grid.SetDimensions(int(dims[2]), int(dims[1]), int(dims[0]))
-
+        # --- Step 4: Create PyVista grid (Hugging Face compatible) ---
+        start_xvfb()
+        grid = pv.ImageData(dimensions=seg.shape)
         grid.spacing = (1, 1, 1)
         grid.origin = (0, 0, 0)
+        grid.point_data["values"] = seg.flatten(order="F")
 
-        try:
-            grid.cell_data["values"] = seg.flatten(order="F")
-        except Exception:
-            # fallback for vtk.vtkImageData
-            arr_flat = seg.flatten(order="F")
-            for i, val in enumerate(arr_flat):
-                grid.GetPointData().GetScalars().SetTuple1(i, int(val))
-
-        pv_grid = pyvista.wrap(grid)
-
-        # --- Step 5: Render rotating 3D volume (headless safe mode) ---
-        try:
-            start_xvfb()  # enable virtual framebuffer if available
-        except OSError:
-            print("⚠️ Xvfb not available — running headless fallback mode.")
-
-        os.environ["PYVISTA_OFF_SCREEN"] = "true"
-
-        plotter = Plotter(off_screen=True)
+        # --- Step 5: Enhanced 3D visualization with lighting ---
+        plotter = pv.Plotter(off_screen=True, window_size=[512, 512])
         plotter.add_volume(
-            pv_grid,
+            grid,
             cmap="bone",
             opacity="sigmoid_5",
             shade=True,
+            diffuse=1.0,
+            specular=0.3,
+            specular_power=15,
         )
         plotter.set_background("black")
-        plotter.enable_eye_dome_lighting()
+
+        # Add realistic light
+        light = pv.Light(
+            position=(seg.shape[2]*2, seg.shape[1]*2, seg.shape[0]),
+            color='white',
+            intensity=1.5,
+        )
+        plotter.add_light(light)
 
         frames = []
         for angle in range(0, 360, 5):
             plotter.camera_position = [
-                (seg.shape[2] * 2, 0, seg.shape[0] / 2),
-                (seg.shape[2] / 2, seg.shape[1] / 2, seg.shape[0] / 2),
+                (seg.shape[2]*2, 0, seg.shape[0]/2),
+                (seg.shape[2]/2, seg.shape[1]/2, seg.shape[0]/2),
                 (0, 0, 1),
             ]
             plotter.camera.azimuth = angle
@@ -774,14 +758,10 @@ async def generate_3d_monai(file: UploadFile = File(...)):
             frames.append(img_frame)
         plotter.close()
 
-        # --- Step 6: Save rotating video ---
+        # --- Step 6: Save rotating MP4 video ---
         os.makedirs("static", exist_ok=True)
         out_path = "static/xray_monai_3d_fullbody.mp4"
-
-        writer = imageio.get_writer(out_path, fps=12, codec="libx264", quality=8)
-        for f in frames:
-            writer.append_data(f)
-        writer.close()
+        iio.imwrite(out_path, frames, fps=12, codec="libx264")
 
         return JSONResponse({
             "ok": True,
@@ -793,6 +773,7 @@ async def generate_3d_monai(file: UploadFile = File(...)):
         import traceback; traceback.print_exc()
         return JSONResponse({"error": str(e)}, status_code=500)
 
+