comic_panel_extractor/llm_panel_extractor.py

from .config import Config
from ultralytics import YOLO 
from PIL import Image 
import cv2
from . import constant
from . import utils
import os
import shutil
import requests
from pathlib import Path
from . import common

class LLMPanelExtractor:
	"""Handles image preprocessing operations."""
	
	def __init__(self, config: Config = None):
		self.config = config or Config()

		# Check if YOLO model exists; if not, download it to the specified path
		yolo_base_model_path = f'{self.config.yolo_base_model_path}_best.pt'
		# yolo_base_model_path = f'{self.config.yolo_trained_model_path}'
		if not os.path.exists(yolo_base_model_path):
			url = "https://huggingface.co/mosesb/best-comic-panel-detection/resolve/main/best.pt"
			print(f"Downloading YOLO model to {yolo_base_model_path}...")
			response = requests.get(url)
			response.raise_for_status()  # Raise an error if the download fails
			with open(yolo_base_model_path, "wb") as f:
				f.write(response.content)
			print("YOLO model downloaded successfully.")

		self.yolo_model = YOLO(yolo_base_model_path)
		os.makedirs(self.config.output_folder, exist_ok=True)

	def extract_bounding_boxes(self, detection_result_boxes):
		"""Extract bounding box coordinates from YOLO detection results."""
		bounding_boxes = []
		for detection_box in detection_result_boxes.xyxy:
			# Extract coordinates
			x_min, y_min, x_max, y_max = map(int, detection_box)
			bounding_boxes.append((x_min, y_min, x_max, y_max))

		return bounding_boxes

	def crop_and_save_detected_panels(self, detected_boxes):
		"""Crop detected boxes and save them in separate folders"""
		if len(detected_boxes) == 0:
			print(f"No boxes detected for {self.config.org_input_path}")
			return

		source_image = cv2.imread(self.config.org_input_path)
		for box_coordinates in detected_boxes:
			# Extract coordinates
			x_min, y_min, x_max, y_max = box_coordinates
			
			# Crop the image
			cropped_panel = source_image[y_min:y_max, x_min:x_max]
			
			# Save cropped image
			constant.INDEX += 1
			panel_output_path = f"{self.config.output_folder}/{constant.INDEX:04d}_panel_{x_min, y_min, x_max, y_max}.jpg"
			cv2.imwrite(panel_output_path, cropped_panel)

	def pre_all_processed_boxes(self, all_processed_boxes, image_width, image_height):
		all_processed_boxes = utils.extend_boxes_to_image_border(
			all_processed_boxes, 
			(image_width, image_height), 
			self.config.min_width_ratio, 
			self.config.min_height_ratio
		)
		all_processed_boxes = sorted(all_processed_boxes, key=lambda box: (box[1], box[0]))  # sort by y_min, then x_min
		all_processed_boxes = utils.extend_to_nearby_boxes(
			all_processed_boxes, 
			(image_width, image_height), 
			self.config.min_width_ratio, 
			self.config.min_height_ratio
		)
		return all_processed_boxes

	def detect_and_extract_panels(self, input_image_path=None, existing_boxes=None, confidence_threshold=0.9):
		"""Main method to detect and extract panels from an image."""
		if not input_image_path:
			input_image_path = self.config.org_input_path

		# Get image dimensions
		with Image.open(input_image_path) as input_image:
			image_width, image_height = input_image.size

		# Run YOLO detection
		detection_results = self.yolo_model.predict(source=input_image_path, device=common.get_device())
		first_detection_result = detection_results[0]
		newly_detected_boxes = None
		all_processed_boxes = []
		
		# Add existing boxes if provided
		if existing_boxes:
			all_processed_boxes.extend(existing_boxes)
		
		# Filter boxes by confidence threshold
		if first_detection_result.boxes is not None:
			high_confidence_filter = first_detection_result.boxes.conf >= confidence_threshold
			if high_confidence_filter.sum() > 0:
				first_detection_result.boxes = first_detection_result.boxes[high_confidence_filter]
				newly_detected_boxes = self.extract_bounding_boxes(first_detection_result.boxes)
				newly_detected_boxes = utils.is_valid_panel((image_width, image_height), newly_detected_boxes, self.config.min_width_ratio, self.config.min_height_ratio)
				if newly_detected_boxes:
					all_processed_boxes.extend(self.extract_bounding_boxes(first_detection_result.boxes))

					# Process and extend boxes
					all_processed_boxes = self.pre_all_processed_boxes(all_processed_boxes, image_width, image_height)

					# Crop and save detected panels
					self.crop_and_save_detected_panels(newly_detected_boxes)
					
					# Save prediction visualization
					visualization_result = first_detection_result.plot(masks=False)
					constant.INDEX += 1
					debug_output_path = f"{self.config.output_folder}/{constant.INDEX:04d}_debug.jpg"
					Image.fromarray(visualization_result[..., ::-1]).save(debug_output_path)
					
					# Create black and white mask
					constant.INDEX += 1
					masked_output_path = f"{self.config.output_folder}/{constant.INDEX:04d}_draw_black.jpg"
					masked_image_path = utils.draw_black(self.config.org_input_path, all_processed_boxes, masked_output_path, stripe=False)
					return masked_image_path, newly_detected_boxes

		# Process boxes even if no new detections
		all_processed_boxes = self.pre_all_processed_boxes(all_processed_boxes, image_width, image_height)

		constant.INDEX += 1
		masked_output_path = f"{self.config.output_folder}/{constant.INDEX:04d}_draw_black.jpg"
		masked_image_path = utils.draw_black(self.config.org_input_path, all_processed_boxes, masked_output_path, stripe=False)
		return masked_image_path, newly_detected_boxes

	def check_for_remaining_similarity(self, current_processed_image_path, existing_boxes):
		# Get image dimensions
		with Image.open(self.config.org_input_path) as input_image:
			image_width, image_height = input_image.size

		all_processed_boxes = self.pre_all_processed_boxes(existing_boxes, image_width, image_height)

		constant.INDEX += 1
		similar_remaining_regions_path = f"{self.config.output_folder}/{constant.INDEX:04d}_remaining_similarity_debug.jpg"
		similar_remaining_box = utils.find_similar_remaining_regions(all_processed_boxes, (image_width, image_height), similar_remaining_regions_path)
		if similar_remaining_box:
			similar_remaining_box = utils.is_valid_panel((image_width, image_height), similar_remaining_box, self.config.min_width_ratio, self.config.min_height_ratio)
			if similar_remaining_box:
				self.crop_and_save_detected_panels(similar_remaining_box)
				existing_boxes.extend(similar_remaining_box)

				all_processed_boxes = self.pre_all_processed_boxes(existing_boxes, image_width, image_height)

				constant.INDEX += 1
				current_processed_image_path = f"{self.config.output_folder}/{constant.INDEX:04d}_remaining_similarity_draw_black.jpg"
				current_processed_image_path = utils.draw_black(self.config.org_input_path, all_processed_boxes, current_processed_image_path, stripe=False)

		return current_processed_image_path, existing_boxes

