app.py

import streamlit as st
from PIL import Image
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from io import BytesIO
import uuid
import gc

import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from segmentation_model import load_model,transform_image, run_inference, save_input_image, save_objects_and_metadata, extract_object
from identification_model import load_yolov8_model, run_object_detection
from text_extraction_model import extract_text
# from models.summarization_model import generate_summary
# from utils.data_mapping import create_summary_table



model = load_model()
detection_model = load_yolov8_model()

def resize_image(image, size=(800, 800)):
    return image.resize(size, Image.ANTIALIAS)

def display_masks(outputs, image, threshold=0.5):
    masks = outputs[0]['masks']
    scores = outputs[0]['scores']
    
    fig, ax = plt.subplots()
    ax.imshow(np.array(image))

    extracted_objects = []
    
    for i in range(len(scores)):
        if scores[i] > threshold:
            mask = masks[i].squeeze().cpu().numpy()
            mask = np.where(mask > 0.5, 1, 0).astype(np.uint8)

            object_img = extract_object(image,mask)
            extracted_objects.append(object_img)
            #Display the mask
            ax.imshow(mask, cmap='jet', alpha=0.5)  # Overlay mask on image
    
    st.pyplot(fig)

    return extracted_objects

    

st.title("AI Pipeline: Image Segmentation, Object Detection, and Text Extraction")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Convert uploaded file to PIL Image
    image = uploaded_file
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    image = Image.open(uploaded_file).convert('RGB')
    # Generate a unique master ID for the image
    master_id = str(uuid.uuid4())
    extracted_text = extract_text(image)
    if extracted_text:
        st.write(extracted_text)
    else:
        st.write("No text was detected")
    
    # Save the input image
    save_input_image(image, master_id)
    # Transform image
    image_tensor = transform_image(image)
    outputs = run_inference(model, image_tensor)

    extracted_objects = display_masks(outputs, image)
    
    if extracted_objects:
    # Save the extracted objects and their metadata
        metadata = save_objects_and_metadata(extracted_objects, master_id)
    
    # Display metadata as a JSON output
        #st.write("Metadata for extracted objects:")
        #st.json(metadata)
    
    # Display each extracted object
        st.write("Extracted Objects:")
        for i, obj_img in enumerate(extracted_objects):
         st.image(obj_img, caption=f'Object {i+1}', use_column_width=True)
         # Convert the object image to a numpy array for YOLO inference
         obj_img_np = np.array(obj_img)
         # Run object detection on each extracted object
         detection_results = run_object_detection(detection_model, obj_img_np)
         st.write(f"Detection results for Object {i+1}:")
         st.json(detection_results)

         extracted_text = extract_text(obj_img)
         
         if extracted_text:
            st.write(f"Extracted Text for Object {i+1}:")
            st.json(extracted_text)
         else:
            st.write("No text was detected")

        # summary = generate_summary(obj_img,text_data)
        # st.write(f"Object Summary:\n{summary}")

    else:
        st.write("No objects were detected")


    # del extracted_objects
    # gc.collect()
    
    # Display results
    #display_masks(outputs, image) 





# if uploaded_file is not None:
#     image = Image.open(uploaded_file).convert("RGB")
#     st.image(image, caption='Uploaded Image.', use_column_width=True)
    
#     image_tensor = transform_image(image)
#     outputs = run_inference(model, image_tensor)
    
#     display_masks(outputs, image)



# def upload_image():
#     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
#     if uploaded_file is not None:
#         image = Image.open(uploaded_file)
#         return image
#     return None


# # def display_segmentation(image):
# #     st.image(image, caption="Original Image", use_column_width=True)
    
#     # Transform and run inference
#     # image_tensor = transform_image(image)
#     # outputs = run_inference(image_tensor)

#     # # Save segmented objects
#     # output_dir = 'segmented_objects/'
#     # save_segmented_objects(image, outputs, output_dir)
    
#     # segmented_images = [Image.open(f"{output_dir}object_{i+1}.png") for i in range(len(outputs[0]['scores']))]
#     # for img in segmented_images:
#     #     st.image(img, caption="Segmented Object", use_column_width=True)




# def main():
#     st.title("Image Processing Pipeline")

# #     uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png"])
# #     if uploaded_file:
# #         image_path = f"data/input_images/{uploaded_file.name}"
# #         image = Image.open(uploaded_file)
# #         image.save(image_path)  # Save the uploaded image for further processing
# #         st.image(image, caption="Uploaded Image")

# #         if st.button("Segment Image"):
# #             segmented = segment_image(image_path)
# #             st.image(segmented, caption="Segmented Image", use_column_width=True)

# #         if st.button("Identify and Extract Objects"):
# #             objects_data = identify_and_extract_objects(image_path)
# #             extracted_objects = []

# #             for obj_data in objects_data:
# #                 object_image = Image.open(obj_data['Image Path'])
# #                 text = extract_text(object_image)
# #                 summary = summarize_text(text)
# #                 obj_data['Text'] = text
# #                 obj_data['Summary'] = summary
# #                 extracted_objects.append(obj_data)

# #                 st.image(object_image, caption=f"Object {obj_data['ID']} - Label {obj_data['Label']}")

# #             summary_file = create_summary_table(extracted_objects)
# #             st.write(pd.DataFrame(extracted_objects))
# #             st.download_button(label="Download Summary Table", data=open(summary_file).read(), file_name="summary.csv")

# if __name__ == "__main__":
#     main()