Update main.py

main.py

@@ -9,6 +9,7 @@ from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 from torch_geometric.data import Data
 from torch_geometric.nn import GATv2Conv, BatchNorm
+from torch_geometric.utils import scatter # Needed for density calculation
 from insightface.app import FaceAnalysis
 
 # ==========================================
@@ -26,34 +27,91 @@ REGION_DATA = {
 }
 
 def get_region_tensor(device):
-    # Initialize with 8 (Other)
     region_map = torch.full((468,), 8, dtype=torch.long, device=device)
     for region_id, points in REGION_DATA.items():
         for p in points:
             if p < 468: region_map[p] = region_id
     return region_map
 
+def calculate_density(coords, edge_index):
+    # Calculates how "crowded" each point is (Key feature for Location Finder)
+    row, col = edge_index
+    dist = torch.norm(coords[row] - coords[col], dim=1)
+    mean_dist = scatter(dist, row, dim=0, reduce='mean', dim_size=468)
+    density = 1.0 / (mean_dist + 1e-6)
+    return density.unsqueeze(1) # [468, 1]
+
 # ==========================================
 # 2. MODEL ARCHITECTURES
 # ==========================================
 
+# --- A. ATTRIBUTE EXPERT (RegionAwareExpert) ---
+class RegionAwareExpert(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.visual_proj = nn.Linear(64, 64)
+        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_tech = nn.Linear(128, 3)    
+        self.head_dosage = nn.Linear(128, 8)  
+        self.head_depth = nn.Linear(128, 4)   
+        self.head_prod = nn.Linear(128, 8)    
+
+    def forward(self, features, coords, point_ids, region_ids):
+        vis = self.visual_proj(features)
+        pid = self.point_id_emb(point_ids)
+        reg = self.region_emb(region_ids)
+        xyz = self.coord_proj(coords)
+        combined = torch.cat([vis, pid, reg, xyz], dim=-1)
+        x = self.neck(combined.view(-1, 192))
+        return {
+            "tech": self.head_tech(x), 
+            "dosage": self.head_dosage(x),
+            "depth": self.head_depth(x), 
+            "product": self.head_prod(x)
+        }
+
+# --- B. LOCATION FINDER (AnatomyLocationNet) [NEW!] ---
+class AnatomyLocationNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.coord_proj = nn.Linear(3, 32)
+        self.point_id_emb = nn.Embedding(468, 64)
+        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)
+        )
+
+    def forward(self, coords, pids, rids, den, ctx):
+        B, N, _ = coords.shape
+        c_emb = self.context_proj(ctx).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__()
         self.context_adapter = nn.Linear(context_dim, geo_dim)
