v1

P3-SAM/utils/chamfer3D/chamfer3D.cu

@@ -1,196 +0,0 @@
-
-#include <stdio.h>
-#include <ATen/ATen.h>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#include <vector>
-
-
-
-__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){
-	const int batch=512;
-	__shared__ float buf[batch*3];
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int k2=0;k2<m;k2+=batch){
-			int end_k=min(m,k2+batch)-k2;
-			for (int j=threadIdx.x;j<end_k*3;j+=blockDim.x){
-				buf[j]=xyz2[(i*m+k2)*3+j];
-			}
-			__syncthreads();
-			for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-				float x1=xyz[(i*n+j)*3+0];
-				float y1=xyz[(i*n+j)*3+1];
-				float z1=xyz[(i*n+j)*3+2];
-				int best_i=0;
-				float best=0;
-				int end_ka=end_k-(end_k&3);
-				if (end_ka==batch){
-					for (int k=0;k<batch;k+=4){
-						{
-							float x2=buf[k*3+0]-x1;
-							float y2=buf[k*3+1]-y1;
-							float z2=buf[k*3+2]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*3+3]-x1;
-							float y2=buf[k*3+4]-y1;
-							float z2=buf[k*3+5]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*3+6]-x1;
-							float y2=buf[k*3+7]-y1;
-							float z2=buf[k*3+8]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*3+9]-x1;
-							float y2=buf[k*3+10]-y1;
-							float z2=buf[k*3+11]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}else{
-					for (int k=0;k<end_ka;k+=4){
-						{
-							float x2=buf[k*3+0]-x1;
-							float y2=buf[k*3+1]-y1;
-							float z2=buf[k*3+2]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (k==0 || d<best){
-								best=d;
-								best_i=k+k2;
-							}
-						}
-						{
-							float x2=buf[k*3+3]-x1;
-							float y2=buf[k*3+4]-y1;
-							float z2=buf[k*3+5]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+1;
-							}
-						}
-						{
-							float x2=buf[k*3+6]-x1;
-							float y2=buf[k*3+7]-y1;
-							float z2=buf[k*3+8]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+2;
-							}
-						}
-						{
-							float x2=buf[k*3+9]-x1;
-							float y2=buf[k*3+10]-y1;
-							float z2=buf[k*3+11]-z1;
-							float d=x2*x2+y2*y2+z2*z2;
-							if (d<best){
-								best=d;
-								best_i=k+k2+3;
-							}
-						}
-					}
-				}
-				for (int k=end_ka;k<end_k;k++){
-					float x2=buf[k*3+0]-x1;
-					float y2=buf[k*3+1]-y1;
-					float z2=buf[k*3+2]-z1;
-					float d=x2*x2+y2*y2+z2*z2;
-					if (k==0 || d<best){
-						best=d;
-						best_i=k+k2;
-					}
-				}
-				if (k2==0 || result[(i*n+j)]>best){
-					result[(i*n+j)]=best;
-					result_i[(i*n+j)]=best_i;
-				}
-			}
-			__syncthreads();
-		}
-	}
-}
-// int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){
-int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){
-
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-
-	NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, n, xyz1.data<float>(), m, xyz2.data<float>(), dist1.data<float>(), idx1.data<int>());
-	NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, m, xyz2.data<float>(), n, xyz1.data<float>(), dist2.data<float>(), idx2.data<int>());
-
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err));
-	    //THError("aborting");
-	    return 0;
-	  }
-	  return 1;
-
-
-}
-__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){
-	for (int i=blockIdx.x;i<b;i+=gridDim.x){
-		for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
-			float x1=xyz1[(i*n+j)*3+0];
-			float y1=xyz1[(i*n+j)*3+1];
-			float z1=xyz1[(i*n+j)*3+2];
-			int j2=idx1[i*n+j];
-			float x2=xyz2[(i*m+j2)*3+0];
-			float y2=xyz2[(i*m+j2)*3+1];
-			float z2=xyz2[(i*m+j2)*3+2];
-			float g=grad_dist1[i*n+j]*2;
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+0]),g*(x1-x2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+1]),g*(y1-y2));
-			atomicAdd(&(grad_xyz1[(i*n+j)*3+2]),g*(z1-z2));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+0]),-(g*(x1-x2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+1]),-(g*(y1-y2)));
-			atomicAdd(&(grad_xyz2[(i*m+j2)*3+2]),-(g*(z1-z2)));
-		}
-	}
-}
-// int chamfer_cuda_backward(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream){
-int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2){
-	// cudaMemset(grad_xyz1,0,b*n*3*4);
-	// cudaMemset(grad_xyz2,0,b*m*3*4);
-	
-	const auto batch_size = xyz1.size(0);
-	const auto n = xyz1.size(1); //num_points point cloud A
-	const auto m = xyz2.size(1); //num_points point cloud B
-
-	NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,n,xyz1.data<float>(),m,xyz2.data<float>(),graddist1.data<float>(),idx1.data<int>(),gradxyz1.data<float>(),gradxyz2.data<float>());
-	NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,m,xyz2.data<float>(),n,xyz1.data<float>(),graddist2.data<float>(),idx2.data<int>(),gradxyz2.data<float>(),gradxyz1.data<float>());
-	
-	cudaError_t err = cudaGetLastError();
-	  if (err != cudaSuccess) {
-	    printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
-	    //THError("aborting");
-	    return 0;
-	  }
-	  return 1;
-	
-}
-

