refactor based on changes on the github project

fonts/fonts_animeace_i.ttf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a210dee89d1c9e27c5c6ebc2b8141fb9335ae751c773b137ce6b3252c72cac2
+size 28820

requirements.txt

@@ -6,4 +6,5 @@ googletrans==4.0.0-rc1
 deep-translator==1.11.4
 fastapi==0.110.3
 uvicorn==0.30.0
-
+opencv-python==4.9.0.80
+numpy==1.26.4

server.py

@@ -5,7 +5,9 @@ This file contains the FastAPI application that serves the web interface and han
 import os
 import io
 import base64
+from typing import Dict
 
+import numpy as np
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
@@ -16,9 +18,11 @@ import uvicorn
 from ultralytics import YOLO
 
 from utils.predict_bounding_boxes import predict_bounding_boxes
-from utils.manga_ocr import get_text_from_images
+from utils.manga_ocr import get_text_from_image
 from utils.translate_manga import translate_manga
-from utils.write_text_on_image import write_text
+from utils.process_contour import process_contour
+from utils.write_text_on_image import add_text
+
 
 # Load the object detection model
 best_model_path = "./model_creation/runs/detect/train5"
@@ -45,36 +49,42 @@ def home(request: Request):
 
 
 @app.post("/predict")
-def predict(request: dict):
+def predict(request: Dict):
 	image = request["image"]
 
 	# Decode base64-encoded image
 	image = base64.b64decode(image)
 	image = Image.open(io.BytesIO(image))
+	image_path = "image.png"
+	translated_image_path = "translated_image.png"
 
-	# Save the image
-	image.save("image.jpg")
-
-	# Perform object detection
-	result = predict_bounding_boxes(object_detection_model, "image.jpg")
+	# Save the image locally
+	image.save(image_path)
 
-	# Extract text from images
-	text_list = get_text_from_images("./bounding_box_images")
+	results = predict_bounding_boxes(object_detection_model, image_path)
+	image = np.array(image)
 
-	# Translate the manga
-	translated_text_list = translate_manga(text_list)
+	for result in results:
+			x1, y1, x2, y2, _, _ = result
+			detected_image = image[int(y1):int(y2), int(x1):int(x2)]
+			im = Image.fromarray(np.uint8((detected_image)*255))
+			text = get_text_from_image(im)
+			detected_image, cont = process_contour(detected_image)
+			text_translated = translate_manga(text)
+			add_text(detected_image, text_translated, cont)
 
-	# Write translation text on image
-	translated_image = write_text(result, translated_text_list, "image.jpg")
+	# Display the translated image
+	result_image = Image.fromarray(image, 'RGB')
+	result_image.save(translated_image_path)
 
 	# Convert the image to base64
 	buff = io.BytesIO()
-	translated_image.save(buff, format="JPEG")
+	result_image.save(buff, format="PNG")
 	img_str = base64.b64encode(buff.getvalue()).decode("utf-8")
 
 	# Clean up
-	os.remove("image.jpg")
-	os.remove("translated_image.png")
+	os.remove(image_path)
+	os.remove(translated_image_path)
 
 	return {"image": img_str}
 

utils/manga_ocr.py

@@ -1,22 +1,14 @@
 """
-This module is used to extract text from images using pytesseract.
+This module is used to extract text from images using manga_ocr.
 """
 
-import os
-from typing import List
 from manga_ocr import MangaOcr
 
-def get_text_from_images(bounding_box_images_path: str) -> List[str]:
+
+def get_text_from_image(image):
 	"""
-	Extract text from images using pytesseract.
+	Extract text from images using manga_ocr.
 	"""
 	mocr = MangaOcr()
 
-	text_list = []
-
-	for image_path in os.listdir(bounding_box_images_path):
-		image_path = os.path.join(bounding_box_images_path, image_path)
-		text = mocr(image_path)
-		text_list.append(text)
-
-	return text_list
+	return mocr(image)

utils/predict_bounding_boxes.py

@@ -4,10 +4,11 @@ in images using the trained Object Detection model.
 """
 import uuid
 import os
+from typing import List
 from PIL import Image
 from ultralytics import YOLO
 
-def predict_bounding_boxes(model: YOLO, image_path: str) -> YOLO:
+def predict_bounding_boxes(model: YOLO, image_path: str) -> List:
 	"""
 	Predict bounding boxes for text in images using the trained Object Detection model.
 	"""
@@ -36,4 +37,4 @@ def predict_bounding_boxes(model: YOLO, image_path: str) -> YOLO:
 		# save each image under a unique name
 		cropped_image.save(f"{bounding_box_images_path}/{uuid.uuid4()}.png")
 	
-	return result
+	return result.boxes.data.tolist()

utils/process_contour.py

@@ -0,0 +1,24 @@
+"""
+This module contains the function to process the contour in the image.
+"""
+from typing import Tuple
+import cv2
+import numpy as np
+
+
+def process_contour(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Process the contour in the image.
+    """
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    _, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
+
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    largest_contour = max(contours, key=cv2.contourArea)
+
+    mask = np.zeros_like(gray)
+    cv2.drawContours(mask, [largest_contour], -1, 255, cv2.FILLED)
+
+    image[mask == 255] = (255, 255, 255)
+
+    return image, largest_contour

