walaa2022 commited on
Commit
e111427
Β·
verified Β·
1 Parent(s): d95e280

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +86 -104
app.py CHANGED
@@ -1,4 +1,4 @@
1
- # app.py - MedGemma with Authentication
2
  import gradio as gr
3
  import torch
4
  from transformers import AutoProcessor, AutoModelForImageTextToText
@@ -51,27 +51,33 @@ def load_model():
51
  device = "cuda" if torch.cuda.is_available() else "cpu"
52
  logger.info(f"Using device: {device}")
53
 
 
 
 
 
 
 
 
 
 
54
  # Load model with authentication
 
55
  model = AutoModelForImageTextToText.from_pretrained(
56
  MODEL_ID,
57
- torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
58
- device_map="auto" if device == "cuda" else None,
59
  trust_remote_code=True,
60
  low_cpu_mem_usage=True,
61
- token=True # Use authenticated token
62
  )
63
-
64
- processor = AutoProcessor.from_pretrained(
65
- MODEL_ID,
66
- trust_remote_code=True,
67
- token=True # Use authenticated token
68
- )
69
-
70
  logger.info("βœ… Model loaded successfully!")
 
71
  return True
72
 
73
  except Exception as e:
74
  logger.error(f"❌ Error loading model: {str(e)}")
 
 
75
  return False
76
 
77
  # Initialize model at startup
@@ -83,15 +89,15 @@ def analyze_medical_image(image, clinical_question, patient_history=""):
83
 
84
  # Check if model is loaded
85
  if not model_loaded or model is None or processor is None:
86
- return """❌ **Model Authentication Issue**
87
 
88
- MedGemma requires authentication. Please ensure:
89
 
90
- 1. **HF_TOKEN is set**: The Space owner needs to add their Hugging Face token to Space Settings β†’ Repository secrets
91
- 2. **Model access approved**: Make sure you have access to MedGemma at https://huggingface.co/google/medgemma-4b-it
92
- 3. **Space restart**: After adding the token, restart the Space
93
 
94
- **Current Status**: Authentication failed - model cannot load without proper token."""
95
 
96
  if image is None:
97
  return "⚠️ Please upload a medical image first."
@@ -100,11 +106,11 @@ MedGemma requires authentication. Please ensure:
100
  return "⚠️ Please provide a clinical question."
101
 
102
  try:
103
- # Prepare the conversation with proper medical context
104
  messages = [
105
  {
106
  "role": "system",
107
- "content": [{"type": "text", "text": "You are MedGemma, an expert medical AI assistant specialized in medical image analysis. Provide detailed, structured analysis while emphasizing that this is for educational purposes only and should not replace professional medical diagnosis. Be thorough but clear in your explanations."}]
108
  }
109
  ]
110
 
@@ -127,6 +133,7 @@ MedGemma requires authentication. Please ensure:
127
  })
128
 
129
  # Process inputs
 
130
  inputs = processor.apply_chat_template(
131
  messages,
132
  add_generation_prompt=True,
@@ -135,32 +142,32 @@ MedGemma requires authentication. Please ensure:
135
  return_tensors="pt"
136
  )
137
 
138
- # Move to appropriate device
139
- device = next(model.parameters()).device
140
- inputs = {k: v.to(device) for k, v in inputs.items()}
 
141
 
142
  input_len = inputs["input_ids"].shape[-1]
143
 
144
- # Generate response with appropriate settings
 
145
  with torch.inference_mode():
146
  generation = model.generate(
147
  **inputs,
148
- max_new_tokens=1500,
149
  do_sample=True,
150
- temperature=0.3, # Lower temperature for more focused medical analysis
151
  top_p=0.95,
152
  repetition_penalty=1.1,
153
- pad_token_id=processor.tokenizer.eos_token_id
154
  )
155
  generation = generation[0][input_len:]
156
 
157
  # Decode response
158
  response = processor.decode(generation, skip_special_tokens=True)
159
-
160
- # Clean up response
161
  response = response.strip()
162
 
