Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PIL import Image | |
| import pandas as pd | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import uuid | |
| 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 summarization_model import generate_description, summarize_text_and_image | |
| # from data_mapping import map_object_data, create_summary_table, save_mapping_to_json | |
| # from visualization import generate_output | |
| #loading the required models | |
| 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") | |
| st.sidebar.header("Options") | |
| st.sidebar.text("Upload an image to start processing.") | |
| 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()) | |
| st.header("Results") #header for results | |
| description = generate_description(image) | |
| st.subheader("Image Description") | |
| st.write("Generated Description:", description) | |
| #extract text for the entire image | |
| extracted_text = extract_text(image) | |
| st.subheader("Extracted Text") | |
| if extracted_text: | |
| st.write(extracted_text) | |
| else: | |
| st.write("No text was detected") | |
| #summarize the entire image | |
| summary = summarize_text_and_image(description, extracted_text) | |
| st.subheader("Image Summary") | |
| st.write("Generated Summary:", summary) | |
| # 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) | |
| objects_data = [] | |
| if extracted_objects: | |
| # Save the extracted objects and their metadata | |
| metadata = save_objects_and_metadata(extracted_objects, master_id) | |
| # 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) | |
| obj_description = generate_description(obj_img) | |
| st.write("Generated Description:", description) | |
| # 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) | |
| st.write(f"Detection results for Object {i+1}:") | |
| st.json(detection_results) | |
| obj_text = extract_text(obj_img) | |
| if obj_text: | |
| st.write(f"Extracted Text for Object {i+1}:") | |
| st.json(obj_text) | |
| else: | |
| st.write("No text was detected") | |
| obj_summary = summarize_text_and_image(obj_description, obj_text) | |
| st.write(f"Object Summary:\n{obj_summary}") | |
| # object_id = str(uuid.uuid4()) | |
| # object_data = map_object_data(object_id, obj_description, obj_text, obj_summary) | |
| # objects_data.append(object_data) | |
| # data_mapping = create_summary_table(objects_data) | |
| # output_path = os.path.join("data", "output", f"{master_id}_data_mapping.json") | |
| # save_mapping_to_json(data_mapping, output_path) | |
| # # Generate the final output image with annotations and summary table | |
| # annotated_image_path, summary_table_path = generate_output(image, outputs[0]['masks'], objects_data, master_id) | |
| # st.subheader("Final Output") | |
| # # Display the annotated image | |
| # st.image(annotated_image_path, caption='Annotated Image', use_column_width=True) | |
| # # Provide a download link for the summary table | |
| # st.write("Summary Table:") | |
| # st.write(f"Download the summary table [here](data/output/{master_id}_summary.csv)") | |
| # # Display the mapped data | |
| # st.write("Mapped Data:") | |
| # st.json(data_mapping) | |
| # # # Display the JSON data | |
| # # st.write("Mapped Data:") | |
| # # st.json(data_mapping) | |
| else: | |
| st.write("No objects were detected") | |