Update src/ui_components_original.py

src/ui_components_original.py

@@ -11,6 +11,7 @@ import html
 from typing import Optional, Dict, Any
 import numpy as np
 import cv2
+import tempfile
 
 # ---- Safe imports for local vs package execution ----
 try:
@@ -719,6 +720,7 @@ button.gr-button:hover, button.gr-button-primary:hover {
                                 with gr.Row():
                                     accept_segmentation_btn = gr.Button("✅ Accept & Generate Full Report", variant="primary")
                                     manual_edit_btn = gr.Button("✏️ Manual Edit", variant="secondary")
+                                    segmentation_only_btn = gr.Button("🎯 Get Segmentation Only", variant="secondary")
                                 
                                 # Manual editing section (initially hidden)
                                 with gr.Group(visible=False) as manual_edit_group:
@@ -743,7 +745,9 @@ button.gr-button:hover, button.gr-button-primary:hover {
                                         interactive=True
                                     )
                                     
-                                    process_manual_btn = gr.Button("🔬 Generate Report with Manual Mask", variant="primary")
+                                    with gr.Row():
+                                        process_manual_btn = gr.Button("🔬 Generate Report with Manual Mask", variant="primary")
+                                        manual_segmentation_only_btn = gr.Button("🎯 Get Manual Segmentation Only", variant="secondary")
 
                             analysis_output = gr.HTML("")
 
@@ -954,25 +958,35 @@ button.gr-button:hover, button.gr-button-primary:hover {
                 return {
                     manual_edit_group: gr.update(visible=True),
                     manual_mask_input: img_path,  # Load the original image for manual editing
-                    analysis_output: "<div class='status-warning'>⚠️ Use the drawing tool to manually mark the wound area, then click 'Generate Report with Manual Mask'.</div>"
+                    analysis_output: "<div class='status-warning'>⚠️ Use the drawing tool to manually mark the wound area, then click your desired action.</div>"
                 }
 
             def process_manual_mask(mask_data):
                 """Process the manual mask from ImageMask component"""
                 if not mask_data:
-                    return "<div class='status-error'>❌ No manual mask provided.</div>"
+                    return None
                 
                 try:
                     # Extract the mask from the ImageMask component
                     # The mask_data contains both the background image and the drawn mask
-                    if isinstance(mask_data, dict) and "layers" in mask_data:
-                        # Get the alpha channel from the first layer (the drawn mask)
-                        alpha_channel = mask_data["layers"][0][:, :, 3]
-                        # Convert to binary mask
-                        mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
+                    if isinstance(mask_data, dict):
+                        # Check if composite exists (newer format)
+                        if "composite" in mask_data:
+                            composite_img = mask_data["composite"]
+                            # Convert to grayscale and extract the drawn areas
+                            gray = cv2.cvtColor(composite_img, cv2.COLOR_RGB2GRAY)
+                            # Create mask where drawn areas are white (255) and background is black (0)
+                            _, mask = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
+                        # Check if layers exist (older format)
+                        elif "layers" in mask_data and len(mask_data["layers"]) > 0:
+                            # Get the alpha channel from the first layer (the drawn mask)
+                            alpha_channel = mask_data["layers"][0][:, :, 3]
+                            # Convert to binary mask - drawn areas have alpha > 0
+                            mask = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
+                        else:
+                            return None
                         
                         # Save the mask temporarily
-                        import tempfile
                         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                             cv2.imwrite(tmp.name, mask)
                             manual_mask_path = tmp.name
@@ -985,6 +999,87 @@ button.gr-button:hover, button.gr-button-primary:hover {
                     logging.error(f"Manual mask processing error: {e}")
                     return None
 
+            def get_segmentation_only(img_path, manual_mask_path=None):
+                """Get only the segmentation mask without full analysis"""
+                try:
+                    if not img_path:
+                        return "<div class='status-error'>❌ No image provided.</div>"
+                    
+                    # Run visual analysis to get segmentation
+                    image_pil = Image.open(img_path)
+                    
+                    if manual_mask_path:
+                        # Use manual mask
+                        visual_results = self.wound_analyzer.analyze_wound(
+                            img_path, {}, seg_adjust=0.0, manual_mask_path=manual_mask_path
+                        )
+                        mask_type = "Manual"
+                    else:
+                        # Use automatic segmentation
+                        visual_results = self.wound_analyzer.perform_visual_analysis(image_pil)
+                        mask_type = "Automatic"
+                    
+                    if not visual_results:
+                        return "<div class='status-error'>❌ Failed to generate segmentation.</div>"
+                    
+                    # Get the segmentation mask path
+                    roi_mask_path = visual_results.get("roi_mask_path")
+                    seg_path = visual_results.get("segmentation_image_path")
+                    
+                    if not roi_mask_path or not os.path.exists(roi_mask_path):
+                        return "<div class='status-error'>❌ Segmentation mask not found.</div>"
+                    
+                    # Convert mask to base64 for display
+                    mask_b64 = self.image_to_base64(roi_mask_path)
+                    seg_b64 = self.image_to_base64(seg_path) if seg_path and os.path.exists(seg_path) else None
+                    
+                    html_output = f"""
+                    <div style="max-width: 800px; margin: 0 auto; background: white; border-radius: 16px; box-shadow: 0 8px 32px rgba(0,0,0,0.1); overflow: hidden;">
+                        <div style="background: linear-gradient(135deg, #3182ce 0%, #2c5aa0 100%); color: white; padding: 30px; text-align: center;">
+                            <h1 style="margin: 0; font-size: 28px; font-weight: 700;">🎯 {mask_type} Wound Segmentation</h1>
+                            <p style="margin: 10px 0 0 0; opacity: 0.9; font-size: 16px;">Binary mask showing wound boundaries</p>
+                        </div>
+                        
+                        <div style="padding: 30px;">
+                            <div class="status-success" style="margin-bottom: 20px;">
+                                <strong>✅ Segmentation Status:</strong> {mask_type} segmentation completed successfully
+                            </div>
+                            
+                            <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin: 20px 0;">
+                                <div style="background: #f8f9fa; padding: 20px; border-radius: 12px; text-align: center;">
+                                    <h3 style="color: #2d3748; margin-top: 0;">Binary Mask</h3>
+                                    {f'<img src="{mask_b64}" style="width: 100%; border-radius: 8px;" alt="Segmentation Mask">' if mask_b64 else '<p>Mask not available</p>'}
+                                    <p style="color: #4a5568; margin: 10px 0 0 0; font-size: 14px;">White = Wound, Black = Background</p>
+                                </div>
+                                
+                                <div style="background: #f8f9fa; padding: 20px; border-radius: 12px; text-align: center;">
+                                    <h3 style="color: #2d3748; margin-top: 0;">Overlay Visualization</h3>
+                                    {f'<img src="{seg_b64}" style="width: 100%; border-radius: 8px;" alt="Segmentation Overlay">' if seg_b64 else '<p>Overlay not available</p>'}
+                                    <p style="color: #4a5568; margin: 10px 0 0 0; font-size: 14px;">Red overlay shows detected wound area</p>
+                                </div>
+                            </div>
+                            
+                            <div style="background: #fff5f5; border: 2px solid #feb2b2; padding: 20px; border-radius: 12px; margin: 20px 0;">
+                                <h3 style="color: #c53030; margin-top: 0;">📥 Download Instructions</h3>
+                                <p style="color: #742a2a; margin: 0;">Right-click on the binary mask image above and select "Save image as..." to download the segmentation mask for your use.</p>
+                            </div>
+                        </div>
+                    </div>
+                    """
+                    
+                    # Clean up temporary file if it exists
+                    if manual_mask_path and os.path.exists(manual_mask_path):
+                        try:
+                            os.unlink(manual_mask_path)
+                        except:
+                            pass
+                    
+                    return html_output
+                    
+                except Exception as e:
+                    logging.error(f"Segmentation only error: {e}")
+                    return f"<div class='status-error'>❌ Error: {html.escape(str(e))}</div>"
+
             def run_full_analysis_with_manual_mask(
                 mode, existing_label, np_name, np_age, np_gender,
                 w_loc, w_dur, pain, moist, infect, diabetic,
@@ -1023,6 +1118,42 @@ button.gr-button:hover, button.gr-button-primary:hover {
                         analysis_output: f"<div class='status-error'>❌ Analysis failed: {html.escape(str(e))}</div>"
                     }
 
+            def get_manual_segmentation_only(mask_data, img_path):
+                """Get only the manual segmentation mask"""
+                try:
+                    manual_mask_path = process_manual_mask(mask_data)
+                    if not manual_mask_path:
+                        return "<div class='status-error'>❌ Failed to process manual mask.</div>"
+                    
+                    result = get_segmentation_only(img_path, manual_mask_path)
+                    return {
+                        analysis_output: result,
+                        segmentation_preview_group: gr.update(visible=False),
+                        manual_edit_group: gr.update(visible=False)
+                    }
+                    
+                except Exception as e:
+                    logging.error(f"Manual segmentation only error: {e}")
+                    return {
+                        analysis_output: f"<div class='status-error'>❌ Error: {html.escape(str(e))}</div>"
+                    }
+
+            def get_auto_segmentation_only(img_path):
+                """Get only the automatic segmentation mask"""
+                try:
+                    result = get_segmentation_only(img_path, None)
+                    return {
+                        analysis_output: result,
+                        segmentation_preview_group: gr.update(visible=False),
+                        manual_edit_group: gr.update(visible=False)
+                    }
+                    
+                except Exception as e:
+                    logging.error(f"Auto segmentation only error: {e}")
+                    return {
+                        analysis_output: f"<div class='status-error'>❌ Error: {html.escape(str(e))}</div>"
+                    }
+
             def run_full_analysis_accept_segmentation(
                 mode, existing_label, np_name, np_age, np_gender,
                 w_loc, w_dur, pain, moist, infect, diabetic,
@@ -1156,6 +1287,13 @@ button.gr-button:hover, button.gr-button-primary:hover {
                 [analysis_output, segmentation_preview_group, manual_edit_group]
             )
             
+            # Get segmentation only (automatic)
+            segmentation_only_btn.click(
+                get_auto_segmentation_only,
+                [wound_image],
+                [analysis_output, segmentation_preview_group, manual_edit_group]
+            )
+            
             # Show manual edit interface
             manual_edit_btn.click(
                 show_manual_edit_interface,
@@ -1172,6 +1310,13 @@ button.gr-button:hover, button.gr-button-primary:hover {
                 [analysis_output, segmentation_preview_group, manual_edit_group]
             )
             
+            # Get manual segmentation only
+            manual_segmentation_only_btn.click(
+                get_manual_segmentation_only,
+                [manual_mask_input, wound_image],
+                [analysis_output, segmentation_preview_group, manual_edit_group]
+            )
+            
             history_btn.click(load_patient_history, [], [patient_history_output])
             search_patient_btn.click(search_patient_by_name, [search_patient_name], [specific_patient_output])
             view_details_btn.click(view_patient_details, [view_details_dd], [view_details_output])
@@ -1498,4 +1643,5 @@ button.gr-button:hover, button.gr-button-primary:hover {
                 'risk_score': 0,
                 'risk_level': 'Unknown',
                 'risk_factors': ['Unable to assess risk due to data processing error']
-            }
+            }
+