Spaces:

CI-DH-I
/

Construction-Image-Identification

Sleeping

Construction-Image-Identification / processor.py

Deploy snapshot for HF Space

71efe81 4 months ago

5.48 kB

import streamlit as st

import io
import pandas as pd
import numpy as np
import cv2
import pytesseract
from PIL import Image

from ultralytics import YOLO
from sahi import AutoDetectionModel
from sahi.predict import get_prediction, get_sliced_prediction
from sahi.utils.cv import visualize_object_predictions, read_image


def split_pdf_to_images(uploaded_file, read_status=None):
    """
    Takes a PDF file, converts each page to a PNG image, 
    and returns a list of byte objects.
    """
    try:
        # Convert PDF bytes to a list of PIL Image objects
        images = convert_from_bytes(uploaded_file.read())
        
        png_files = []
        for i, image in enumerate(images):
            
            if read_status is not None:
                try:
                    read_status.update(label=f"Downloading image {i+1}/{len(images)}")
                except Exception:
                    # DeltaGenerator may expose .text; fallback to that, then to st.write
                    try:
                        read_status.text(f"Downloading image {i+1}/{len(images)}")
                    except Exception:
                        st.write(f"Downloading Image: {i+1}/{len(images)}")
            else:
                st.write(f"Downloading Image: {i+1}/{len(images)}")
            # Save PIL image to a byte buffer as a PNG
            buf = io.BytesIO()
            image.save(buf, format="PNG")
            byte_im = buf.getvalue()
            png_files.append({"name": f"page_{i+1}.png", "content": byte_im})
            
        # Reset pointer for other functions
        uploaded_file.seek(0)
        return png_files
    except Exception as e:
        return f"Error processing PDF: {e}"


def analyze_data_types(uploaded_file):
    # (Keeping your second function for the other file drop)

    try:
        df = pd.read_csv(uploaded_file)
        uploaded_file.seek(0)
        return df.dtypes
    except Exception as e:
        return f"Error: {e}"


def run_computer_vision(image_list, analysis_status=None):
    """
    Runs CV model on images, draws bounding boxes, 
    and returns analysis with processed images.
    """
    results = []
    
    # Assuming 'model' is pre-loaded (e.g., model = YOLO('yolov8n.pt'))
    
    detection_model = AutoDetectionModel.from_pretrained(
                        model_type="ultralytics",
                        model_path="Model/Prod/10_epoch_model_fixed_text_size.pt",
                        confidence_threshold=0.3,
                        device="cpu",  # or 'cuda:0'
                    )
    
    results = []
    for i in range(len(image_list)):
        
        # Update the UI via the passed-in status placeholder when available.
        if analysis_status is not None:
            try:
                analysis_status.update(label=f"Processing image {i+1}/{len(image_list)}")
            except Exception:
                # DeltaGenerator may expose .text; fallback to that, then to st.write
                try:
                    analysis_status.text(f"Processing image {i+1}/{len(image_list)}")
                except Exception:
                    st.write(f"Processing Image: {i+1}/{len(image_list)}")
        else:
            st.write(f"Processing Image: {i+1}/{len(image_list)}")
        
        image_mem = Image.open(io.BytesIO(image_list[i]["content"])).convert("RGB")
        
        sahi_result = get_sliced_prediction(
        image_mem,
        detection_model,
        slice_height=256,
        slice_width=256,
        overlap_height_ratio=0.2,
        overlap_width_ratio=0.2,
        )
        
        object_prediction_list = sahi_result.object_prediction_list

        visualization_result = visualize_object_predictions(
            image=np.array(image_mem),
            object_prediction_list=object_prediction_list,
            hide_labels=False, # This removes the text labels entirely
            hide_conf=False,  # This removes the confidence scores
            rect_th=2         # Optional: Adjust box thickness
        )

        # 5. Convert the resulting numpy array back to a PIL Image
        annotated_canvas = visualization_result["image"]
        final_image = Image.fromarray(annotated_canvas)
        results.append(final_image.copy())
        
    return memory_images_to_pdf(results) # Returning the list of dicts containing image bytes and text


def memory_images_to_pdf(pil_image_list):
    if not pil_image_list:
        return None

    # Create an in-memory buffer for the PDF
    pdf_buffer = io.BytesIO()
    
    # Ensure all images are RGB
    rgb_images = [img.convert("RGB") for img in pil_image_list]
    
    # Save to the buffer
    rgb_images[0].save(
        pdf_buffer, 
        format="PDF", 
        save_all=True, 
        append_images=rgb_images[1:]
    )
    
    # Reset buffer position to the start so it can be read
    pdf_buffer.seek(0)
    return pdf_buffer

