Update main.py

main.py

@@ -8,10 +8,20 @@ import cv2
 import numpy as np
 from PIL import Image
 import io
-from typing import Dict, Any, Optional
+import json
+import os
+from typing import Dict, Any, Optional, List
 
 app = FastAPI(title="ScanAssured OCR & NER API")
 
+# --- DRUG INTERACTIONS DATABASE ---
+DRUG_INTERACTIONS = {}
+interactions_path = os.path.join(os.path.dirname(__file__), 'interactions_data.json')
+if os.path.exists(interactions_path):
+    with open(interactions_path, 'r') as f:
+        DRUG_INTERACTIONS = json.load(f)
+    print(f"Loaded {len(DRUG_INTERACTIONS)} drug interaction entries")
+
 # Enable CORS for Flutter app
 app.add_middleware(
     CORSMiddleware,
@@ -192,9 +202,6 @@ def extract_text_structured(result) -> str:
     all_words = []
 
     for page in result.pages:
-        page_height = page.dimensions[0]
-        page_width = page.dimensions[1]
-
         for block in page.blocks:
             for line in block.lines:
                 line_text = ""
@@ -202,7 +209,6 @@ def extract_text_structured(result) -> str:
 
                 for word in line.words:
                     line_text += word.value + " "
-                    # Get vertical position (y coordinate)
                     min_y = min(min_y, word.geometry[0][1])
 
                 if line_text.strip():
@@ -212,10 +218,8 @@ def extract_text_structured(result) -> str:
                         'x': line.geometry[0][0] if hasattr(line, 'geometry') else 0
                     })
 
-    # Sort by vertical position (top to bottom), then horizontal (left to right)
-    all_words.sort(key=lambda w: (round(w['y'] * 20) / 20, w['x']))  # Group similar y positions
+    all_words.sort(key=lambda w: (round(w['y'] * 20) / 20, w['x']))
 
-    # Join with newlines for lines at different vertical positions
     result_text = ""
     prev_y = -1
     for word_info in all_words:
@@ -227,6 +231,106 @@ def extract_text_structured(result) -> str:
 
     return result_text.strip()
 
+def extract_words_with_boxes(result) -> list:
+    """
+    Extract all words with their bounding boxes from docTR result.
+    Returns list of {word, bbox} where bbox is [[x0,y0], [x1,y1]] normalized 0-1.
+    """
+    words_with_boxes = []
+
+    for page in result.pages:
+        for block in page.blocks:
+            for line in block.lines:
+                for word in line.words:
+                    # geometry is ((x0, y0), (x1, y1)) normalized
+                    bbox = [
+                        [word.geometry[0][0], word.geometry[0][1]],
+                        [word.geometry[1][0], word.geometry[1][1]]
+                    ]
+                    words_with_boxes.append({
+                        'word': word.value,
+                        'bbox': bbox
+                    })
+
+    return words_with_boxes
+
+def check_drug_interactions(detected_drugs: List[str]) -> List[Dict]:
+    """
+    Check for known interactions between detected drugs.
+    Returns list of interaction warnings.
+    """
+    interactions = []
+    drugs_lower = [d.lower().strip() for d in detected_drugs]
+
+    # Check each pair of drugs
+    for i, drug1 in enumerate(drugs_lower):
+        for drug2 in drugs_lower[i+1:]:
+            # Check if drug1 interacts with drug2
+            if drug1 in DRUG_INTERACTIONS:
+                if drug2 in DRUG_INTERACTIONS[drug1]:
+                    interaction = DRUG_INTERACTIONS[drug1][drug2]
+                    interactions.append({
+                        'drug1': detected_drugs[i],
+                        'drug2': detected_drugs[drugs_lower.index(drug2)],
+                        'severity': interaction.get('severity', 'info'),
+                        'description': interaction.get('description', ''),
+                        'recommendation': interaction.get('recommendation'),
+                    })
+            # Check reverse (drug2 interacts with drug1)
+            elif drug2 in DRUG_INTERACTIONS:
+                if drug1 in DRUG_INTERACTIONS[drug2]:
+                    interaction = DRUG_INTERACTIONS[drug2][drug1]
+                    interactions.append({
+                        'drug1': detected_drugs[drugs_lower.index(drug2)],
+                        'drug2': detected_drugs[i],
+                        'severity': interaction.get('severity', 'info'),
+                        'description': interaction.get('description', ''),
+                        'recommendation': interaction.get('recommendation'),
+                    })
+
+    return interactions
+
+def map_entities_to_boxes(entities: list, words_with_boxes: list, cleaned_text: str) -> list:
+    """
+    Map NER entities back to word bounding boxes.
+    Uses fuzzy matching to find entity words in OCR words.
+    """
+    entities_with_boxes = []
+
+    for entity in entities:
+        entity_word = entity['word'].lower().strip()
+        entity_parts = entity_word.split()
+
+        # Find matching word(s) in OCR output
+        matched_boxes = []
+        for word_info in words_with_boxes:
+            ocr_word = word_info['word'].lower().strip()
+            # Check if OCR word matches any part of entity
+            for part in entity_parts:
+                if part in ocr_word or ocr_word in part:
+                    matched_boxes.append(word_info['bbox'])
+                    break
+
+        # Combine bounding boxes if multiple matches
+        if matched_boxes:
+            # Get bounding box that encompasses all matched words
+            min_x = min(box[0][0] for box in matched_boxes)
+            min_y = min(box[0][1] for box in matched_boxes)
+            max_x = max(box[1][0] for box in matched_boxes)
+            max_y = max(box[1][1] for box in matched_boxes)
+            combined_bbox = [[min_x, min_y], [max_x, max_y]]
+        else:
+            combined_bbox = None
+
+        entities_with_boxes.append({
+            'entity_group': entity['entity_group'],
+            'score': entity['score'],
+            'word': entity['word'],
+            'bbox': combined_bbox
+        })
+
+    return entities_with_boxes
+
 # --- FastAPI Routes ---
 
 @app.get("/")