utils/translate_manga.py

@@ -1,17 +1,15 @@
 """
 This module is used to translate manga from Japanese to English.
 """
-from typing import List
+
 from deep_translator import GoogleTranslator
 
-def translate_manga(text_list: List) -> List:
+def translate_manga(text: str) -> str:
 	"""
 	Translate manga from Japanese to English.
 	"""
-	for i, text in enumerate(text_list):
-		translated_text = GoogleTranslator(source="ja", target="en").translate(text)
-		print("Original text:", text)
-		print("Translated text:", translated_text)
-		text_list[i] = translated_text
+	translated_text = GoogleTranslator(source="ja", target="en").translate(text)
+	print("Original text:", text)
+	print("Translated text:", translated_text)
 	
-	return text_list
+	return translated_text

utils/write_text_on_image.py

@@ -1,227 +1,59 @@
 """
-This module is used to write text on an image.
+This module contains a function to add text to an image with a bounding box.
 """
-import os
-from typing import List, Tuple
-from PIL import Image, ImageFont, ImageDraw
-from PIL.ImageFont import FreeTypeFont
-from ultralytics import YOLO
+import textwrap
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import cv2
 
 
-def get_font(image: Image, text: str, width: int, height: int) -> Tuple[FreeTypeFont, int, int]:
+def add_text(image: np.ndarray, text: str, contour: np.ndarray):
 	"""
-	Get the font size and position for the text.
+	Add text to an image with a bounding box.
 	"""
 
-	# Default values at start
-	font_size = None 
-	font = None
-	box = None  # For version 8.0.0
-	x = 0
-	y = 0
+	font_path = "./fonts/fonts_animeace_i.ttf"
+	pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+	draw = ImageDraw.Draw(pil_image)
 
-	draw = ImageDraw.Draw(image)  # Create a draw object
+	x, y, w, h = cv2.boundingRect(contour)
 
-	# Test for different font sizes
-	for size in range(1, 500):
+	line_height = 16
+	font_size = 14
+	wrapping_ratio = 0.075
 
-		# Create new font
-		new_font = ImageFont.load_default(size=font_size)
+	wrapped_text = textwrap.fill(text, width=int(w * wrapping_ratio),
+									break_long_words=True)
 
-		# Calculate bbox for version 8.0.0
-		new_box = draw.textbbox((0, 0), text, font=new_font)
+	font = ImageFont.truetype(font_path, size=font_size)
 
-		# Calculate width and height
-		new_w = new_box[2] - new_box[0]  # Bottom - Top
-		new_h = new_box[3] - new_box[1]  # Right - Left
+	lines = wrapped_text.split('\n')
+	total_text_height = (len(lines)) * line_height
 
-		# If too big then exit with previous values
-		if new_w > width or new_h > height:
-			break
+	while total_text_height > h:
+		line_height -= 2
+		font_size -= 2
+		wrapping_ratio += 0.025
 
-		# Set new current values as current values
-		font_size = size
-		font = new_font
-		box = new_box
-		w = new_w
-		h = new_h
+		wrapped_text = textwrap.fill(text, width=int(w * wrapping_ratio),
+										break_long_words=True)
 
-		# Calculate position (minus margins in box)
-		x = (width - w) // 2 - box[0]  # Minus left margin
-		y = (height - h) // 2 - box[1]  # Minus top margin
+		font = ImageFont.truetype(font_path, size=font_size)
 
-	return font, x, y
+		lines = wrapped_text.split('\n')
+		total_text_height = (len(lines)) * line_height
 
+	# Vertical centering
+	text_y = y + (h - total_text_height) // 2
 
