Upload 6 files

app.py

@@ -0,0 +1,211 @@
+import shutil
+import streamlit as st
+import os
+import sys
+import pandas as pd
+import json
+from PIL import Image
+import logging
+
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models.segmentation_model import SegmentationModel
+from models.identification_model import IdentificationModel
+from models.text_extraction_model import TextExtractionModel
+from models.summarization_model import SummarizationModel
+from utils.postprocessing import save_segmented_objects
+from utils.data_mapping import map_data, save_mapped_data
+from utils.visualization import visualize_detections, visualize_segmentation, create_summary_table
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
+@st.cache_resource
+def load_segmentation_model():
+    return SegmentationModel()
+
+@st.cache_resource
+def load_identification_model():
+    return IdentificationModel()
+
+@st.cache_resource
+def load_text_extraction_model():
+    return TextExtractionModel()
+
+@st.cache_resource
+def load_summarization_model():
+    return SummarizationModel()
+
+def main():
+    st.set_page_config(layout="wide")
+    st.markdown("""
+    <style>
+    .stImage > div {
+        margin-left: auto;
+        margin-right: auto;
+    }
+    .stTable > div {
+        margin-left: auto;
+        margin-right: auto;
+    }
+    h1{ /* Title style */
+            text-align: center;
+        }
+    </style>
+    """, unsafe_allow_html=True)
+
+    def clear_segmented_objects_folder(folder_path):
+        # Remove all files in the segmented_objects folder
+        if os.path.exists(folder_path) and os.path.isdir(folder_path):
+            for filename in os.listdir(folder_path):
+                file_path = os.path.join(folder_path, filename)
+                try:
+                    if os.path.isfile(file_path) or os.path.islink(file_path):
+                        os.unlink(file_path)  # Remove the file
+                    elif os.path.isdir(file_path):
+                        shutil.rmtree(file_path)  # Remove the directory
+                except Exception as e:
+                    st.error(f'Failed to delete {file_path}. Reason: {e}')
+        else:
+            print(f"Folder '{folder_path}' does not exist, skipping the clearing step.")
+        
+    clear_segmented_objects_folder("data/segmented_objects")
+    
+    st.title("Image Processing Pipeline 🤖")
+
+    # File upload
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
+    logging.debug(f"Uploaded file: {uploaded_file}")
+
+    if uploaded_file is not None:
+        # Save uploaded file
+        input_path = os.path.join("data", "input_images", uploaded_file.name)
+        with open(input_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        logging.debug(f"File saved to: {input_path}")
+
+        image = Image.open(input_path)
+
+        # Segmentation
+        segmentation_model = load_segmentation_model()
+        masks, boxes, labels, class_name = segmentation_model.segment_image(input_path)
+        logging.debug(f"Segmentation results: {len(masks)} masks, {len(boxes)} boxes, {len(labels)} labels")
+        
+        # Save segmented objects
+        objects = save_segmented_objects(image, masks, boxes, "data/segmented_objects")
+        logging.debug(f"Saved {len(objects)} segmented objects")
+
+        # Object identification
+        identification_model = load_identification_model()
+        detections = []
+        for file in sorted(os.listdir("data/segmented_objects")):
+            f = os.path.join("data/segmented_objects", file)
+            obj_detections = identification_model.identify_objects(f, class_name)
+            if obj_detections:  # Only append if the object was identified
+                class_name.remove(obj_detections[0]['description'])
+                detections.extend(obj_detections)
+        logging.debug(f"Detections: {len(detections)} objects identified")
+
+        # Match detections to segmented objects
+        object_descriptions = []
+        for obj, det in zip(objects, detections):
+            if det:
+                object_descriptions.append(f"This is a {det['description']} with confidence {det['probability']:.2f}")
+            else:
+                object_descriptions.append("Unidentified object")
+        logging.debug(f"Object description: {detections}")
+
+        output_dir = "data/output"
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+        # Save detections
+        with open("data/output/detections.json", "w") as f:
+            json.dump(detections, f)
+        logging.debug("Detections saved to data/output/detections.json")
+
+        # Text extraction
+        text_extraction_model = load_text_extraction_model()
+        extracted_texts = [text_extraction_model.extract_text(obj[1]) for obj in objects]
+        logging.debug(f"Extracted texts: {extracted_texts}")
+
+        # Summarization
+        summarization_model = load_summarization_model()
+        summaries = [summarization_model.summarize(f"{desc} {text}") for desc, text in zip(object_descriptions, extracted_texts)]
+        logging.debug(f"Summaries: {summaries}")
+
+        # Data mapping
+        mapped_data = map_data(objects, detections, object_descriptions, extracted_texts, summaries)
+        save_mapped_data(mapped_data, "data/output/mapped_data.json")
+
+        # Visualization
+        visualize_segmentation(image, masks, "data/output/segmented_image.png")
+        visualize_detections(input_path, "data/output/detected_objects.png")
+        create_summary_table(mapped_data, "data/output/summary_table.csv")
+    
+        # Load the images and table
+
+        # Initialize session state if not already done
+        if 'show_original_image' not in st.session_state:
+            st.session_state.show_original_image = False
+        if 'show_segmented_image' not in st.session_state:
+            st.session_state.show_segmented_image = False
+        if 'show_detected_objects' not in st.session_state:
+            st.session_state.show_detected_objects = False
+        if 'show_summary_table' not in st.session_state:
+            st.session_state.show_summary_table = False
+
+        button_col1, button_col2, button_col3, button_col4 = st.columns(4)
+
+        with button_col1:
+            if st.button("Show Original Image"):
+                st.session_state.show_original_image = not st.session_state.show_original_image
+
+        with button_col2:
+            if st.button("Show Segmented Image"):
+                st.session_state.show_segmented_image = not st.session_state.show_segmented_image
+
+        with button_col3:
+            if st.button("Show Detected Objects"):
+                st.session_state.show_detected_objects = not st.session_state.show_detected_objects
+
+        with button_col4:
+            if st.button("Show Summary Table"):
+                st.session_state.show_summary_table = not st.session_state.show_summary_table
+
+        # Display components based on session state
+        def resize_image(image_path, target_width, target_height):
+            image = Image.open(image_path)
+            resized_image = image.resize((target_width, target_height))
+            return resized_image
+
+        # Set desired width and height
+        IMAGE_WIDTH = 600
+        IMAGE_HEIGHT = 400
+        
+        if st.session_state.show_original_image:
+            col1, col2, col3 = st.columns([0.3, 0.4, 0.3])
+            with col2:
+                resized_image = resize_image(input_path, IMAGE_WIDTH, IMAGE_HEIGHT)
+                st.image(resized_image, caption="Original Image", use_column_width=True)
+
+        if st.session_state.show_segmented_image:
+            col1, col2, col3 = st.columns([0.3, 0.4, 0.3])
+            with col2:
+                resized_image = resize_image("data/output/segmented_image.png", IMAGE_WIDTH, IMAGE_HEIGHT)
+                st.image(resized_image, caption="Segmented Image", use_column_width=True)
+
+        if st.session_state.show_detected_objects:
+            col1, col2, col3 = st.columns([0.3, 0.4, 0.3])
+            with col2:
+                resized_image = resize_image("data/output/detected_objects.png", IMAGE_WIDTH, IMAGE_HEIGHT)
+                st.image(resized_image, caption="Detected Objects", use_column_width=True)
+
+        if st.session_state.show_summary_table:
+            col1, col2, col3 = st.columns([1, 3, 1])
+            with col2:
+                summary_table = pd.read_csv("data/output/summary_table.csv")
+                st.table(summary_table)
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,11 @@
+torch
+torchvision
+clip
+easyocr
+transformers
+matplotlib
+pandas
+streamlit
+Pillow
+ultralytics
+opencv-python-headless

utils/data_mapping.py

@@ -0,0 +1,18 @@
+import json
+
+def map_data(objects,detections, descriptions, extracted_texts, summaries):
+    mapped_data = {}
+    for (obj_id, file_path, box),det, description, text, summary in zip(objects,detections, descriptions, extracted_texts, summaries):
+            mapped_data[obj_id] = {
+            "file_path": file_path,
+            "box": box,
+            "description": description,
+            "extracted_text": text,
+            "summary": summary
+        }
+
+    return mapped_data
+
+def save_mapped_data(mapped_data, output_file):
+    with open(output_file, "w") as f:
+        json.dump(mapped_data, f, indent=2)

utils/postprocessing.py

@@ -0,0 +1,12 @@
+import os
+from PIL import Image
+
+def save_segmented_objects(image, masks, boxes, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+    objects = []
+    for i, (mask, box) in enumerate(zip(masks, boxes)):
+        obj_image = image.crop(box)
+        file_path = os.path.join(output_dir, f"object_{i}.png")
+        obj_image.save(file_path)
+        objects.append((f"object_{i}", file_path, box.tolist()))
+    return objects

utils/visualization.py

@@ -0,0 +1,53 @@
+import matplotlib.pyplot as plt
+import pandas as pd
+import cv2
+from ultralytics import YOLO
+from PIL import Image
+def visualize_detections(image_path, output_path):
+    
+    model = YOLO('yolov8s.pt')  # You can change this to other YOLOv8 models as needed
+    # Read the image
+    image = cv2.imread(image_path)
+
+    # Run YOLOv8 inference on the image
+    results = model(image)
+
+    # Process the results and draw bounding boxes
+    for result in results:
+        boxes = result.boxes.cpu().numpy()
+        for box in boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            confidence = float(box.conf[0])
+            class_id = int(box.cls[0])
+            class_name = model.names[class_id]
+
+            # Draw bounding box
+            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+
+            # Prepare label
+            label = f"{class_name}"
+            
+            # Get label size
+            (label_width, label_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+
+            # Draw filled rectangle for label background
+            cv2.rectangle(image, (x1, y1 - label_height - baseline), (x1 + label_width, y1), (0, 255, 0), cv2.FILLED)
+
+            # Put label text
+            cv2.putText(image, label, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
+
+    # Save the output image
+    cv2.imwrite(output_path, image)
+
+def visualize_segmentation(image, masks, output_file):
+    #plt.imshow(image)
+    for mask in masks:
+        plt.imshow(mask, alpha=0.5)
+    plt.axis('off')
+    plt.savefig(output_file,bbox_inches='tight', pad_inches=0)
+    plt.close()
+
+
+def create_summary_table(mapped_data, output_file):
+    df = pd.DataFrame.from_dict(mapped_data, orient='index')
+    df.to_csv(output_file)