Create app.py

app.py

@@ -0,0 +1,175 @@
+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 models.text_extraction_model import extract_text
+# from models.summarization_model import summarize_text
+# 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("Image Segmentation with Mask R-CNN and Object Detection")
+
+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())
+    
+    # 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)
+    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()