-def add_discoloration(color: Tuple[int, int, int], strength: int) -> Tuple[int, int, int]:
-	"""
-	Add discoloration to the background color.
-	"""
-
-	r, g, b = color
-	r = max(0, min(255, r + strength))
-	g = max(0, min(255, g + strength))
-	b = max(0, min(255, b + strength))
-
-	if r == 255 and g == 255 and b == 255:
-		r, g, b = 245, 245, 245
-
-	return (r, g, b)
-
-
-def get_background_color(image: Image, x_min: float,
-						y_min: float, x_max: float, y_max: float) -> Tuple[int, int, int]:
-	"""
-	Determine the background color of the text box.
-	"""
-
-	margin = 10
-	edge_region = image.crop(
-		(x_min - margin, y_min - margin, x_max + margin, y_max + margin))
-
-	# Find the most common color in the cropped region
-	edge_colors = edge_region.getcolors(
-		edge_region.size[0] * edge_region.size[1])
-	background_color = max(edge_colors, key=lambda x: x[0])[1]
-
-	# Add a bit of discoloration to the background color
-	background_color = add_discoloration(background_color, 40)
-
-	return background_color
-
-
-def get_text_fill_color(background_color: Tuple[int, int, int]):
-	"""
-	Determine the text color based on the background color.
-	"""
-	luminance = (
-		0.299 * background_color[0]
-		+ 0.587 * background_color[1]
-		+ 0.114 * background_color[2]
-	) / 255
-
-	# Determine the text color based on the background luminance
-	if luminance > 0.5:
-		return "black"
-	else:
-		return "white"
-
-
-def translated_words_fit(translated_text: str, font: ImageFont,
-						width: float, draw: ImageDraw) -> Tuple[bool, float]:
-	"""
-	Check if the translated words fit within the bounding box.
-	"""
-	words = translated_text.split()
-	total_width = 0
-
-	for word in words:
-		bbox = draw.textbbox((0, 0), word, font=font)
-		word_width = bbox[2] - bbox[0]
-		total_width += word_width
-
-	return total_width <= width, total_width
-
-
-def recalculate_words_and_size(translated_text: List, font: ImageFont,
-							image: Image, total_width: float,
-							draw: ImageDraw, x_min: float, y_min: float,
-							x_max: float, y_max: float) -> List:
-	"""
-	Recalculate the words and font size to fit within the bounding box.
-	"""
-	words = translated_text.split()
-	even_split = total_width // len(words)
-	lines = []
-	current_line = ""
-
-	for word in words:
-		bbox = draw.textbbox((0, 0), word, font=font)
-		word_width = bbox[2] - bbox[0]
-		if len(current_line) + word_width > even_split:
-			lines.append(current_line)
-			current_line = word
-		else:
-			current_line += " " + word
-	lines.append(current_line)
-	translated_text = "\n".join(lines)
-
-	# Calculate font size and position again
-	font, x, y = get_font(image, translated_text, x_max - x_min, y_max - y_min)
-	return translated_text, font, x, y
-
-
-def replace_text_with_translation(image: Image, text_boxes: List, translated_texts: List) -> Image:
-	"""
-	Replace the text in the bounding boxes with the translated text.
-	"""
-
-	draw = ImageDraw.Draw(image)
-
-	for text_box, translated_text in zip(text_boxes, translated_texts):
-		x_min, y_min, x_max, y_max = text_box
-
-		if translated_text is None:
-			continue
-
-		# Find the most common color in the text region
-		background_color = get_background_color(
-			image, x_min, y_min, x_max, y_max)
-
-		# Draw a rectangle to cover the text region with the original background color
-		draw.rectangle(((x_min, y_min), (x_max, y_max)), fill=background_color)
-
-		# Calculate font size and position
-		font, x, y = get_font(image, translated_text,
-							x_max - x_min, y_max - y_min)
-
-		# Split the words if they are too long
-		fits, total_width = translated_words_fit(
-			translated_text, font, x_max - x_min, draw)
-
-		if not fits or font.size < 30:
-			translated_text, font, x, y = recalculate_words_and_size(
-				translated_text, font, image, total_width, draw, x_min, y_min, x_max, y_max)
-
-		# Use local font
-		# Get full path to the font file
-		font_path = os.path.join(os.path.dirname(__file__), "../fonts/mangat.ttf")
-		font = ImageFont.truetype(
-			font=font_path, size=font.size, encoding="unic")
-
-		# Draw the translated text within the box
-		draw.text(
-			(x_min + x, y_min + y),
-			translated_text,
-			fill=get_text_fill_color(background_color),
-			font=font,
-		)
-
-	return image
-
-
-def write_text(result: YOLO, text_list: List[str], image_path: str) -> None:
-	"""
-	Replace the text in the bounding boxes with the translated text.
-	"""
-
-	all_boxes = []
-	all_translated_texts = []
-
-	for i, box in enumerate(result.boxes):
-		coords = [round(x) for x in box.xyxy[0].tolist()]
-		text = text_list[i]
-
-		all_boxes.append(coords)
-		all_translated_texts.append(text)
+	for line in lines:
+		text_length = draw.textlength(line, font=font)
 
-	image = Image.open(image_path)
+		# Horizontal centering
+		text_x = x + (w - text_length) // 2
 
-	# Draw text within all bounding boxes
-	image = replace_text_with_translation(
-		image, all_boxes, all_translated_texts)
+		draw.text((text_x, text_y), line, font=font, fill=(0, 0, 0))
 
-	# save the image
-	image.save("translated_image.png")
+		text_y += line_height
 
-	return image
+	image[:, :, :] = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)