P3-SAM/utils/chamfer3D/chamfer_cuda.cpp

@@ -1,29 +0,0 @@
-#include <torch/torch.h>
-#include <vector>
-
-/// TMP
-// #include "common.h"
-/// NOT TMP
-
-int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2,
-                         at::Tensor idx1, at::Tensor idx2);
-
-int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1,
-                          at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2,
-                          at::Tensor idx1, at::Tensor idx2);
-
-int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2,
-                    at::Tensor idx1, at::Tensor idx2) {
-  return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2);
-}
-
-int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2,
-                     at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) {
-
-  return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("forward", &chamfer_forward, "chamfer forward (CUDA)");
-  m.def("backward", &chamfer_backward, "chamfer backward (CUDA)");
-}

P3-SAM/utils/chamfer3D/dist_chamfer_3D.py

@@ -1,81 +0,0 @@
-from torch import nn
-from torch.autograd import Function
-import torch
-import importlib
-import os
-chamfer_found = importlib.find_loader("chamfer_3D") is not None
-if not chamfer_found:
-    ## Cool trick from https://github.com/chrdiller
-    print("Jitting Chamfer 3D")
-    cur_path = os.path.dirname(os.path.abspath(__file__))
-    build_path = cur_path.replace('chamfer3D', 'tmp')
-    os.makedirs(build_path, exist_ok=True)
-
-    from torch.utils.cpp_extension import load
-    chamfer_3D = load(name="chamfer_3D",
-          sources=[
-              "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]),
-              "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]),
-              ], build_directory=build_path)
-    print("Loaded JIT 3D CUDA chamfer distance")
-
-else:
-    import chamfer_3D
-    print("Loaded compiled 3D CUDA chamfer distance")
-
-
-# Chamfer's distance module @thibaultgroueix
-# GPU tensors only
-class chamfer_3DFunction(Function):
-    @staticmethod
-    def forward(ctx, xyz1, xyz2):
-        batchsize, n, dim = xyz1.size()
-        assert dim==3, "Wrong last dimension for the chamfer distance 's input! Check with .size()"
-        _, m, dim = xyz2.size()
-        assert dim==3, "Wrong last dimension for the chamfer distance 's input! Check with .size()"
-        device = xyz1.device
-
-        device = xyz1.device
-
-        dist1 = torch.zeros(batchsize, n)
-        dist2 = torch.zeros(batchsize, m)
-
-        idx1 = torch.zeros(batchsize, n).type(torch.IntTensor)
-        idx2 = torch.zeros(batchsize, m).type(torch.IntTensor)
-
-        dist1 = dist1.to(device)
-        dist2 = dist2.to(device)
-        idx1 = idx1.to(device)
-        idx2 = idx2.to(device)
-        torch.cuda.set_device(device)
-
-        chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
-        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
-        return dist1, dist2, idx1, idx2
-
-    @staticmethod
-    def backward(ctx, graddist1, graddist2, gradidx1, gradidx2):
-        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
-        graddist1 = graddist1.contiguous()
-        graddist2 = graddist2.contiguous()
-        device = graddist1.device
-
-        gradxyz1 = torch.zeros(xyz1.size())
-        gradxyz2 = torch.zeros(xyz2.size())
-
-        gradxyz1 = gradxyz1.to(device)
-        gradxyz2 = gradxyz2.to(device)
-        chamfer_3D.backward(
-            xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
-        )
-        return gradxyz1, gradxyz2
-
-
-class chamfer_3DDist(nn.Module):
-    def __init__(self):
-        super(chamfer_3DDist, self).__init__()
-
-    def forward(self, input1, input2):
-        input1 = input1.contiguous()
-        input2 = input2.contiguous()
-        return chamfer_3DFunction.apply(input1, input2)

P3-SAM/utils/chamfer3D/setup.py