def extract_panel_via_llm(input_image_path, config=None, reset=True):
	"""Main function to extract panels using various image processing techniques."""
	# Initialize configuration
	extractor_config = config or Config()
	extractor_config.org_input_path = input_image_path

	# Clean output folder
	if reset:
		if Path(extractor_config.output_folder).exists():
			shutil.rmtree(extractor_config.output_folder, ignore_errors=True)
		Path(extractor_config.output_folder).mkdir(exist_ok=True)

	# Initialize extractor
	panel_extractor = LLMPanelExtractor(extractor_config)

	current_processed_image_path = extractor_config.org_input_path
	accumulated_detected_boxes = []
	all_processed_boxes = []
	
	# Get original image dimensions
	with Image.open(current_processed_image_path) as original_image:
		original_width, original_height = original_image.size

	# Define image processing techniques to try
	processing_techniques = [
		{
			'name': 'clahe',
			'function': utils.convert_to_clahe,
			'confidence_level': 1.0,
			'description': 'CLAHE (Contrast Limited Adaptive Histogram Equalization)'
		},
		{
			'name': 'grayscale',
			'function': utils.convert_to_grayscale_pil,
			'confidence_level': 1.0,
			'description': 'Grayscale conversion'
		},
		{
			'name': 'lab_l',
			'function': utils.convert_to_lab_l,
			'confidence_level': 1.0,
			'description': 'LAB L-channel extraction'
		},
		{
			'name': 'group_color',
			'function': utils.convert_to_group_colors,
			'confidence_level': 0.1,
			'image_path': extractor_config.org_input_path,
			'description': 'Group Color extraction'
		}
	]

	# Process with different techniques until white ratio threshold is met
	for technique in processing_techniques:
		iteration_count = 0
		confidence_level = technique["confidence_level"]
		if technique.get("image_path", None) and utils.box_covered_ratio(panel_extractor.pre_all_processed_boxes(accumulated_detected_boxes, original_width, original_height), (original_width, original_height)) < 0.95:
			current_processed_image_path = technique.get("image_path")
		
		while (utils.box_covered_ratio(panel_extractor.pre_all_processed_boxes(accumulated_detected_boxes, original_width, original_height), (original_width, original_height)) < 0.95 and confidence_level > 0):
			
			print(f"\n{technique['description']} process iteration: {iteration_count} confidence level: {confidence_level}")
			iteration_count += 1
			confidence_level -= 0.1
			
			# Apply image processing technique
			constant.INDEX += 1
			processed_output_path = f"{extractor_config.output_folder}/{constant.INDEX:04d}_convert_to_{technique['name']}.jpg"
			current_processed_image_path = technique['function'](current_processed_image_path, processed_output_path)
			
			# Run panel detection on processed image
			current_processed_image_path, newly_detected_boxes = panel_extractor.detect_and_extract_panels(
				input_image_path=current_processed_image_path, 
				existing_boxes=accumulated_detected_boxes, 
				confidence_threshold=confidence_level
			)
			if newly_detected_boxes:
				accumulated_detected_boxes.extend(newly_detected_boxes)

			current_processed_image_path, accumulated_detected_boxes = panel_extractor.check_for_remaining_similarity(current_processed_image_path, accumulated_detected_boxes)
			all_processed_boxes = panel_extractor.pre_all_processed_boxes(accumulated_detected_boxes, original_width, original_height)

	remain_boxes = utils.get_remaining_areas((original_width, original_height), all_processed_boxes)
	if remain_boxes:
		remain_boxes = utils.is_valid_panel((original_width, original_height), remain_boxes, extractor_config.min_width_ratio, extractor_config.min_height_ratio)
		if remain_boxes:
			panel_extractor.crop_and_save_detected_panels(remain_boxes)
			all_processed_boxes.extend(remain_boxes)
			accumulated_detected_boxes.extend(remain_boxes)

	all_path = [file for file in os.listdir(extractor_config.output_folder) if "_panel_" in file]
	print(f"Processing complete. Final result saved to: {extractor_config.output_folder}")
	print(f"Total panels detected: {len(all_path)}")
	return all_path, accumulated_detected_boxes, all_processed_boxes


if __name__ == "__main__":
	import argparse

	# Parse command-line arguments
	argument_parser = argparse.ArgumentParser(description="Run panel extractor on an image")
	argument_parser.add_argument("--input", type=str, required=True, help="Path to input image")
	parsed_arguments = argument_parser.parse_args()
	
	final_result_path, total_detected_boxes = extract_panel_via_llm(parsed_arguments.input)