glb support

app.py

@@ -6,146 +6,72 @@ import plotly.graph_objects as go
 from torch_geometric.utils import to_dense_batch
 from model import HierarchicalFPSCliffordNet
 
-# --- SAF PYTORCH C++ YEDEKLERİ (Hugging Face için Özel Çözüm) ---
-def knn_pure(x, y, k):
+# --- SAF PYTORCH C++ YEDEKLERİ ---
+def knn_pure(x, y, k=3):
     dist = torch.cdist(x, y)
     _, topk_idx = torch.topk(dist, k, dim=1, largest=False)
-    col = torch.arange(x.size(0)).view(-1, 1).expand(-1, k).reshape(-1)
-    row = topk_idx.reshape(-1)
-    return torch.stack([row, col], dim=0)
+    return torch.stack([torch.arange(x.size(0)).view(-1, 1).expand(-1, k).reshape(-1), topk_idx.reshape(-1)], dim=0)
 
-def radius_graph_pure(pos, r, max_num_neighbors=32):
+def radius_graph_pure(pos, r, max_n=16):
     dist = torch.cdist(pos, pos)
-    n = pos.size(0)
-    k = min(max_num_neighbors, n)
-    _, target = torch.topk(dist, k, dim=1, largest=False)
-    source = torch.arange(n).view(-1, 1).expand(-1, k)
-    valid = dist[source, target] < r
-    return torch.stack([source[valid], target[valid]], dim=0)
+    _, target = torch.topk(dist, min(max_n, pos.size(0)), dim=1, largest=False)
+    source = torch.arange(pos.size(0)).view(-1, 1).expand(-1, min(max_n, pos.size(0)))
+    mask = dist[source, target] < r
+    return torch.stack([source[mask], target[mask]], dim=0)
 
 def fps_pure(pos, ratio=0.5):
-    n = pos.size(0)
-    k = max(1, int(n * ratio))
-    idx = torch.zeros(k, dtype=torch.long, device=pos.device)
-    dist = torch.full((n,), 1e10, device=pos.device)
+    k = max(1, int(pos.size(0) * ratio))
+    idx = torch.zeros(k, dtype=torch.long)
+    dist = torch.full((pos.size(0),), 1e10)
     farthest = 0
     for i in range(k):
         idx[i] = farthest
-        farthest_point = pos[farthest].view(1, 3)
-        dist_to_farthest = torch.cdist(pos, farthest_point).squeeze(-1)
-        dist = torch.min(dist, dist_to_farthest)
+        dist = torch.min(dist, torch.cdist(pos, pos[farthest].view(1, 3)).squeeze())
         farthest = torch.argmax(dist).item()
     return idx
 
