Spaces:
Sleeping
Sleeping
| 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() |