@@ -1,14 +0,0 @@
-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='chamfer_3D',
-    ext_modules=[
-        CUDAExtension('chamfer_3D', [
-            "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']),
-            "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']),
-        ]),
-    ],
-    cmdclass={
-        'build_ext': BuildExtension
-    })

app.py

@@ -11,13 +11,11 @@ import spaces
 
 sys.path.append('P3-SAM')
 from demo.auto_mask import AutoMask
-# from demo.auto_mask_no_postprocess import AutoMask as AutoMaskNoPostProcess
 sys.path.append('XPart')
 from partgen.partformer_pipeline import PartFormerPipeline
 from partgen.utils.misc import get_config_from_file
 
 automask = AutoMask()
-# automask_no_postprocess = AutoMaskNoPostProcess(automask_instance=automask)
 
 def _load_pipeline():
     pl.seed_everything(2026, workers=True)
@@ -41,17 +39,21 @@ _PIPELINE = _load_pipeline()
 output_path = 'P3-SAM/results/gradio'
 os.makedirs(output_path, exist_ok=True)
 
+def is_supported_3d_file(filename):
+    # 获取文件扩展名（小写），并去除开头的点
+    ext = os.path.splitext(filename)[1].lower()
+    return ext in ['.glb', '.ply', '.obj']
+
 @spaces.GPU
-def segment(mesh_path, connectivity=True, postprocess=True, postprocess_threshold=0.95, seed=42):
+def segment(mesh_path, postprocess=True, postprocess_threshold=0.95, seed=42):
     if mesh_path is None:
         gr.Warning("No Input Mesh")
-        gr_state[0] = (None, None)
+        return None, None
+    if not is_supported_3d_file(mesh_path):
+        gr.Warning("Only support glb ply obj.")
         return None, None
     mesh = trimesh.load(mesh_path, force='mesh', process=False)
-    if connectivity:
-        aabb, face_ids, mesh = automask.predict_aabb(mesh, seed=seed, is_parallel=False, post_process=postprocess, threshold=postprocess_threshold)
-    else:
-        aabb, face_ids, mesh = automask_no_postprocess.predict_aabb(mesh, seed=seed, is_parallel=False, post_process=False)
+    aabb, face_ids, mesh = automask.predict_aabb(mesh, seed=seed, is_parallel=False, post_process=postprocess, threshold=postprocess_threshold)
     color_map = {}
     unique_ids = np.unique(face_ids)
     for i in unique_ids:
@@ -76,7 +78,7 @@ def segment(mesh_path, connectivity=True, postprocess=True, postprocess_threshol
     gr_state = [(aabb, mesh_path)]
     return file_path, face_id_save_path, gr_state
 
-@spaces.GPU
+@spaces.GPU(duration=150)
 def generate(mesh_path, seed=42, gr_state=None):
     if mesh_path is None:
         gr.Warning("No Input Mesh")
@@ -114,6 +116,7 @@ with gr.Blocks() as demo:
 # ☯️ Hunyuan3D Part:P3-SAM&XPart
 This demo allows you to generate parts given a 3D model using Hunyuan3D-Part.
 First segment the 3D model using P3-SAM and then generate parts using XPart.
+Please upload glb ply or obj 3D model files.
 '''
     )
     with gr.Row():
@@ -133,7 +136,6 @@ Input a mesh and push the "Segment" button to get the segmentation results.
             p3sam_input = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Input Mesh")
             p3sam_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Segmentation Result")
             p3sam_face_id_output = gr.File(label='Face ID')
-            p3sam_conectivity = gr.Checkbox(value=True, label="Connectivity")
             p3sam_postprocess = gr.Checkbox(value=True, label="Post-processing")
             p3sam_postprocess_threshold = gr.Number(value=0.95, label="Post-processing Threshold")
             p3sam_seed = gr.Number(value=42, label="Random Seed")
@@ -183,7 +185,7 @@ Input a mesh, segment it using P3-SAM on the left, and push the "Generate" butto
             xpart_output_exploded = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Exploded Object")
             xpart_seed = gr.Number(value=42, label="Random Seed")
         gr_state = gr.State(value=[(None, None)])
-    p3sam_button.click(segment, inputs=[p3sam_input, p3sam_conectivity, p3sam_postprocess, p3sam_postprocess_threshold, p3sam_seed], outputs=[p3sam_output, p3sam_face_id_output, gr_state])
+    p3sam_button.click(segment, inputs=[p3sam_input, p3sam_postprocess, p3sam_postprocess_threshold, p3sam_seed], outputs=[p3sam_output, p3sam_face_id_output, gr_state])
     xpart_button.click(generate, inputs=[p3sam_input, xpart_seed, gr_state], outputs=[xpart_output, xpart_output_bbox, xpart_output_exploded])