-        self.gate_net = nn.Sequential(
-            nn.Linear(geo_dim * 2, geo_dim // 2),
-            nn.ReLU(),
-            nn.Linear(geo_dim // 2, geo_dim),
-            nn.Sigmoid() 
-        )
+        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)
-        num_nodes = x_geo.shape[1]
-        ctx_expanded = ctx_adapted.expand(-1, num_nodes, -1)
+        ctx_expanded = ctx_adapted.expand(-1, 468, -1)
         combined = torch.cat([x_geo, ctx_expanded], dim=-1)
-        gate = self.gate_net(combined)
-        return x_geo + (gate * ctx_expanded)
+        return x_geo + (self.gate_net(combined) * ctx_expanded)
 
 class SmartClinicalNet(nn.Module):
     def __init__(self, hidden=32, heads=2):
@@ -65,67 +123,20 @@ class SmartClinicalNet(nn.Module):
         self.bn1 = BatchNorm(hidden * heads)
         self.conv2 = GATv2Conv(hidden * heads, hidden, heads=heads, concat=True)
         self.bn2 = BatchNorm(hidden * heads)
-        # Note: self.head is removed/ignored for feature extraction
-        
+
     def get_features(self, data):
-        # Extracts the 64-dim Visual + Geometric features
         x, edges = data.x, data.edge_index
         batch_size = data.embedding.shape[0]
         global_ctx = torch.cat([data.embedding, data.virt_prof], dim=1).unsqueeze(1)
-        
         ids = torch.arange(468, device=x.device).repeat(batch_size)
         if len(ids) > x.shape[0]: ids = ids[:x.shape[0]]
-        
         geo_feat = self.geo_proj(torch.cat([x, self.id_emb(ids)], dim=1))
         geo_reshaped = geo_feat.view(batch_size, 468, -1)
-        
         fused = self.fusion(geo_reshaped, global_ctx)
         fused_flat = fused.view(-1, 32)
-        
         h = F.elu(self.bn1(self.conv1(fused_flat, edges)))
         h = F.elu(self.bn2(self.conv2(h, edges)))
-        return h # [Batch*468, 64]
-
-class RegionAwareExpert(nn.Module):
-    def __init__(self):
-        super().__init__()
-        # Inputs: Visual(64) + PointID(64) + RegionID(32) + Coords(32)
-        self.visual_proj = nn.Linear(64, 64)
-        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)
-        )
-        
-        self.head_gate = nn.Linear(128, 1) # This is now the "Location/Probability" model
-        self.head_tech = nn.Linear(128, 3)    
-        self.head_dosage = nn.Linear(128, 8)  
-        self.head_depth = nn.Linear(128, 4)   
-        self.head_prod = nn.Linear(128, 8)    
-
-    def forward(self, features, coords, point_ids, region_ids):
-        vis = self.visual_proj(features)
-        pid = self.point_id_emb(point_ids)
-        reg = self.region_emb(region_ids)
-        xyz = self.coord_proj(coords)
-        
-        combined = torch.cat([vis, pid, reg, xyz], dim=-1)
-        x = self.neck(combined.view(-1, 192))
-        
-        return {
-            "is_injection": torch.sigmoid(self.head_gate(x)),
-            "tech": self.head_tech(x), 
-            "dosage": self.head_dosage(x),
-            "depth": self.head_depth(x), 
-            "product": self.head_prod(x)
-        }
+        return h 
 
 # ==========================================
 # 3. INITIALIZATION
@@ -133,40 +144,42 @@ class RegionAwareExpert(nn.Module):
 app = FastAPI()
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-print("🔄 Initializing Clinical AI...")
+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 (Pre-computed for speed)
-# Graph Edges
+# 2. Static Data
 mp_edges = mp.solutions.face_mesh.FACEMESH_TESSELATION
 s, t = [], []
 for src, dst in mp_edges:
     s.extend([src, dst]); t.extend([dst, src])
 GLOBAL_EDGE_INDEX = torch.tensor([s, t], dtype=torch.long).to(DEVICE)
-
-# Region & Point Tensors [1, 468] for batching
 STATIC_REGION_IDS = get_region_tensor(DEVICE).unsqueeze(0)
 STATIC_POINT_IDS = torch.arange(468, dtype=torch.long, device=DEVICE).unsqueeze(0)
 
-# 3. Load Models
+# 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
 
 # Load Weights
 if os.path.exists("smart_clinical_model.pth"):
-    print("📂 Loading Backbone...")
-    # strict=False allows us to ignore the old 'head' layer
     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"):
-    print("📂 Loading God Mode Expert...")
     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.")
 
@@ -174,7 +187,6 @@ print("✅ System Ready.")
 # 4. HELPERS
 # ==========================================
 def get_virtual_profile_norm(x_tensor):
-    # Calculates relative facial depth metrics
     NOSE=1; LIP=13; CHIN=152; L_CHEEK=234; L_JAW=172; L_EYE=33; R_EYE=263
     eye_dist = torch.norm(x_tensor[L_EYE] - x_tensor[R_EYE]) + 1e-6
     chin = (x_tensor[CHIN, 2] - x_tensor[LIP, 2]) / eye_dist
