mystdreamm/Image_Segmentation

1import streamlit as st2import os3import json4from PIL import Image5from pipeline_module import EnhancedMaskRCNNPipeline6 7# Initialize the pipeline8@st.cache_resource9def load_pipeline():10    config_file = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"11    weights_file = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"12    return EnhancedMaskRCNNPipeline(config_file, weights_file)13 14# Streamlit app15def run_streamlit_app():16    st.title("AI Pipeline for Image Segmentation and Object Analysis")17    st.write("Upload an image to segment objects, extract text, and analyze!")18 19    # Load the pipeline20    pipeline = load_pipeline()21 22    # Output directory23    base_output_dir = "data"24    input_images_dir = os.path.join(base_output_dir, "input_images")25    segmented_objects_dir = os.path.join(base_output_dir, "segmented_objects", "segmented_objects")26    mapped_files_dir = os.path.join(base_output_dir, "segmented_objects", "mapped_files")27    output_visualizations_dir = os.path.join(base_output_dir, "output", "visualizations")28    output_summaries_dir = os.path.join(base_output_dir, "output", "summaries")29    output_texts_dir = os.path.join(base_output_dir, "output", "extracted_texts")30 31    # Create directories if they don't exist (before file uploading and processing)32    for directory in [input_images_dir, segmented_objects_dir, mapped_files_dir, output_visualizations_dir, output_summaries_dir, output_texts_dir]:33        os.makedirs(directory, exist_ok=True)34 35    # File uploader36    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])37 38    if uploaded_file is not None:39        try:40            # Save the uploaded image to 'input_images' folder within the output directory41            image_path = os.path.join(input_images_dir, uploaded_file.name)42            with open(image_path, "wb") as f:43                f.write(uploaded_file.getbuffer())44 45            # Process the image46            with st.spinner("Processing image..."):47                result = pipeline.process_image(image_path, base_output_dir)48 49            # Display the segmented image50            st.subheader("Segmented Image")51            segmented_image_path = os.path.join(output_visualizations_dir, f"{result['master_id']}_visualized.jpg")52            st.image(segmented_image_path, use_column_width=True)53 54            # Display whole image summary55            st.subheader("Whole Image Summary")56            st.write(result['whole_image_summary'])57 58            # Display extracted text for the whole image59            st.subheader("Extracted Text")60            if result['objects']:61                object_text_path = os.path.join(output_texts_dir, f"{result['objects'][0]['object_id']}_text.txt")62                with open(object_text_path, 'r') as f:63                    extracted_text = f.read().strip()64                    if extracted_text:65                        st.write(extracted_text)66                    else:67                        st.write("No text was extracted from the image.")68            else:69                st.write("No objects detected in the image.")70 71            # Display object analysis results72            st.subheader("Object Analysis")73            for obj in result['objects']:74                with st.expander(f"{obj['class']} (Confidence: {obj['score']:.2f})"):75                    st.image(obj['object_path'], use_column_width=True)76 77            # Display links to output folders78            st.subheader("Output Folders")79            st.write(f"- Input Images: {os.path.abspath(input_images_dir)}")80            st.write(f"- Segmented Objects: {os.path.abspath(segmented_objects_dir)}")81            st.write(f"- Mapped Files: {os.path.abspath(mapped_files_dir)}")82            st.write(f"- Output (Visualizations): {os.path.abspath(output_visualizations_dir)}")83            st.write(f"- Output (Summaries): {os.path.abspath(output_summaries_dir)}")84            st.write(f"- Output (Extracted Texts): {os.path.abspath(output_texts_dir)}")85 86            st.success(f"Processing complete! Results saved in: {base_output_dir}")87        except Exception as e:88            st.error(f"An error occurred during processing: {str(e)}")89    else:90        st.write("👆 Upload an image to get started!")91 92# Run the Streamlit app93if __name__ == "__main__":94    run_streamlit_app()