resolved the PermissionError

medrax/tools/report_generation.py

@@ -2,6 +2,7 @@ from typing import Any, Dict, Optional, Tuple, Type
 from pydantic import BaseModel, Field
 
 import torch
+import os  # Added to create local cache dir
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForToolRun,
@@ -28,18 +29,6 @@ class ChestXRayInput(BaseModel):
 
 
 class ChestXRayReportGeneratorTool(BaseTool):
-    """Tool that generates comprehensive chest X-ray reports with both findings and impressions.
-
-    This tool uses two Vision-Encoder-Decoder models (ViT-BERT) trained on CheXpert
-    and MIMIC-CXR datasets to generate structured radiology reports. It automatically
-    generates both detailed findings and impression summaries for each chest X-ray,
-    following standard radiological reporting format.
-
-    The tool uses:
-    - Findings model: Generates detailed observations of all visible structures
-    - Impression model: Provides concise clinical interpretation and key diagnoses
-    """
-
     name: str = "chest_xray_report_generator"
     description: str = (
         "A tool that analyzes chest X-ray images and generates comprehensive radiology reports "
@@ -47,7 +36,7 @@ class ChestXRayReportGeneratorTool(BaseTool):
         "to a chest X-ray image file. Output is a structured report with both detailed "
         "observations and key clinical conclusions."
     )
-    device: Optional[str] = "cpu"  # Change the device to "cpu"
+    device: Optional[str] = "cpu"
     args_schema: Type[BaseModel] = ChestXRayInput
     findings_model: VisionEncoderDecoderModel = None
     impression_model: VisionEncoderDecoderModel = None
@@ -57,12 +46,12 @@ class ChestXRayReportGeneratorTool(BaseTool):
     impression_processor: ViTImageProcessor = None
     generation_args: Dict[str, Any] = None
 
-    def __init__(self, cache_dir: str = "/model-weights", device: Optional[str] = "cpu"):
-        """Initialize the ChestXRayReportGeneratorTool with both findings and impression models."""
+    def __init__(self, cache_dir: str = "./model_weights", device: Optional[str] = "cpu"):
         super().__init__()
-        self.device = torch.device(device) if device else torch.device("cpu")  # Ensure CPU is used
+        os.makedirs(cache_dir, exist_ok=True)  # ✅ Ensure local folder exists
+        self.device = torch.device(device) if device else torch.device("cpu")
 
-        # Initialize findings model
+        # Load findings model
         self.findings_model = VisionEncoderDecoderModel.from_pretrained(
             "IAMJB/chexpert-mimic-cxr-findings-baseline", cache_dir=cache_dir
         ).eval()
@@ -73,7 +62,7 @@ class ChestXRayReportGeneratorTool(BaseTool):
             "IAMJB/chexpert-mimic-cxr-findings-baseline", cache_dir=cache_dir
         )
 
-        # Initialize impression model
+        # Load impression model
         self.impression_model = VisionEncoderDecoderModel.from_pretrained(
             "IAMJB/chexpert-mimic-cxr-impression-baseline", cache_dir=cache_dir
         ).eval()
@@ -84,11 +73,10 @@ class ChestXRayReportGeneratorTool(BaseTool):
             "IAMJB/chexpert-mimic-cxr-impression-baseline", cache_dir=cache_dir
         )
 
-        # Move models to device (CPU)
+        # Move models to CPU
         self.findings_model = self.findings_model.to(self.device)
         self.impression_model = self.impression_model.to(self.device)
 
-        # Default generation arguments
         self.generation_args = {
             "num_return_sequences": 1,
             "max_length": 128,
@@ -99,19 +87,8 @@ class ChestXRayReportGeneratorTool(BaseTool):
     def _process_image(
         self, image_path: str, processor: ViTImageProcessor, model: VisionEncoderDecoderModel
     ) -> torch.Tensor:
-        """Process the input image for a specific model.
-
-        Args:
-            image_path (str): Path to the input image.
-            processor: Image processor for the specific model.
-            model: The model to process the image for.
-
-        Returns:
-            torch.Tensor: Processed image tensor ready for model input.
-        """
         image = Image.open(image_path).convert("RGB")
         pixel_values = processor(image, return_tensors="pt").pixel_values
-
         expected_size = model.config.encoder.image_size
         actual_size = pixel_values.shape[-1]
 
@@ -123,23 +100,11 @@ class ChestXRayReportGeneratorTool(BaseTool):
                 align_corners=False,
             )
 
-        pixel_values = pixel_values.to(self.device)
-
-        return pixel_values
+        return pixel_values.to(self.device)
 
     def _generate_report_section(
         self, pixel_values: torch.Tensor, model: VisionEncoderDecoderModel, tokenizer: BertTokenizer
     ) -> str:
-        """Generate a report section using the specified model.
-
-        Args:
-            pixel_values: Processed image tensor.
-            model: The model to use for generation.
-            tokenizer: The tokenizer for the model.
-
-        Returns:
-            str: Generated text for the report section.
-        """
         generation_config = GenerationConfig(
             **{
                 **self.generation_args,
@@ -149,9 +114,7 @@ class ChestXRayReportGeneratorTool(BaseTool):
                 "decoder_start_token_id": tokenizer.cls_token_id,
             }
         )
-
         generated_ids = model.generate(pixel_values, generation_config=generation_config)
-
         return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
     def _run(
@@ -159,17 +122,7 @@ class ChestXRayReportGeneratorTool(BaseTool):
         image_path: str,
         run_manager: Optional[CallbackManagerForToolRun] = None,
     ) -> Tuple[str, Dict]:
-        """Generate a comprehensive chest X-ray report containing both findings and impression.
-
-        Args:
-            image_path (str): The path to the chest X-ray image file.
-            run_manager (Optional[CallbackManagerForToolRun]): The callback manager.
-
-        Returns:
-            Tuple[str, Dict]: A tuple containing the complete report and metadata.
-        """
         try:
-            # Process image for both models
             findings_pixels = self._process_image(
                 image_path, self.findings_processor, self.findings_model
             )
@@ -177,7 +130,6 @@ class ChestXRayReportGeneratorTool(BaseTool):
                 image_path, self.impression_processor, self.impression_model
             )
 
-            # Generate both sections
             with torch.inference_mode():
                 findings_text = self._generate_report_section(
                     findings_pixels, self.findings_model, self.findings_tokenizer
@@ -186,19 +138,16 @@ class ChestXRayReportGeneratorTool(BaseTool):
                     impression_pixels, self.impression_model, self.impression_tokenizer
                 )
 
-            # Combine into formatted report
             report = (
                 "CHEST X-RAY REPORT\n\n"
                 f"FINDINGS:\n{findings_text}\n\n"
                 f"IMPRESSION:\n{impression_text}"
             )
-
             metadata = {
                 "image_path": image_path,
                 "analysis_status": "completed",
                 "sections_generated": ["findings", "impression"],
             }
-
             return report, metadata
 
         except Exception as e:
@@ -213,5 +162,4 @@ class ChestXRayReportGeneratorTool(BaseTool):
         image_path: str,
         run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
     ) -> Tuple[str, Dict]:
-        """Asynchronously generate a comprehensive chest X-ray report."""
         return self._run(image_path)