def get_codes_from_image():

    # *Optional: Set the path to the Tesseract executable if it's not in your system's PATH*
    # pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

    # Load the image using OpenCV
    image = cv2.imread('your_image.png')

    # Convert the image to grayscale (improves accuracy)
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Use Tesseract to extract text from the image
    text = pytesseract.image_to_string(gray_image)

    # Print the extracted text
    return text

	import streamlit as st

	import io
	import pandas as pd
	import numpy as np
	import cv2
	import pytesseract
	from PIL import Image

	from ultralytics import YOLO
	from sahi import AutoDetectionModel
	from sahi.predict import get_prediction, get_sliced_prediction
	from sahi.utils.cv import visualize_object_predictions, read_image


	def split_pdf_to_images(uploaded_file, read_status=None):
	"""
	Takes a PDF file, converts each page to a PNG image,
	and returns a list of byte objects.
	"""
	try:
	# Convert PDF bytes to a list of PIL Image objects
	images = convert_from_bytes(uploaded_file.read())

	png_files = []
	for i, image in enumerate(images):

	if read_status is not None:
	try:
	read_status.update(label=f"Downloading image {i+1}/{len(images)}")
	except Exception:
	# DeltaGenerator may expose .text; fallback to that, then to st.write
	try:
	read_status.text(f"Downloading image {i+1}/{len(images)}")
	except Exception:
	st.write(f"Downloading Image: {i+1}/{len(images)}")
	else:
	st.write(f"Downloading Image: {i+1}/{len(images)}")
	# Save PIL image to a byte buffer as a PNG
	buf = io.BytesIO()
	image.save(buf, format="PNG")
	byte_im = buf.getvalue()
	png_files.append({"name": f"page_{i+1}.png", "content": byte_im})

	# Reset pointer for other functions
	uploaded_file.seek(0)
	return png_files
	except Exception as e:
	return f"Error processing PDF: {e}"


	def analyze_data_types(uploaded_file):
	# (Keeping your second function for the other file drop)

	try:
	df = pd.read_csv(uploaded_file)
	uploaded_file.seek(0)
	return df.dtypes
	except Exception as e:
	return f"Error: {e}"


	def run_computer_vision(image_list, analysis_status=None):
	"""
	Runs CV model on images, draws bounding boxes,
	and returns analysis with processed images.
	"""
	results = []

	# Assuming 'model' is pre-loaded (e.g., model = YOLO('yolov8n.pt'))

	detection_model = AutoDetectionModel.from_pretrained(
	model_type="ultralytics",
	model_path="Model/Prod/10_epoch_model_fixed_text_size.pt",
	confidence_threshold=0.3,
	device="cpu", # or 'cuda:0'
	)

	results = []
	for i in range(len(image_list)):

	# Update the UI via the passed-in status placeholder when available.
	if analysis_status is not None:
	try:
	analysis_status.update(label=f"Processing image {i+1}/{len(image_list)}")
	except Exception:
	# DeltaGenerator may expose .text; fallback to that, then to st.write
	try:
	analysis_status.text(f"Processing image {i+1}/{len(image_list)}")
	except Exception:
	st.write(f"Processing Image: {i+1}/{len(image_list)}")
	else:
	st.write(f"Processing Image: {i+1}/{len(image_list)}")

	image_mem = Image.open(io.BytesIO(image_list[i]["content"])).convert("RGB")

	sahi_result = get_sliced_prediction(
	image_mem,
	detection_model,
	slice_height=256,
	slice_width=256,
	overlap_height_ratio=0.2,
	overlap_width_ratio=0.2,
	)

	object_prediction_list = sahi_result.object_prediction_list

	visualization_result = visualize_object_predictions(
	image=np.array(image_mem),
	object_prediction_list=object_prediction_list,
	hide_labels=False, # This removes the text labels entirely
	hide_conf=False, # This removes the confidence scores
	rect_th=2 # Optional: Adjust box thickness
	)

	# 5. Convert the resulting numpy array back to a PIL Image
	annotated_canvas = visualization_result["image"]
	final_image = Image.fromarray(annotated_canvas)
	results.append(final_image.copy())

	return memory_images_to_pdf(results) # Returning the list of dicts containing image bytes and text


	def memory_images_to_pdf(pil_image_list):
	if not pil_image_list:
	return None

	# Create an in-memory buffer for the PDF
	pdf_buffer = io.BytesIO()

	# Ensure all images are RGB
	rgb_images = [img.convert("RGB") for img in pil_image_list]

	# Save to the buffer
	rgb_images[0].save(
	pdf_buffer,
	format="PDF",
	save_all=True,
	append_images=rgb_images[1:]
	)

	# Reset buffer position to the start so it can be read
	pdf_buffer.seek(0)
	return pdf_buffer

	def get_codes_from_image():

	# Optional: Set the path to the Tesseract executable if it's not in your system's PATH
	# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

	# Load the image using OpenCV
	image = cv2.imread('your_image.png')

	# Convert the image to grayscale (improves accuracy)
	gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

	# Use Tesseract to extract text from the image
	text = pytesseract.image_to_string(gray_image)

	# Print the extracted text
	return text