-# --- Kategori Sözlükleri ---
-CATEGORIES = {
-    'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4, 'Earphone': 5, 'Guitar': 6,
-    'Knife': 7, 'Lamp': 8, 'Laptop': 9, 'Motorbike': 10, 'Mug': 11, 'Pistol': 12,
-    'Rocket': 13, 'Skateboard': 14, 'Table': 15
-}
-SEG_CLASSES = {
-    'Airplane': [0, 1, 2, 3], 'Bag': [4, 5], 'Cap': [6, 7], 'Car': [8, 9, 10, 11],
-    'Chair': [12, 13, 14, 15], 'Earphone': [16, 17, 18], 'Guitar': [19, 20, 21],
-    'Knife': [22, 23], 'Lamp': [24, 25, 26, 27], 'Laptop': [28, 29],
-    'Motorbike': [30, 31, 32, 33, 34, 35], 'Mug': [36, 37], 'Pistol': [38, 39, 40],
-    'Rocket': [41, 42, 43], 'Skateboard': [44, 45, 46], 'Table': [47, 48, 49]
-}
+# --- Ayarlar ---
+CATEGORIES = {'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4, 'Earphone': 5, 'Guitar': 6, 'Knife': 7, 'Lamp': 8, 'Laptop': 9, 'Motorbike': 10, 'Mug': 11, 'Pistol': 12, 'Rocket': 13, 'Skateboard': 14, 'Table': 15}
+SEG_CLASSES = {'Airplane': list(range(0,4)), 'Chair': list(range(12,16)), 'Guitar': list(range(19,22)), 'Laptop': list(range(28,30))} # Örnek kısıtlı liste
 
-# --- Modeli Yükleme ---
 device = torch.device("cpu")
 model = HierarchicalFPSCliffordNet(base_channels=12).to(device)
-
 try:
     checkpoint = torch.load("best_all_categories_clifford.pt", map_location=device, weights_only=True)
-    clean_sd = {k.replace("_orig_mod.", ""): v for k, v in checkpoint.items()}
-    model.load_state_dict(clean_sd, strict=False)
+    model.load_state_dict({k.replace("_orig_mod.", ""): v for k, v in checkpoint.items()}, strict=False)
     model.eval()
-    print("Model ağırlıkları başarıyla yüklendi!")
-except Exception as e:
-    print("Ağırlık yükleme hatası (Demo rastgele ağırlıklarla çalışacak):", e)
-
-# --- Veri Hazırlama Fonksiyonu ---
-def process_point_cloud(file_path, category_name):
-    mesh = trimesh.load(file_path, force='mesh')
-    points = np.array(mesh.vertices)
-    
-    if len(points) > 1024:
-        idx = np.random.choice(len(points), 1024, replace=False)
-        points = points[idx]
-    elif len(points) < 1024:
-        idx = np.random.choice(len(points), 1024, replace=True)
-        points = points[idx]
-        
-    pos = torch.tensor(points, dtype=torch.float32)
-    pos = pos - pos.mean(dim=0, keepdim=True)
-    scale = pos.norm(dim=1).max().clamp(min=1e-8)
-    pos = pos / scale
-    
-    batch = torch.zeros(pos.size(0), dtype=torch.long)
-    cat_idx = torch.tensor([CATEGORIES[category_name]], dtype=torch.long)
-    
-    # Saf PyTorch ile Hızlı Hesaplamalar (C++ Yok)
-    fps_idx_1 = fps_pure(pos, ratio=0.5)
-    pos2 = pos[fps_idx_1]
-    batch2 = torch.zeros(pos2.size(0), dtype=torch.long)
-    
-    fps_idx_2 = fps_pure(pos2, ratio=0.25)
-    pos3 = pos2[fps_idx_2]
-    batch3 = torch.zeros(pos3.size(0), dtype=torch.long)
-
-    edge_index_1 = radius_graph_pure(pos, r=0.15, max_num_neighbors=16)
-    edge_index_2 = radius_graph_pure(pos2, r=0.30, max_num_neighbors=32)
-
-    assign_index_32 = knn_pure(pos3, pos2, k=3)
-    assign_index_21 = knn_pure(pos2, pos, k=3)
+except: print("Ağırlıklar bulunamadı.")
 
-    dummy_x = torch.zeros(pos3.size(0), 1)
-    _, x_dense_mask = to_dense_batch(dummy_x, batch3, max_num_nodes=int(torch.bincount(batch3).max().item()))
-    
-    return pos, batch, cat_idx, fps_idx_1, fps_idx_2, pos2, batch2, pos3, batch3, edge_index_1, edge_index_2, assign_index_32, assign_index_21, x_dense_mask
-
-# --- Çıkarım Fonksiyonu ---
 def predict(file, category_name):
-    if file is None:
-        return None
-        
-    (pos, batch, cat_idx, fps_idx_1, fps_idx_2, pos2, batch2, pos3, batch3, 
-     edge_index_1, edge_index_2, assign_index_32, assign_index_21, x_dense_mask) = process_point_cloud(file.name, category_name)
+    if not file: return None
+    mesh = trimesh.load(file.name, force='mesh')
+    p = np.array(mesh.vertices)
+    p = p[np.random.choice(len(p), 1024, replace=len(p)<1024)]
+    pos = torch.tensor(p, dtype=torch.float32)
+    pos = (pos - pos.mean(0)) / pos.norm(dim=1).max().clamp(1e-8)
     
+    # İşleme
+    f1 = fps_pure(pos, 0.5); p2 = pos[f1]
+    f2 = fps_pure(p2, 0.25); p3 = p2[f2]
+    e1 = radius_graph_pure(pos, 0.15); e2 = radius_graph_pure(p2, 0.30)
+    a32 = knn_pure(p3, p2); a21 = knn_pure(p2, pos)
+    _, m = to_dense_batch(torch.zeros(p3.size(0), 1), torch.zeros(p3.size(0), dtype=torch.long))
+
     with torch.no_grad():
-        logits = model(pos, batch, cat_idx, fps_idx_1, fps_idx_2, pos2, batch2, pos3, batch3,
-                       edge_index_1, edge_index_2, assign_index_32, assign_index_21, x_dense_mask)
-        
-        valid_parts = SEG_CLASSES[category_name]
-        mask = torch.zeros_like(logits, dtype=torch.bool)
-        mask[:, valid_parts] = True
-        logits = logits.masked_fill(~mask, -10000.0)
-        preds = logits.argmax(dim=-1).numpy()
-    
-    pos_np = pos.numpy()
-    fig = go.Figure(data=[go.Scatter3d(
-        x=pos_np[:, 0], y=pos_np[:, 1], z=pos_np[:, 2],
-        mode='markers',
-        marker=dict(size=4, color=preds, colorscale='Viridis', opacity=0.9, line=dict(width=0.5, color='black'))
-    )])
-    
-    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0), scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False)))
+        out = model(pos, torch.zeros(1024, dtype=torch.long), torch.tensor([CATEGORIES[category_name]]), f1, f2, p2, torch.zeros(p2.size(0), dtype=torch.long), p3, torch.zeros(p3.size(0), dtype=torch.long), e1, e2, a32, a21, m)
+        res = out.argmax(-1).numpy()
+
+    fig = go.Figure(data=[go.Scatter3d(x=p[:,0], y=p[:,1], z=p[:,2], mode='markers', marker=dict(size=4, color=res, colorscale='Viridis', opacity=0.8))])
+    fig.update_layout(margin=dict(l=0,r=0,b=0,t=0), scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False)))
     return fig
 