163
- # Add structured disclaimer
164
  disclaimer = """
165
 
166
  ---
@@ -173,11 +180,14 @@ MedGemma requires authentication. Please ensure:
173
  ---
174
  """
175
 
 
176
  return response + disclaimer
177
 
178
  except Exception as e:
179
- logger.error(f"Error in analyze_medical_image: {str(e)}")
180
- return f"❌ Error processing request: {str(e)}\n\nPlease try again or contact support if the issue persists."
 
 
181
 
182
  # Create Gradio interface
183
  def create_interface():
@@ -195,9 +205,9 @@ def create_interface():
195
  padding: 16px;
196
  margin: 16px 0;
197
  }
198
- .auth-warning {
199
- background-color: #fffbeb;
200
- border: 1px solid #fed7aa;
201
  border-radius: 8px;
202
  padding: 16px;
203
  margin: 16px 0;
@@ -211,25 +221,23 @@ def create_interface():
211
 
212
  **Advanced Medical AI Assistant powered by Google's MedGemma-4B**
213
 
214
- This tool can analyze various medical imaging modalities including:
215
- - 🫁 **Chest X-rays** - Pneumonia, COVID-19, lung pathology
216
- - 🧠 **CT Scans** - Brain, chest, abdomen imaging
217
- - πŸ”¬ **Histopathology** - Microscopic tissue analysis
218
- - πŸ‘οΈ **Ophthalmology** - Retinal imaging, eye conditions
219
- - 🩺 **Dermatology** - Skin lesions and conditions
220
  """)
221
 
222
- # Authentication status
223
- if not model_loaded:
 
 
 
 
 
 
 
224
  gr.Markdown("""
225
- <div class="auth-warning">
226
- πŸ” <strong>AUTHENTICATION REQUIRED</strong><br>
227
- MedGemma model requires authentication. Please:
228
- <ol>
229
- <li>Ensure you have access to the model at <a href="https://huggingface.co/google/medgemma-4b-it">MedGemma page</a></li>
230
- <li>Add your HF_TOKEN to Space Settings β†’ Repository secrets</li>
231
- <li>Restart the Space</li>
232
- </ol>
233
  </div>
234
  """)
235
 
@@ -238,93 +246,70 @@ def create_interface():
238
  <div class="disclaimer">
239
  ⚠️ <strong>IMPORTANT MEDICAL DISCLAIMER</strong><br>
240
  This tool is for <strong>educational and research purposes only</strong>.
241
- Do not upload real patient data or use for actual medical diagnosis.
242
- Always consult qualified healthcare professionals.
243
  </div>
244
  """)
245
 
246
  with gr.Row():
247
  # Left column - Inputs
248
  with gr.Column(scale=1):
249
- gr.Markdown("## πŸ“€ Upload & Configure")
250
 
251
  image_input = gr.Image(
252
  label="Medical Image",
253
  type="pil",
254
- height=350,
255
  sources=["upload", "clipboard"]
256
  )
257
 
258
  clinical_question = gr.Textbox(
259
  label="Clinical Question *",
260
- placeholder="Examples:\nβ€’ Describe the findings in this chest X-ray\nβ€’ What pathological changes are visible?\nβ€’ Provide differential diagnosis based on imaging\nβ€’ Identify any abnormalities present",
261
- lines=4,
262
- max_lines=6
263
  )
264
 
265
  patient_history = gr.Textbox(
266
  label="Patient History (Optional)",
267
- placeholder="Example: 65-year-old male presenting with chronic cough, shortness of breath, and chest pain. History of smoking for 30 years.",
268
- lines=3,
269
- max_lines=5
270
  )
271
 
272
  with gr.Row():
273
- clear_btn = gr.Button("πŸ—‘οΈ Clear All", variant="secondary")
274
  analyze_btn = gr.Button("πŸ” Analyze Image", variant="primary", size="lg")
275
 
276
- # Model status
277
- auth_status = "βœ… Authenticated" if model_loaded else "πŸ” Authentication Required"
278
- model_status = "βœ… Loaded" if model_loaded else "❌ Not Loaded"
279
 
