Update main.py

main.py

@@ -33,7 +33,7 @@ def get_region_tensor(device):
     return region_map
 
 def calculate_density(coords, edge_index):
-    # Pure PyTorch version (No extra dependencies required)
+    # Pure PyTorch: No scatter dependency issues
     row, col = edge_index
     dist = torch.norm(coords[row] - coords[col], dim=1)
     sum_dist = torch.zeros(468, device=coords.device)
@@ -87,6 +87,7 @@ 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)
@@ -94,7 +95,9 @@ class AnatomyLocationNet(nn.Module):
 
     def forward(self, coords, pids, rids, den, ctx):
         B, N, _ = coords.shape
-        c_emb = self.context_proj(ctx).unsqueeze(1).expand(-1, N, -1)
+        # ✅ FIXED: Safe expansion. Unsqueeze dim 1, then expand to N.
+        c_emb = self.context_proj(ctx).unsqueeze(1).expand(B, 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)
@@ -106,6 +109,7 @@ 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)
+        # Expand context to match geometry
         ctx_expanded = ctx_adapted.expand(-1, 468, -1)
         combined = torch.cat([x_geo, ctx_expanded], dim=-1)
         return x_geo + (self.gate_net(combined) * ctx_expanded)
@@ -159,14 +163,13 @@ backbone = SmartClinicalNet(hidden=32, heads=2).to(DEVICE)
 expert_model = RegionAwareExpert().to(DEVICE)
 location_model = AnatomyLocationNet().to(DEVICE)
 
+# Load weights safely
 if os.path.exists("smart_clinical_model.pth"):
     backbone.load_state_dict(torch.load("smart_clinical_model.pth", map_location=DEVICE), strict=False)
     backbone.eval()
-
 if os.path.exists("region_expert.pth"):
     expert_model.load_state_dict(torch.load("region_expert.pth", map_location=DEVICE))
     expert_model.eval()
-    
 if os.path.exists("anatomy_location.pth"):
     location_model.load_state_dict(torch.load("anatomy_location.pth", map_location=DEVICE))
     location_model.eval()
@@ -203,17 +206,16 @@ def apply_nms_indices(landmarks_np, probs_np, spacing_mm=12.0, threshold=0.5):
         sorted_idx = rest[survivors]
     return keep
 
-# Helper to format Top 3 probabilities
 def get_top_probs(probs_array, class_list):
     sorted_idxs = np.argsort(probs_array)[::-1][:3]
     return {class_list[i]: float(f"{probs_array[i]:.4f}") for i in sorted_idxs}
 
 # ==========================================
-# 5. API ENDPOINT (UPDATED FOR ALL POINTS)
+# 5. API ENDPOINT (FINAL VERSION)
 # ==========================================
 @app.get("/")
 def home():
-    return {"message": "SOTA Clinical AI (V8 Full Data)"}
+    return {"message": "SOTA Clinical AI - Ready"}
 
 @app.post("/predict")
 async def predict_injections(file: UploadFile = File(...)):
@@ -224,7 +226,7 @@ async def predict_injections(file: UploadFile = File(...)):
     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 Base Features
+    # 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"})
@@ -239,15 +241,16 @@ async def predict_injections(file: UploadFile = File(...)):
     virt_prof = get_virtual_profile_norm(x_geo_norm.cpu()).to(DEVICE)
     
     # 3. Prepare Inputs
+    # unsqueeze(0) for batch size of 1
     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. Run Location Finder (for Optimal Points)
+        # 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. Run Attribute Expert (for ALL points now)
+        # 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)
@@ -257,38 +260,35 @@ 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. Determine Optimal Indices (for frontend highlighting)
+    # 4. Optimal Indices (For UI Highlighting)
     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)
     optimal_set = set(optimal_indices_list)
     
+    # Classes
     classes_tech = ["Bolus", "Fanning", "Microbolus"]
     classes_dosage = ["0.01ml", "0.02ml", "0.05ml", "0.1ml", "0.2ml", "0.3ml", "0.5ml", "1.0ml"]
     classes_depth = ["Periosteal", "Subdermal", "Hypodermic", "Dermal"]
     classes_prod = ["XXL", "XL", "L", "M", "S", "Hydro", "Induce", "Lips"]
 
-    # 5. Build Response for ALL 468 Points
+    # 5. Build Response
     all_points_list = []
     
-    # Iterate through ALL points, not just optimal ones
     for idx in range(468):
-        # Extract distributions for this specific point
         p_t, p_d, p_de, p_p = prob_t[idx], prob_d[idx], prob_de[idx], prob_p[idx]
         
         pt_info = {
             "point_id": int(idx),
             "confidence": float(f"{probs_loc[idx]:.4f}"),
-            "is_optimal": idx in optimal_set, # Flag for the frontend to know if this is a "star" point
+            "is_optimal": idx in optimal_set,
             "coordinates": {"x": lms[idx].x, "y": lms[idx].y},
             "attributes": {
-                # Top Picks
                 "technique": classes_tech[np.argmax(p_t)],
                 "dosage": classes_dosage[np.argmax(p_d)],
                 "depth": classes_depth[np.argmax(p_de)],
                 "product": classes_prod[np.argmax(p_p)],
                 
-                # Full Probabilities (Fixes the 0% UI bug)
                 "technique_probs": get_top_probs(p_t, classes_tech),
                 "dosage_probs": get_top_probs(p_d, classes_dosage),
                 "depth_probs": get_top_probs(p_de, classes_depth),
@@ -306,7 +306,6 @@ async def predict_injections(file: UploadFile = File(...)):
             "total_optimal": len(optimal_set),
             "max_confidence": float(probs_loc.max())
         },
-        # NOTE: Returned key is 'injection_sites' to represent full dataset
         "injection_sites": all_points_list,
         "all_probabilities": [float(f"{p:.4f}") for p in probs_loc.tolist()],
         "all_coordinates": all_coords_list