Spaces:
Build error
Build error
| 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) |