280
  gr.Markdown(f"""
281
  **Authentication:** {auth_status}
282
  **Model Status:** {model_status}
283
- **Model:** {MODEL_ID}
284
  **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}
285
  """)
286
 
287
  # Right column - Output
288
  with gr.Column(scale=1):
289
- gr.Markdown("## πŸ“‹ Medical Analysis Results")
290
 
291
  output = gr.Textbox(
292
- label="AI Medical Analysis",
293
- lines=25,
294
- max_lines=35,
295
  show_copy_button=True,
296
- placeholder="Upload an image and ask a clinical question to get started..."
297
  )
298
 
299
- # Example cases section
300
- gr.Markdown("## πŸ’‘ Example Use Cases")
301
-
302
- with gr.Accordion("Click to see example cases", open=False):
303
  examples = gr.Examples(
304
  examples=[
305
  [
306
  "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
307
- "Analyze this chest X-ray and describe any abnormal findings. Comment on the heart size, lung fields, and overall chest anatomy.",
308
- "Adult patient presenting with respiratory symptoms including cough and shortness of breath."
309
- ],
310
- [
311
- None,
312
- "What pathological changes are visible in this medical image? Provide a structured analysis including anatomical observations and potential diagnoses.",
313
- "Patient with acute onset symptoms"
314
- ],
315
- [
316
- None,
317
- "Perform a systematic review of this imaging study. Include: 1) Technical quality assessment, 2) Normal anatomical structures, 3) Abnormal findings, 4) Clinical significance.",
318
- ""
319
- ],
320
- [
321
- None,
322
- "Compare the findings in this image to normal anatomy. What are the key differences and what might they suggest clinically?",
323
- "Follow-up imaging for known condition"
324
  ]
325
  ],
326
- inputs=[image_input, clinical_question, patient_history],
327
- label="Click any example to load it"
328
  )
329
 
330
  # Event handlers
@@ -343,23 +328,20 @@ def create_interface():
343
  outputs=[image_input, clinical_question, patient_history, output]
344
  )
345
 
346
- # Footer information
347
  gr.Markdown("""
348
  ---
349
  ### πŸ”¬ About MedGemma
350
 
351
- MedGemma is Google's specialized medical AI model trained on medical imaging and clinical text.
352
- **Note**: This model requires authentication and access approval.
353
-
354
- ### πŸ”’ Privacy & Data Policy
355
- - **No data storage**: Images and text are processed in real-time and not saved
356
- - **No patient data**: Use only synthetic, anonymized, or educational images
357
- - **Educational use**: This tool is designed for learning and research purposes
358
 
359
- ### πŸ“ž Support
360
- For technical issues or questions, please create an issue in the [Hugging Face Space repository](https://huggingface.co/spaces).
 
 
361
 
362
- **Model**: Google MedGemma-4B | **Framework**: Transformers + Gradio | **License**: Apache 2.0
363
  """)
364
 
365
  return demo
 
1
+ # app.py - MedGemma with Fixed Authentication
2
  import gradio as gr
3
  import torch
4
  from transformers import AutoProcessor, AutoModelForImageTextToText
 
51
  device = "cuda" if torch.cuda.is_available() else "cpu"
52
  logger.info(f"Using device: {device}")
53
 
54
+ # Load processor first
55
+ logger.info("Loading processor...")
56
+ processor = AutoProcessor.from_pretrained(
57
+ MODEL_ID,
58
+ trust_remote_code=True,
59
+ token=True
60
+ )
61
+ logger.info("βœ… Processor loaded successfully")
62
+
63
  # Load model with authentication
64
+ logger.info("Loading model...")
65
  model = AutoModelForImageTextToText.from_pretrained(
66
  MODEL_ID,
67
+ torch_dtype=torch.float32, # Use float32 for CPU compatibility
68
+ device_map=None, # Let PyTorch handle device placement
69
  trust_remote_code=True,
70
  low_cpu_mem_usage=True,
71
+ token=True
72
  )
 
 
 
 
 
 
 
73
  logger.info("βœ… Model loaded successfully!")
74
+
75
  return True
76
 
77
  except Exception as e:
78
  logger.error(f"❌ Error loading model: {str(e)}")
79
+ import traceback
80
+ logger.error(f"Full traceback: {traceback.format_exc()}")
81
  return False
82
 
83
  # Initialize model at startup
 
89
 
90
  # Check if model is loaded
91
  if not model_loaded or model is None or processor is None:
92
+ return """❌ **Model Loading Issue**
93
 
94
+ The model failed to load properly. This could be due to:
95
 
96
+ 1. **Memory constraints**: The model requires significant RAM
97
+ 2. **Hardware limitations**: Consider upgrading to GPU hardware
98
+ 3. **Temporary issue**: Try refreshing the page
99
 
100
+ **Current Status**: Model loading failed - please try again or contact support."""
101
 
102
  if image is None:
103
  return "⚠️ Please upload a medical image first."
 
106
  return "⚠️ Please provide a clinical question."
107
 
108
  try:
109
+ # Prepare the conversation
110
  messages = [
111
  {
112
  "role": "system",
113
+ "content": [{"type": "text", "text": "You are MedGemma, an expert medical AI assistant specialized in medical image analysis. Provide detailed, structured analysis while emphasizing that this is for educational purposes only and should not replace professional medical diagnosis."}]
114
  }
115
  ]
116
 
 
133
  })
134
 
135
  # Process inputs
136
+ logger.info("Processing input...")
137
  inputs = processor.apply_chat_template(
138
  messages,
139
  add_generation_prompt=True,
 
142
  return_tensors="pt"
143
  )
144
 
145
+ # Move to appropriate device if model is on GPU
146
+ if torch.cuda.is_available() and next(model.parameters()).is_cuda:
147
+ device = next(model.parameters()).device
148
+ inputs = {k: v.to(device) for k, v in inputs.items()}
149
 
150
  input_len = inputs["input_ids"].shape[-1]
151
 
152
+ # Generate response
153
+ logger.info("Generating response...")
154
  with torch.inference_mode():
155
  generation = model.generate(
156
  **inputs,
157
+ max_new_tokens=1000, # Reduced for stability
158
  do_sample=True,
159
+ temperature=0.3,
160
  top_p=0.95,
161
  repetition_penalty=1.1,
162
+ pad_token_id=processor.tokenizer.eos_token_id if hasattr(processor, 'tokenizer') else None
163
  )
164
  generation = generation[0][input_len:]
165
 
166
  # Decode response
167
  response = processor.decode(generation, skip_special_tokens=True)
 
 
168
  response = response.strip()
169
 
170
+ # Add medical disclaimer
171
  disclaimer = """
172
 
173
  ---
 
180
  ---
181
  """
182
 
183
+ logger.info("βœ… Analysis completed successfully")
184
  return response + disclaimer
185
 
186
  except Exception as e:
187
+ logger.error(f"❌ Error in analysis: {str(e)}")
188
+ import traceback
189
+ logger.error(f"Full traceback: {traceback.format_exc()}")
190
+ return f"❌ Analysis failed: {str(e)}\n\nPlease try with a different image or question."
191
 
192
  # Create Gradio interface
193
  def create_interface():
 
205
  padding: 16px;
206
  margin: 16px 0;
207
  }
