Spaces:

yansari
/

Tesseract

Running

Tesseract / Models /Text_Models /text_bboxer.py

Deploy: floorplan rendered as DOM element, edges always visible

1e3df84 9 months ago

4.02 kB

import os
import torch
import cv2
import numpy as np
from torch.autograd import Variable
from collections import OrderedDict
import sys

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
# Import utility modules (ensure these are available in your PYTHONPATH or adjust imports)
import craft_utils
import imgproc
import file_utils
from craft import CRAFT

# Helper function to load model weights
def copyStateDict(state_dict):
    if list(state_dict.keys())[0].startswith("module"):
        start_idx = 1
    else:
        start_idx = 0
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = ".".join(k.split(".")[start_idx:])
        new_state_dict[name] = v
    return new_state_dict

def get_Textboxes(input_image_path, model_weights_dir, output_directory):
    """
    Perform text detection and return the path to the text file containing bounding box coordinates.
    Args:
        input_image_path (str): Path to the input image.
        model_checkpoint (str): Path to the model checkpoint file.
        output_directory (str): Path to the directory where results should be saved.
    Returns:
        str: Path to the text file containing bounding box coordinates.
    """
    # Check for CUDA availability
    cuda_available = torch.cuda.is_available()
    model_checkpoint = os.path.join(model_weights_dir, "craft_mlt_25k.pth")
    if not os.path.exists(model_checkpoint):
        raise FileNotFoundError(f"Error: Model checkpoint not found at '{model_checkpoint}'.")

    #print(f"CUDA available: {cuda_available}")

    # Create output directory if it doesn't exist
    os.makedirs(output_directory, exist_ok=True)

    # Derive image-specific folder
    image_name_no_ext = os.path.splitext(os.path.basename(input_image_path))[0]
    image_specific_folder = os.path.join(output_directory, image_name_no_ext)
    os.makedirs(image_specific_folder, exist_ok=True)
    
    # Load the input image
    image = imgproc.loadImage(input_image_path)

    # Load the CRAFT model 
    net = CRAFT()
    if cuda_available:
        net.load_state_dict(copyStateDict(torch.load(model_checkpoint, weights_only=False)))
        net = net.cuda()
        net = torch.nn.DataParallel(net)
    else:
        net.load_state_dict(copyStateDict(torch.load(model_checkpoint, map_location="cpu", weights_only=False)))
    net.eval()

    # Preprocess the image 
    img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
        image, 1280, interpolation=cv2.INTER_LINEAR, mag_ratio=1.5
    )
    ratio_h = ratio_w = 1 / target_ratio
    x = imgproc.normalizeMeanVariance(img_resized)
    x = torch.from_numpy(x).permute(2, 0, 1)
    x = Variable(x.unsqueeze(0))
    if cuda_available:
        x = x.cuda()

    # Forward pass
    with torch.no_grad():
        y, _ = net(x)

    # Generate score maps
    score_text = y[0, :, :, 0].cpu().data.numpy()
    score_link = y[0, :, :, 1].cpu().data.numpy()

    # Post-processing
    boxes, polys = craft_utils.getDetBoxes(score_text, score_link, 0.7, 0.4, 0.4, False)
    boxes = craft_utils.adjustResultCoordinates(boxes, ratio_w, ratio_h)
    polys = craft_utils.adjustResultCoordinates(polys, ratio_w, ratio_h)

    # Handle None polygons
    for i in range(len(polys)):
        if polys[i] is None:
            polys[i] = boxes[i]

    # Save the mask image
    mask_output_path = os.path.join(image_specific_folder, f"{image_name_no_ext}_mask.jpg")
    render_img = np.hstack((score_text, score_link))
    cv2.imwrite(mask_output_path, imgproc.cvt2HeatmapImg(render_img))

    # Save the detected text regions and results text file
    file_utils.saveResult(input_image_path, image[:, :, ::-1], polys, dirname=image_specific_folder)

    # Return the path to the bounding boxes text file
    print("\nCreating Text bboxes...")
    text_file_path = os.path.join(image_specific_folder, f"res_{image_name_no_ext}.txt") 
    print(f"{len(boxes)} Text bounding boxes saved in: {text_file_path}")
    return text_file_path 

	import os
	import torch
	import cv2
	import numpy as np
	from torch.autograd import Variable
	from collections import OrderedDict
	import sys

	sys.path.append(os.path.dirname(os.path.abspath(__file__)))
	# Import utility modules (ensure these are available in your PYTHONPATH or adjust imports)
	import craft_utils
	import imgproc
	import file_utils
	from craft import CRAFT

	# Helper function to load model weights
	def copyStateDict(state_dict):
	if list(state_dict.keys())[0].startswith("module"):
	start_idx = 1
	else:
	start_idx = 0
	new_state_dict = OrderedDict()
	for k, v in state_dict.items():
	name = ".".join(k.split(".")[start_idx:])
	new_state_dict[name] = v
	return new_state_dict

	def get_Textboxes(input_image_path, model_weights_dir, output_directory):
	"""
	Perform text detection and return the path to the text file containing bounding box coordinates.

	Args:
	input_image_path (str): Path to the input image.
	model_checkpoint (str): Path to the model checkpoint file.
	output_directory (str): Path to the directory where results should be saved.

	Returns:
	str: Path to the text file containing bounding box coordinates.
	"""
	# Check for CUDA availability
	cuda_available = torch.cuda.is_available()
	model_checkpoint = os.path.join(model_weights_dir, "craft_mlt_25k.pth")
	if not os.path.exists(model_checkpoint):
	raise FileNotFoundError(f"Error: Model checkpoint not found at '{model_checkpoint}'.")

	#print(f"CUDA available: {cuda_available}")

	# Create output directory if it doesn't exist
	os.makedirs(output_directory, exist_ok=True)

	# Derive image-specific folder
	image_name_no_ext = os.path.splitext(os.path.basename(input_image_path))[0]
	image_specific_folder = os.path.join(output_directory, image_name_no_ext)
	os.makedirs(image_specific_folder, exist_ok=True)

	# Load the input image
	image = imgproc.loadImage(input_image_path)

	# Load the CRAFT model
	net = CRAFT()
	if cuda_available:
	net.load_state_dict(copyStateDict(torch.load(model_checkpoint, weights_only=False)))
	net = net.cuda()
	net = torch.nn.DataParallel(net)
	else:
	net.load_state_dict(copyStateDict(torch.load(model_checkpoint, map_location="cpu", weights_only=False)))
	net.eval()

	# Preprocess the image
	img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(
	image, 1280, interpolation=cv2.INTER_LINEAR, mag_ratio=1.5
	)
	ratio_h = ratio_w = 1 / target_ratio
	x = imgproc.normalizeMeanVariance(img_resized)
	x = torch.from_numpy(x).permute(2, 0, 1)
	x = Variable(x.unsqueeze(0))
	if cuda_available:
	x = x.cuda()

	# Forward pass
	with torch.no_grad():
	y, _ = net(x)

	# Generate score maps
	score_text = y[0, :, :, 0].cpu().data.numpy()
	score_link = y[0, :, :, 1].cpu().data.numpy()

	# Post-processing
	boxes, polys = craft_utils.getDetBoxes(score_text, score_link, 0.7, 0.4, 0.4, False)
	boxes = craft_utils.adjustResultCoordinates(boxes, ratio_w, ratio_h)
	polys = craft_utils.adjustResultCoordinates(polys, ratio_w, ratio_h)

	# Handle None polygons
	for i in range(len(polys)):
	if polys[i] is None:
	polys[i] = boxes[i]

	# Save the mask image
	mask_output_path = os.path.join(image_specific_folder, f"{image_name_no_ext}_mask.jpg")
	render_img = np.hstack((score_text, score_link))
	cv2.imwrite(mask_output_path, imgproc.cvt2HeatmapImg(render_img))

	# Save the detected text regions and results text file
	file_utils.saveResult(input_image_path, image[:, :, ::-1], polys, dirname=image_specific_folder)

	# Return the path to the bounding boxes text file
	print("\nCreating Text bboxes...")
	text_file_path = os.path.join(image_specific_folder, f"res_{image_name_no_ext}.txt")
	print(f"{len(boxes)} Text bounding boxes saved in: {text_file_path}")
	return text_file_path