Update main.py

main.py

@@ -34,7 +34,6 @@ def get_region_tensor(device):
     return region_map
 
 def calculate_density(coords, edge_index):
-    # Safe Density Calculation
     row, col = edge_index
     dist = torch.norm(coords[row] - coords[col], dim=1)
     sum_dist = torch.zeros(468, device=coords.device)
@@ -43,12 +42,11 @@ def calculate_density(coords, edge_index):
     count.scatter_add_(0, row, torch.ones_like(dist))
     mean_dist = sum_dist / (count + 1e-6)
     density = 1.0 / (mean_dist + 1e-6)
-    return density.unsqueeze(1) # [1, 468, 1]
+    return density.unsqueeze(1)
 
 # ==========================================
 # 2. MODEL ARCHITECTURES
 # ==========================================
-
 class RegionAwareExpert(nn.Module):
     def __init__(self):
         super().__init__()
@@ -88,7 +86,6 @@ class AnatomyLocationNet(nn.Module):
         self.region_emb = nn.Embedding(9, 32)
         self.density_proj = nn.Linear(1, 16)
         self.context_proj = nn.Linear(512, 32)
-        
         self.neck = nn.Sequential(
             nn.Linear(176, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Dropout(0.4),
             nn.Linear(256, 128), nn.LeakyReLU(0.2), nn.Linear(128, 1)
@@ -96,9 +93,7 @@ class AnatomyLocationNet(nn.Module):
 
     def forward(self, coords, pids, rids, den, ctx):
         B, N, _ = coords.shape
-        # ✅ FIXED: Use repeat instead of expand for safety on all pytorch versions
         c_emb = self.context_proj(ctx).unsqueeze(1).repeat(1, N, 1)
-        
         combined = torch.cat([self.coord_proj(coords), self.point_id_emb(pids), 
                               self.region_emb(rids), self.density_proj(den), c_emb], dim=-1)
         return self.neck(combined.view(-1, 176)).view(B, N, 1)
@@ -110,7 +105,6 @@ class GatedFusion(nn.Module):
         self.gate_net = nn.Sequential(nn.Linear(geo_dim * 2, geo_dim // 2), nn.ReLU(), nn.Linear(geo_dim // 2, geo_dim), nn.Sigmoid())
     def forward(self, x_geo, x_ctx):
         ctx_adapted = self.context_adapter(x_ctx)
-        # Use repeat for safety
         ctx_expanded = ctx_adapted.unsqueeze(1).repeat(1, 468, 1)
         combined = torch.cat([x_geo, ctx_expanded], dim=-1)
         return x_geo + (self.gate_net(combined) * ctx_expanded)
@@ -129,7 +123,6 @@ class SmartClinicalNet(nn.Module):
     def get_features(self, data):
         x, edges = data.x, data.edge_index
         batch_size = data.embedding.shape[0]
-        # Fixed unsqueeze/cat logic
         global_ctx = torch.cat([data.embedding, data.virt_prof], dim=1)
         ids = torch.arange(468, device=x.device).repeat(batch_size)
         if len(ids) > x.shape[0]: ids = ids[:x.shape[0]]
@@ -213,7 +206,7 @@ def get_top_probs(probs_array, class_list):
     return {class_list[i]: float(f"{probs_array[i]:.4f}") for i in sorted_idxs}
 
 # ==========================================
-# 5. API ENDPOINT (CRASH PROOF)
+# 5. API ENDPOINT (FINAL - WITH ALL COORDINATES)
 # ==========================================
 @app.get("/")
 def home():
@@ -222,14 +215,12 @@ def home():
 @app.post("/predict")
 async def predict_injections(file: UploadFile = File(...)):
     try:
-        # 1. Image Processing
         contents = await file.read()
         nparr = np.frombuffer(contents, np.uint8)
         img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
         if img_bgr is None: return JSONResponse(status_code=400, content={"error": "Invalid image"})
         img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
         
-        # 2. Extract Features
         res = mp_mesh.process(img_rgb)
         if not res.multi_face_landmarks:
             return JSONResponse(status_code=400, content={"error": "No face detected"})
@@ -243,16 +234,13 @@ async def predict_injections(file: UploadFile = File(...)):
         emb = torch.tensor(faces[0].embedding).float().to(DEVICE) if faces else torch.zeros(512).to(DEVICE)
         virt_prof = get_virtual_profile_norm(x_geo_norm.cpu()).to(DEVICE)
         
-        # 3. Prepare Inputs
         data = Data(x=x_geo_norm, edge_index=GLOBAL_EDGE_INDEX, embedding=emb.unsqueeze(0), virt_prof=virt_prof.unsqueeze(0))
         density = calculate_density(x_geo_norm, GLOBAL_EDGE_INDEX).unsqueeze(0).to(DEVICE)
         
         with torch.no_grad():
-            # A. Location Finder
             logits_loc = location_model(x_geo_norm.unsqueeze(0), STATIC_POINT_IDS, STATIC_REGION_IDS, density, emb.unsqueeze(0))
             probs_loc = torch.sigmoid(logits_loc).squeeze().cpu().numpy()
             
-            # B. Attribute Expert
             smart_features = backbone.get_features(data).unsqueeze(0)
             coords_input = x_geo_norm.unsqueeze(0)
             preds = expert_model(smart_features, coords_input, STATIC_POINT_IDS, STATIC_REGION_IDS)
@@ -262,7 +250,6 @@ async def predict_injections(file: UploadFile = File(...)):
             prob_de = torch.softmax(preds['depth'], dim=-1).squeeze().cpu().numpy()
             prob_p = torch.softmax(preds['product'], dim=-1).squeeze().cpu().numpy()
 
-        # 4. Build Response
         h, w, _ = img_bgr.shape
         pixel_coords = np.array([[p.x*w, p.y*h, p.z*w] for p in lms[:468]])
         optimal_indices_list = apply_nms_indices(pixel_coords, probs_loc, spacing_mm=12.0, threshold=0.4)
@@ -293,11 +280,15 @@ async def predict_injections(file: UploadFile = File(...)):
                 }
             }
             all_points_list.append(pt_info)
+        
+        # ✅ BACKEND FIX: Explicitly generating the list of landmarks here
+        all_coords_list = [[float(f"{p.x:.4f}"), float(f"{p.y:.4f}")] for p in lms[:468]]
 
         return {
             "status": "success",
-            "message": "SOTA V9: Stable",
+            "message": "SOTA V10: All Data + Landmarks",
             "injection_sites": all_points_list,
+            "all_coordinates": all_coords_list, # <--- SENDING SEPARATE LIST AGAIN
             "summary": {"total_optimal": len(optimal_set)}
         }