File size: 5,479 Bytes
71efe81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | 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
|