@@ -316,32 +420,52 @@ async def process_image(
         doc = DocumentFile.from_images([img_bytes])
         result = ocr_predictor_instance(doc)
 
-        # Use structured extraction for better layout preservation
+        # Get image dimensions for frontend highlighting
+        img_height, img_width = preprocessed_img.shape[:2]
+
+        # Extract text and word bounding boxes
         structured_text = extract_text_structured(result)
         cleaned_text = basic_cleanup(structured_text)
+        words_with_boxes = extract_words_with_boxes(result)
 
         print(f"OCR Structured Text:\n{structured_text[:500]}...")
+        print(f"Extracted {len(words_with_boxes)} words with bounding boxes")
 
         # Perform NER on cleaned text
         print("Running NER...")
         entities = ner_pipeline(cleaned_text)
 
-        # Filter and structure entities
+        # Structure entities (return all with score > 0.1, let frontend filter)
         structured_entities = []
         for entity in entities:
-            if entity.get('score', 0.0) > 0.6:
+            if entity.get('score', 0.0) > 0.1:
                 structured_entities.append({
                     'entity_group': entity['entity_group'],
                     'score': float(entity['score']),
                     'word': entity['word'].strip(),
                 })
 
+        # Map entities to bounding boxes
+        entities_with_boxes = map_entities_to_boxes(structured_entities, words_with_boxes, cleaned_text)
+
+        # Check for drug interactions
+        detected_drugs = []
+        for entity in structured_entities:
+            if entity['entity_group'] in ['CHEM', 'CHEMICAL', 'TREATMENT']:
+                detected_drugs.append(entity['word'])
+
+        interactions = check_drug_interactions(detected_drugs) if detected_drugs else []
+        print(f"Found {len(interactions)} drug interactions")
+
         return {
-            "structured_text": structured_text,  # Preserves layout (newlines, spacing)
-            "cleaned_text": cleaned_text,         # Single line for NER
-            "medical_entities": structured_entities,
+            "structured_text": structured_text,
+            "cleaned_text": cleaned_text,
+            "medical_entities": entities_with_boxes,
+            "interactions": interactions,  # NEW: Drug interaction warnings
             "model_id": NER_MODELS[ner_model_id]["name"],
-            "ocr_model": f"{det_arch} + {reco_arch}"
+            "ocr_model": f"{det_arch} + {reco_arch}",
+            "image_width": img_width,
+            "image_height": img_height
         }
 
     except Exception as e: