Update main.py

main.py

@@ -32,33 +32,22 @@ def get_region_tensor(device):
             if p < 468: region_map[p] = region_id
     return region_map
 
-# --- REPLACE THIS FUNCTION IN MAIN.PY ---
 def calculate_density(coords, edge_index):
     # Pure PyTorch version (No extra dependencies required)
     row, col = edge_index
-    
-    # 1. Calculate distance for every edge
     dist = torch.norm(coords[row] - coords[col], dim=1)
-    
-    # 2. Sum distances for each point (Scatter Add)
     sum_dist = torch.zeros(468, device=coords.device)
     sum_dist.scatter_add_(0, row, dist)
-    
-    # 3. Count neighbors for each point
     count = torch.zeros(468, device=coords.device)
     count.scatter_add_(0, row, torch.ones_like(dist))
-    
-    # 4. Calculate Mean and Density
     mean_dist = sum_dist / (count + 1e-6)
     density = 1.0 / (mean_dist + 1e-6)
-    
     return density.unsqueeze(1) # [1, 468, 1]
 
 # ==========================================
 # 2. MODEL ARCHITECTURES
 # ==========================================
 
-# --- A. ATTRIBUTE EXPERT (RegionAwareExpert) ---
 class RegionAwareExpert(nn.Module):
     def __init__(self):
         super().__init__()
@@ -66,13 +55,11 @@ class RegionAwareExpert(nn.Module):
         self.point_id_emb = nn.Embedding(468, 64)
         self.region_emb = nn.Embedding(9, 32)
         self.coord_proj = nn.Linear(3, 32)
-        
         self.neck = nn.Sequential(
             nn.Linear(192, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Dropout(0.3),
             nn.Linear(256, 128), nn.LeakyReLU(0.2)
         )
-        # Heads
-        self.head_gate = nn.Linear(128, 1) # Ignored in this version (using Location Net instead)
+        self.head_gate = nn.Linear(128, 1)
         self.head_tech = nn.Linear(128, 3)    
         self.head_dosage = nn.Linear(128, 8)  
         self.head_depth = nn.Linear(128, 4)   
@@ -92,7 +79,6 @@ class RegionAwareExpert(nn.Module):
             "product": self.head_prod(x)
         }
 
-# --- B. LOCATION FINDER (AnatomyLocationNet) [NEW!] ---
 class AnatomyLocationNet(nn.Module):
     def __init__(self):
         super().__init__()
@@ -101,7 +87,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)
@@ -109,14 +94,11 @@ class AnatomyLocationNet(nn.Module):
 
     def forward(self, coords, pids, rids, den, ctx):
         B, N, _ = coords.shape
-        # FIX: Add .unsqueeze(1) before .expand
         c_emb = self.context_proj(ctx).unsqueeze(1).expand(-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)
 
-# --- C. BACKBONE (SmartClinicalNet - Feature Extractor) ---
 class GatedFusion(nn.Module):
     def __init__(self, geo_dim, context_dim):
         super().__init__()
@@ -161,12 +143,10 @@ DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 print("🔄 Initializing SOTA Clinical System...")
 
-# 1. Analyzers
 mp_mesh = mp.solutions.face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1)
 face_app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
 face_app.prepare(ctx_id=0, det_size=(640, 640))
 
-# 2. Static Data
 mp_edges = mp.solutions.face_mesh.FACEMESH_TESSELATION
 s, t = [], []
 for src, dst in mp_edges:
@@ -175,26 +155,21 @@ GLOBAL_EDGE_INDEX = torch.tensor([s, t], dtype=torch.long).to(DEVICE)
 STATIC_REGION_IDS = get_region_tensor(DEVICE).unsqueeze(0)
 STATIC_POINT_IDS = torch.arange(468, dtype=torch.long, device=DEVICE).unsqueeze(0)
 
-# 3. Load ALL 3 Models
 backbone = SmartClinicalNet(hidden=32, heads=2).to(DEVICE)
 expert_model = RegionAwareExpert().to(DEVICE)
-location_model = AnatomyLocationNet().to(DEVICE) # New Model
+location_model = AnatomyLocationNet().to(DEVICE)
 
-# Load Weights
 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()
-    print("✅ Backbone Loaded")
 
 if os.path.exists("region_expert.pth"):
     expert_model.load_state_dict(torch.load("region_expert.pth", map_location=DEVICE))
     expert_model.eval()
-    print("✅ Attribute Expert Loaded")
     
 if os.path.exists("anatomy_location.pth"):
     location_model.load_state_dict(torch.load("anatomy_location.pth", map_location=DEVICE))
     location_model.eval()
-    print("✅ Location Finder Loaded")
 
 print("✅ System Ready.")
 
@@ -228,23 +203,26 @@ 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
 # ==========================================
 @app.get("/")
 def home():
-    return {"message": "SOTA Clinical AI (Location + Attributes) V6"}
+    return {"message": "SOTA Clinical AI (V7 Full Data Response)"}
 
 @app.post("/predict")
 async def predict_injections(file: UploadFile = File(...)):
-    # 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 Base Features
     res = mp_mesh.process(img_rgb)
     if not res.multi_face_landmarks:
         return JSONResponse(status_code=400, content={"error": "No face detected"})
@@ -258,34 +236,24 @@ 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))
-    
-    # --- CALCULATE DENSITY (Required for new model) ---
     density = calculate_density(x_geo_norm, GLOBAL_EDGE_INDEX).unsqueeze(0).to(DEVICE)
     
     with torch.no_grad():
-        # A. Run Location Finder (The 94% Model)
-        # Inputs: Coords, IDs, Regions, Density, Context
         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 (Only if we need attributes)
         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)
         
-        # Attributes
         prob_t = torch.softmax(preds['tech'], dim=-1).squeeze().cpu().numpy()
         prob_d = torch.softmax(preds['dosage'], dim=-1).squeeze().cpu().numpy()
         prob_de = torch.softmax(preds['depth'], dim=-1).squeeze().cpu().numpy()
         prob_p = torch.softmax(preds['product'], dim=-1).squeeze().cpu().numpy()
 
-    # 4. Post-Processing & NMS
     h, w, _ = img_bgr.shape
     pixel_coords = np.array([[p.x*w, p.y*h, p.z*w] for p in lms[:468]])
-    
-    # Use the NEW high-accuracy probabilities for NMS
     optimal_indices = apply_nms_indices(pixel_coords, probs_loc, spacing_mm=12.0, threshold=0.4)
     
     classes_tech = ["Bolus", "Fanning", "Microbolus"]
@@ -295,15 +263,25 @@ async def predict_injections(file: UploadFile = File(...)):
 
     optimal_list = []
     for idx in optimal_indices:
+        # 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": idx,
+            "point_id": int(idx),
             "confidence": float(f"{probs_loc[idx]:.4f}"),
             "coordinates": {"x": lms[idx].x, "y": lms[idx].y},
             "attributes": {
-                "technique": classes_tech[np.argmax(prob_t[idx])],
-                "dosage": classes_dosage[np.argmax(prob_d[idx])],
-                "depth": classes_depth[np.argmax(prob_de[idx])],
-                "product": classes_prod[np.argmax(prob_p[idx])]
+                # 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),
+                "product_probs": get_top_probs(p_p, classes_prod)
             }
         }
         optimal_list.append(pt_info)
@@ -312,7 +290,7 @@ async def predict_injections(file: UploadFile = File(...)):
 
     return {
         "status": "success",
-        "message": "SOTA V6: Dual-Model Active",
+        "message": "SOTA V7: Full Stats Active",
         "summary": {
             "total_optimal": len(optimal_list),
             "max_confidence": float(probs_loc.max())