Create app.py

app.py

@@ -0,0 +1,128 @@
+import gradio as gr
+from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+import torch
+import json
+
+class SinogramAnalysisSystem:
+    def __init__(self):
+        print("Initializing system...")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Load analysis models
+        print("Loading tumor detection models...")
+        self.tumor_classifier = AutoModelForImageClassification.from_pretrained(
+            "SIATCN/vit_tumor_classifier"
+        ).to(self.device)
+        self.tumor_processor = AutoImageProcessor.from_pretrained("SIATCN/vit_tumor_classifier")
+        
+        self.size_classifier = AutoModelForImageClassification.from_pretrained(
+            "SIATCN/vit_tumor_radius_detection_finetuned"
+        ).to(self.device)
+        self.size_processor = AutoImageProcessor.from_pretrained(
+            "SIATCN/vit_tumor_radius_detection_finetuned"
+        )
+        
+        # Load Hymba model
+        print("Loading Hymba model...")
+        repo_name = "nvidia/Hymba-1.5B-Base"
+        self.tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
+        self.llm = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
+        self.llm = self.llm.to(self.device).to(torch.bfloat16)
+        
+        print("System ready!")
+
+    def process_sinogram(self, image):
+        if isinstance(image, str):
+            image = Image.open(image)
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        return image.resize((224, 224))
+
+    @torch.no_grad()
+    def analyze_sinogram(self, processed_image):
+        # Detect tumor
+        inputs = self.tumor_processor(processed_image, return_tensors="pt").to(self.device)
+        outputs = self.tumor_classifier(**inputs)
+        tumor_present = outputs.logits.softmax(dim=-1)[0].cpu()
+        has_tumor = tumor_present[1] > tumor_present[0]
+        
+        # Assess size
+        size_inputs = self.size_processor(processed_image, return_tensors="pt").to(self.device)
+        size_outputs = self.size_classifier(**size_inputs)
+        size_pred = size_outputs.logits.softmax(dim=-1)[0].cpu()
+        sizes = ["no-tumor", "0.5", "1.0", "1.5"]
+        tumor_size = sizes[size_pred.argmax().item()]
+        
+        return has_tumor, tumor_size
+
+    def generate_report(self, tumor_present, tumor_size):
+        prompt = f"""As a medical professional, provide a brief analysis of these sinogram findings:
+
+Findings:
+- Tumor Detection: {'Positive' if tumor_present else 'Negative'}
+- Tumor Size: {tumor_size} cm
+
+Please provide:
+1. Brief interpretation
+2. Clinical recommendations
+3. Follow-up plan"""
+
+        # Generate response using Hymba
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        outputs = self.llm.generate(
+            **inputs,
+            max_length=512,
+            do_sample=True,
+            temperature=0.7,
+            use_cache=True
+        )
+        
+        response = self.tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:],
+            skip_special_tokens=True
+        )
+        
+        return response.strip()
+
+    def analyze_image(self, image):
+        try:
+            # Process sinogram
+            processed = self.process_sinogram(image)
+            tumor_present, tumor_size = self.analyze_sinogram(processed)
+            
+            # Generate medical report
+            report = self.generate_report(tumor_present, tumor_size)
+            
+            # Format results
+            return f"""
+SINOGRAM ANALYSIS:
+• Tumor Detection: {'Positive' if tumor_present else 'Negative'}
+• Size Assessment: {tumor_size} cm
+
+MEDICAL REPORT:
+{report}
+"""
+        except Exception as e:
+            return f"Error during analysis: {str(e)}"
+
+def create_interface():
+    system = SinogramAnalysisSystem()
+    
+    iface = gr.Interface(
+        fn=system.analyze_image,
+        inputs=[
+            gr.Image(type="pil", label="Upload Sinogram")
+        ],
+        outputs=[
+            gr.Textbox(label="Analysis Results", lines=15)
+        ],
+        title="Sinogram Analysis System",
+        description="Upload a sinogram for tumor detection and medical assessment."
+    )
+    
+    return iface
+
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch(debug=True, share=True)