Add automatic subject/patient ID detection from filenames and folders for grouping slices

app.py

@@ -412,29 +412,122 @@ def process_sequence(image_files, prompt_text, modality, window_type):
 # Store processed results for interactive viewer
 processed_results_cache = {}
 
-def process_slices_for_viewer(image_files, prompt_text, modality, window_type):
-    """Process all slices and cache results for interactive viewing."""
-    if model is None or processor is None:
-        return None, 0, "❌ Error: Model not loaded.", "No slices loaded"
+def extract_subject_id(file_path):
+    """Extract subject/patient ID from file path.
+    
+    Common patterns:
+    - Folder name: /subject_001/image.png -> subject_001
+    - Filename prefix: subject_001_slice_01.png -> subject_001
+    - Patient ID in filename: patient_123_slice_5.dcm -> patient_123
+    - Study UID in DICOM: extract from DICOM metadata
+    """
+    import re
+    
+    file_path = str(file_path)
+    filename = os.path.basename(file_path)
+    dir_path = os.path.dirname(file_path)
+    
+    # Try to extract from folder name (common in medical datasets)
+    folder_name = os.path.basename(dir_path.rstrip('/'))
+    if folder_name and folder_name not in ['', '.', '..']:
+        # Check if folder name looks like a subject ID
+        if re.match(r'(subject|patient|sub|pat|case|id)[_-]?\d+', folder_name, re.I):
+            return folder_name
+    
+    # Try to extract from filename
+    # Pattern: subject_001, patient_123, sub_001, case_001, etc.
+    patterns = [
+        r'(subject|patient|sub|pat|case|id)[_-]?(\d+)',  # subject_001, patient_123
+        r'(\d{3,})',  # Just numbers (3+ digits) - might be subject ID
+        r'([A-Z]{2,}\d+)',  # BR001, MR123, etc.
+    ]
+    
+    for pattern in patterns:
+        match = re.search(pattern, filename, re.I)
+        if match:
+            if len(match.groups()) > 1:
+                return f"{match.group(1)}_{match.group(2)}"
+            else:
+                return match.group(1)
     
+    # For DICOM files, try to read PatientID or StudyInstanceUID
+    if file_path.lower().endswith('.dcm'):
+        try:
+            ds = pydicom.dcmread(file_path, stop_before_pixels=True)
+            patient_id = getattr(ds, 'PatientID', None)
+            if patient_id:
+                return f"patient_{patient_id}"
+            study_uid = getattr(ds, 'StudyInstanceUID', None)
+            if study_uid:
+                # Use last part of UID as identifier
+                return f"study_{study_uid.split('.')[-1][:8]}"
+        except:
+            pass
+    
+    # Fallback: use filename without extension
+    base_name = os.path.splitext(filename)[0]
+    if len(base_name) > 0:
+        return base_name
+    
+    return "unknown"
+
+def group_images_by_subject(image_files):
+    """Group image files by subject/patient ID."""
     if not image_files:
-        return None, 0, "⚠️ Please upload medical image files.", "No slices loaded"
+        return {}
     
-    # Handle single file or list of files
     if isinstance(image_files, str):
         image_files = [image_files]
     
     # Filter out None files
     image_files = [f for f in image_files if f is not None]
     
+    # Group by subject ID
+    subject_groups = {}
+    for file_path in image_files:
+        subject_id = extract_subject_id(file_path)
+        if subject_id not in subject_groups:
+            subject_groups[subject_id] = []
+        subject_groups[subject_id].append(file_path)
+    
+    # Sort files within each group (by filename)
+    for subject_id in subject_groups:
+        subject_groups[subject_id].sort()
+    
+    return subject_groups
+
+def process_slices_for_viewer(image_files, selected_subject, prompt_text, modality, window_type):
+    """Process all slices for selected subject and cache results for interactive viewing."""
+    if model is None or processor is None:
+        return None, 0, "❌ Error: Model not loaded.", "No slices loaded", gr.Dropdown(choices=[], value=None), ""
+    
     if not image_files:
-        return None, 0, "⚠️ No valid files uploaded.", "No slices loaded"
+        return None, 0, "⚠️ Please upload medical image files.", "No slices loaded", gr.Dropdown(choices=[], value=None), ""
+    
+    # Group by subject
+    subject_groups = group_images_by_subject(image_files)
+    
+    if not subject_groups:
+        return None, 0, "⚠️ Could not detect subjects in uploaded files.", "No slices loaded", gr.Dropdown(choices=[], value=None), ""
+    
+    # Extract subject ID from selection (format: "subject_id (N slices)")
+    if selected_subject:
+        subject_id = selected_subject.split(" (")[0]
+    else:
+        # Use first subject if none selected
+        subject_id = list(subject_groups.keys())[0]
+    
+    if subject_id not in subject_groups:
+        return None, 0, f"⚠️ Subject '{subject_id}' not found.", "No slices loaded", gr.Dropdown(choices=[], value=None), ""
+    
+    # Get files for selected subject
+    subject_files = subject_groups[subject_id]
     
     results = []
     status_messages = []
     
-    for idx, image_file in enumerate(image_files):
-        status_msg = f"Processing slice {idx + 1}/{len(image_files)}..."
+    for idx, image_file in enumerate(subject_files):
+        status_msg = f"Processing slice {idx + 1}/{len(subject_files)}..."
         status_messages.append(status_msg)
         
         result = process_medical_image(image_file, prompt_text, modality, window_type)
@@ -446,47 +539,55 @@ def process_slices_for_viewer(image_files, prompt_text, modality, window_type):
             status_messages.append(f"❌ Failed to process slice {idx + 1}")
     
     if results:
-        # Cache results with a unique key
-        cache_key = f"{len(image_files)}_{prompt_text}_{modality}"
+        # Cache results with a unique key including subject ID
+        cache_key = f"{subject_id}_{len(subject_files)}_{prompt_text}_{modality}"
         processed_results_cache[cache_key] = results
         
         max_slices = len(results) - 1
-        status = f"✅ Processed {len(results)}/{len(image_files)} slices!\nUse slider or buttons to navigate."
-        slice_info = f"Slice 1/{len(results)}"
+        status = f"✅ Processed {len(results)}/{len(subject_files)} slices for {subject_id}!\nUse slider or buttons to navigate."
+        slice_info = f"Slice 1/{len(results)} ({subject_id})"
+        
+        # Update subject dropdown choices
+        choices = []
+        for sid, files in sorted(subject_groups.items()):
+            marker = "✓" if sid == subject_id else ""
+            choices.append(f"{marker} {sid} ({len(files)} slices)")
         
-        return results[0], max_slices, status, slice_info
+        return results[0], max_slices, status, slice_info, gr.Dropdown(choices=choices, value=choices[0] if choices else None), f"Viewing: {subject_id}"
     else:
-        return None, 0, "❌ No slices were processed successfully.", "No slices loaded"
+        return None, 0, "❌ No slices were processed successfully.", "No slices loaded", gr.Dropdown(choices=[], value=None), ""
 
-def navigate_slice(slice_idx, image_files, prompt_text, modality, window_type):
+def navigate_slice(slice_idx, image_files, selected_subject, prompt_text, modality, window_type):
     """Navigate to a specific slice in the sequence."""
     if not image_files:
         return None, "No slices loaded"
     
-    # Handle single file or list of files
-    if isinstance(image_files, str):
-        image_files = [image_files]
+    # Group by subject and get selected subject's files
+    subject_groups = group_images_by_subject(image_files)
     
-    # Filter out None files
-    image_files = [f for f in image_files if f is not None]
+    if selected_subject:
+        subject_id = selected_subject.split(" (")[0]
+    else:
+        subject_id = list(subject_groups.keys())[0] if subject_groups else None
     
-    if not image_files:
+    if not subject_id or subject_id not in subject_groups:
         return None, "No slices loaded"
     
+    subject_files = subject_groups[subject_id]
     slice_idx = int(slice_idx)
-    cache_key = f"{len(image_files)}_{prompt_text}_{modality}"
+    cache_key = f"{subject_id}_{len(subject_files)}_{prompt_text}_{modality}"
     
     if cache_key in processed_results_cache:
         results = processed_results_cache[cache_key]
         if 0 <= slice_idx < len(results):
-            slice_info = f"Slice {slice_idx + 1}/{len(results)}"
+            slice_info = f"Slice {slice_idx + 1}/{len(results)} ({subject_id})"
             return results[slice_idx], slice_info
     
     # If not cached, process on the fly (fallback)
-    if 0 <= slice_idx < len(image_files):
-        result = process_medical_image(image_files[slice_idx], prompt_text, modality, window_type)
+    if 0 <= slice_idx < len(subject_files):
+        result = process_medical_image(subject_files[slice_idx], prompt_text, modality, window_type)
         if result:
-            slice_info = f"Slice {slice_idx + 1}/{len(image_files)}"
+            slice_info = f"Slice {slice_idx + 1}/{len(subject_files)} ({subject_id})"
             return result, slice_info
     
     return None, f"Invalid slice index: {slice_idx}"
@@ -572,6 +673,14 @@ with gr.Blocks() as demo:
         
         with gr.Tab("Interactive Slice Viewer"):
             gr.Markdown("**Scroll through multiple slices/images from the same subject interactively**")
+            gr.Markdown("""
+            **📋 Subject Detection:** The app automatically detects subject/patient IDs from:
+            - Folder names (e.g., `subject_001/`, `patient_123/`)
+            - Filenames (e.g., `subject_001_slice_01.png`, `patient_123.dcm`)
+            - DICOM metadata (PatientID, StudyInstanceUID)
+            
+            **💡 Tip:** Upload images organized by subject folders for best results!
+            """)
             with gr.Row():
                 with gr.Column():
                     files_input = gr.File(
@@ -579,7 +688,15 @@ with gr.Blocks() as demo:
                         file_types=[".dcm", ".png", ".jpg", ".jpeg"],
                         file_count="multiple",
                         type="filepath",
-                        info="Upload multiple slices from the same subject (e.g., axial MRI slices)"
+                        info="Upload multiple slices from the same subject (e.g., axial MRI slices). Images are automatically grouped by subject ID."
+                    )
+                    
+                    subject_dropdown = gr.Dropdown(
+                        label="Select Subject/Patient",
+                        choices=[],
+                        value=None,
+                        interactive=True,
+                        info="Select which subject's slices to view (auto-detected from filenames/folders)"
                     )
                     
                     text_input_batch = gr.Textbox(
@@ -603,6 +720,7 @@ with gr.Blocks() as demo:
                             info="CT windowing preset (ignored for MRI)"
                         )
                     
+                    detect_subjects_btn = gr.Button("🔍 Detect Subjects", variant="secondary", size="sm")
                     submit_batch_btn = gr.Button("Process All Slices", variant="primary", size="lg")
                     
                     gr.Markdown("---")
@@ -636,6 +754,13 @@ with gr.Blocks() as demo:
                         interactive=False
                     )
                     
+                    subject_info_text = gr.Textbox(
+                        label="Subject Info",
+                        value="",
+                        interactive=False,
+                        visible=False
+                    )
+                    
                     gr.Markdown("### Status")
                     status_batch_text = gr.Textbox(
                         label="Processing Status",
@@ -776,46 +901,53 @@ with gr.Blocks() as demo:
         outputs=[image_output, status_text]
     )
     
+    # Detect subjects when files are uploaded
+    detect_subjects_btn.click(
+        fn=detect_subjects,
+        inputs=[files_input],
+        outputs=[subject_dropdown, status_batch_text]
+    )
+    
     # Interactive slice viewer
     submit_batch_btn.click(
         fn=process_slices_for_viewer,
-        inputs=[files_input, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
-        outputs=[current_slice_output, slice_slider, status_batch_text, slice_info_text]
+        inputs=[files_input, subject_dropdown, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
+        outputs=[current_slice_output, slice_slider, status_batch_text, slice_info_text, subject_dropdown, subject_info_text]
     ).then(
         lambda max_val: gr.Slider(maximum=max_val, interactive=True),
         inputs=[slice_slider],
         outputs=[slice_slider]
     )
     
-    def update_slice(slice_num, files, prompt, mod, window):
-        result, info = navigate_slice(int(slice_num), files, prompt, mod, window)
+    def update_slice(slice_num, files, selected_subject, prompt, mod, window):
+        result, info = navigate_slice(int(slice_num), files, selected_subject, prompt, mod, window)
         return result, info
     
     slice_slider.change(
         fn=update_slice,
-        inputs=[slice_slider, files_input, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
+        inputs=[slice_slider, files_input, subject_dropdown, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
         outputs=[current_slice_output, slice_info_text]
     )
     
-    def prev_slice(current, files, prompt, mod, window):
+    def prev_slice(current, files, selected_subject, prompt, mod, window):
         new_val = max(0, current - 1)
-        result, info = navigate_slice(new_val, files, prompt, mod, window)
+        result, info = navigate_slice(new_val, files, selected_subject, prompt, mod, window)
         return new_val, result, info
     
-    def next_slice(current, max_val, files, prompt, mod, window):
+    def next_slice(current, max_val, files, selected_subject, prompt, mod, window):
         new_val = min(max_val, current + 1)
-        result, info = navigate_slice(new_val, files, prompt, mod, window)
+        result, info = navigate_slice(new_val, files, selected_subject, prompt, mod, window)
         return new_val, result, info
     
     prev_btn.click(
         fn=prev_slice,
-        inputs=[slice_slider, files_input, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
+        inputs=[slice_slider, files_input, subject_dropdown, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
         outputs=[slice_slider, current_slice_output, slice_info_text]
     )
     
     next_btn.click(
         fn=next_slice,
-        inputs=[slice_slider, slice_slider, files_input, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
+        inputs=[slice_slider, slice_slider, files_input, subject_dropdown, text_input_batch, modality_dropdown_batch, window_dropdown_batch],
         outputs=[slice_slider, current_slice_output, slice_info_text]
     )