@@ -183,24 +195,18 @@ def get_virtual_profile_norm(x_tensor):
     return torch.tensor([chin, cheek, jaw], dtype=torch.float)
 
 def apply_nms_indices(landmarks_np, probs_np, spacing_mm=12.0, threshold=0.5):
-    # Non-Maximum Suppression to find optimal points
     valid = np.where(probs_np > threshold)[0]
     if len(valid) == 0: return []
-    
     sorted_idx = valid[np.argsort(probs_np[valid])[::-1]]
     keep = []
-    
-    # Approx pixel-to-mm scale using eyes
     eye_l, eye_r = landmarks_np[33], landmarks_np[263]
     eye_dist_px = np.linalg.norm(eye_l - eye_r)
     mm_per_px = 63.0 / (eye_dist_px + 1e-6)
     min_dist_sq = (spacing_mm / mm_per_px) ** 2
-    
     while len(sorted_idx) > 0:
         curr = sorted_idx[0]
         keep.append(int(curr))
         if len(sorted_idx) == 1: break
-        
         rest = sorted_idx[1:]
         dists = np.sum((landmarks_np[rest] - landmarks_np[curr])**2, axis=1)
         survivors = np.where(dists > min_dist_sq)[0]
@@ -212,7 +218,7 @@ def apply_nms_indices(landmarks_np, probs_np, spacing_mm=12.0, threshold=0.5):
 # ==========================================
 @app.get("/")
 def home():
-    return {"message": "Region-Aware Clinical AI is Running"}
+    return {"message": "SOTA Clinical AI (Location + Attributes) V6"}
 
 @app.post("/predict")
 async def predict_injections(file: UploadFile = File(...)):
@@ -231,30 +237,29 @@ async def predict_injections(file: UploadFile = File(...)):
     lms = res.multi_face_landmarks[0].landmark
     coords = [[p.x, p.y, p.z] for p in lms[:468]]
     x_geo = torch.tensor(coords, dtype=torch.float).to(DEVICE)
-    x_geo_norm = x_geo - x_geo.mean(dim=0) # Normalize for model
+    x_geo_norm = x_geo - x_geo.mean(dim=0)
     
     faces = face_app.get(img_bgr)
     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. Backbone Inference
+    # 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. Get Visual Features (64-dim)
-        smart_features = backbone.get_features(data).unsqueeze(0) # [1, 468, 64]
+        # 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. Prepare Inputs for God Mode
-        # [1, 468, 3] Coords
-        coords_input = x_geo_norm.unsqueeze(0) 
-        
-        # C. Expert Inference
+        # 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)
         
-        # D. Decode Outputs
-        # The Gate is now the "Location Probability"
-        probs_loc = preds['is_injection'].squeeze().cpu().numpy()
-        
         # Attributes
         prob_t = torch.softmax(preds['tech'], dim=-1).squeeze().cpu().numpy()
         prob_d = torch.softmax(preds['dosage'], dim=-1).squeeze().cpu().numpy()
@@ -265,12 +270,10 @@ async def predict_injections(file: UploadFile = File(...)):
     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 Gate probability for NMS
+    # Use the NEW high-accuracy probabilities for NMS
     optimal_indices = apply_nms_indices(pixel_coords, probs_loc, spacing_mm=12.0, threshold=0.4)
     
-    # 5. Formatting Response
-    # Labels must match your training mappings exactly
-    classes_tech = ["Bolus", "Fanning", "Microbolus"] # 0, 1, 2
+    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"]
@@ -294,7 +297,7 @@ async def predict_injections(file: UploadFile = File(...)):
 
     return {
         "status": "success",
-        "message": "expert attribute V5: ACTIVE",
+        "message": "SOTA V6: Dual-Model Active",
         "summary": {
             "total_optimal": len(optimal_list),
             "max_confidence": float(probs_loc.max())