main.py

import io
import base64

import numpy as np
from PIL import Image
from ultralytics import YOLO

from utils.predict_bounding_boxes import predict_bounding_boxes
from utils.manga_ocr_utils import get_text_from_image
from utils.translate_manga import translate_manga
from utils.process_contour import process_contour
from utils.write_text_on_image import add_text

MODEL_PATH = "./model_creation/runs/detect/train5/weights/best.pt"
object_detection_model = YOLO(MODEL_PATH)

def extract_text_from_regions(image: np.ndarray, results: list):

	for result in results:
		x1, y1, x2, y2, _, _ = result
		detected_image = image[int(y1):int(y2), int(x1):int(x2)]
		if detected_image.shape[-1] == 4:
			detected_image = detected_image[:, :, :3]
		im = Image.fromarray(np.uint8(detected_image * 255))
		text = get_text_from_image(im)

		processed_image, cont = process_contour(detected_image)
		translated_text = translate_manga(text, source_lang="auto", target_lang="en")
		add_text(processed_image, translated_text, cont)


def convert_image_to_base64(image: Image.Image) -> str:
	buff = io.BytesIO()
	image.save(buff, format="PNG")
	return base64.b64encode(buff.getvalue()).decode("utf-8")


def predict(image: np.ndarray):

	image = Image.fromarray(image)
	image.save("image.png")

	try:
		np_image = np.array(image)

		results = predict_bounding_boxes(object_detection_model, "image.png")
		extract_text_from_regions(np_image, results)

		return np_image

	except Exception as e:
		print(f"Error: {str(e)}")
		return None