208
+ .success {
209
+ background-color: #f0f9ff;
210
+ border: 1px solid #bae6fd;
211
  border-radius: 8px;
212
  padding: 16px;
213
  margin: 16px 0;
 
221
 
222
  **Advanced Medical AI Assistant powered by Google's MedGemma-4B**
223
 
224
+ Specialized in medical imaging across multiple modalities:
225
+ 🫁 **Radiology** β€’ πŸ”¬ **Histopathology** β€’ πŸ‘οΈ **Ophthalmology** β€’ 🩺 **Dermatology**
 
 
 
 
226
  """)
227
 
228
+ # Status display
229
+ if model_loaded:
230
+ gr.Markdown("""
231
+ <div class="success">
232
+ βœ… <strong>SYSTEM READY</strong><br>
233
+ MedGemma model is loaded and authenticated. You can now analyze medical images.
234
+ </div>
235
+ """)
236
+ else:
237
  gr.Markdown("""
238
+ <div class="disclaimer">
239
+ ⚠️ <strong>SYSTEM LOADING</strong><br>
240
+ MedGemma model is still loading. Please wait a few moments and refresh the page.
 
 
 
 
 
241
  </div>
242
  """)
243
 
 
246
  <div class="disclaimer">
247
  ⚠️ <strong>IMPORTANT MEDICAL DISCLAIMER</strong><br>
248
  This tool is for <strong>educational and research purposes only</strong>.
249
+ Do not upload real patient data. Always consult qualified healthcare professionals.
 
250
  </div>
251
  """)
252
 
253
  with gr.Row():
254
  # Left column - Inputs
255
  with gr.Column(scale=1):
256
+ gr.Markdown("## πŸ“€ Upload Medical Image")
257
 
258
  image_input = gr.Image(
259
  label="Medical Image",
260
  type="pil",
261
+ height=300,
262
  sources=["upload", "clipboard"]
263
  )
264
 
265
  clinical_question = gr.Textbox(
266
  label="Clinical Question *",
267
+ placeholder="Examples:\nβ€’ Describe the findings in this chest X-ray\nβ€’ What pathological changes are visible?\nβ€’ Provide differential diagnosis\nβ€’ Identify any abnormalities",
268
+ lines=4
 
269
  )
270
 
271
  patient_history = gr.Textbox(
272
  label="Patient History (Optional)",
273
+ placeholder="e.g., 65-year-old male with chronic cough and dyspnea",
274
+ lines=2
 
275
  )
276
 
277
  with gr.Row():
278
+ clear_btn = gr.Button("πŸ—‘οΈ Clear", variant="secondary")
279
  analyze_btn = gr.Button("πŸ” Analyze Image", variant="primary", size="lg")
280
 
281
+ # System status
282
+ auth_status = "βœ… Authenticated" if model_loaded else "πŸ”„ Loading"
283
+ model_status = "βœ… Ready" if model_loaded else "πŸ”„ Loading"
284
 
285
  gr.Markdown(f"""
286
  **Authentication:** {auth_status}
287
  **Model Status:** {model_status}
 
288
  **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}
289
  """)
290
 
291
  # Right column - Output
292
  with gr.Column(scale=1):
293
+ gr.Markdown("## πŸ“‹ Medical Analysis")
294
 
295
  output = gr.Textbox(
296
+ label="AI Analysis Results",
297
+ lines=20,
 
298
  show_copy_button=True,
299
+ placeholder="Upload a medical image and ask a clinical question to get started..."
300
  )
301
 
302
+ # Example cases
303
+ with gr.Accordion("πŸ“š Example Cases", open=False):
 
 
304
  examples = gr.Examples(
305
  examples=[
306
  [
307
  "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
308
+ "Analyze this chest X-ray for any abnormal findings. Comment on heart size, lung fields, and overall anatomy.",
309
+ "Adult patient with respiratory symptoms"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
310
  ]
311
  ],
312
+ inputs=[image_input, clinical_question, patient_history]
 
313
  )
314
 
315
  # Event handlers
 
328
  outputs=[image_input, clinical_question, patient_history, output]
329
  )
330
 
331
+ # Footer
332
  gr.Markdown("""
333
  ---
334
  ### πŸ”¬ About MedGemma
335
 
336
+ MedGemma-4B is Google's specialized medical AI model for educational medical image analysis.
337
+ It demonstrates strong performance across radiology, pathology, dermatology, and ophthalmology.
 
 
 
 
 
338
 
339
+ ### πŸ”’ Privacy & Ethics
340
+ - Real-time processing with no data retention
341
+ - Designed for educational and research use only
342
+ - No PHI or patient data should be uploaded
343
 
344
+ **Model:** Google MedGemma-4B | **License:** Apache 2.0
345
  """)
346
 
347
  return demo