-# --- Gradio Arayüzü ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Clifford-Dirac 3D Part Segmentation (0.15M Params)")
-    gr.Markdown("Bu model, Geometrik Cebir kullanarak 3D objeleri yüksek hassasiyetle parçalara ayırır. Bir .obj veya .ply dosyası yükleyin ve sonucu keşfedin!")
-    
+    gr.Markdown("# 🚀 Clifford-Dirac 3D Segmentation (0.15M Params)")
     with gr.Row():
-        with gr.Column(scale=1):
-            input_file = gr.File(label="3D Model Yükle (.obj veya .ply)")
-            cat_dropdown = gr.Dropdown(choices=list(CATEGORIES.keys()), label="Obje Kategorisi", value="Chair")
-            run_btn = gr.Button("Segmentasyonu Başlat", variant="primary")
-            
-        with gr.Column(scale=2):
-            output_plot = gr.Plot(label="İnteraktif Segmentasyon Sonucu")
-            
-    run_btn.click(predict, inputs=[input_file, cat_dropdown], outputs=output_plot)
-
+        with gr.Column():
+            inp = gr.Model3D(label="Model Yükle (.glb, .obj)")
+            cat = gr.Dropdown(choices=list(CATEGORIES.keys()), label="Kategori", value="Chair")
+            btn = gr.Button("Tahmin Et", variant="primary")
+        out = gr.Plot(label="Sonuç")
+    btn.click(predict, [inp, cat], out)
 demo.launch()