Spaces:
Running
Running
path isue fix
Browse files- comic_panel_extractor/__init__.py +1 -16
- comic_panel_extractor/annorator_server.py +5 -7
- comic_panel_extractor/config.py +24 -21
- comic_panel_extractor/extractor_server.py +1 -3
- comic_panel_extractor/inference.py +2 -1
- comic_panel_extractor/server.py +2 -2
- comic_panel_extractor/train.py +126 -4
- comic_panel_extractor/utils.py +40 -47
- comic_panel_extractor/yolo_manager.py +5 -18
comic_panel_extractor/__init__.py
CHANGED
|
@@ -1,16 +1 @@
|
|
| 1 |
-
|
| 2 |
-
from .config import Config
|
| 3 |
-
from .text_detector import TextDetector, TextDetection
|
| 4 |
-
from .image_processor import ImageProcessor
|
| 5 |
-
from .panel_extractor import PanelExtractor, PanelData
|
| 6 |
-
|
| 7 |
-
__version__ = "0.1.0"
|
| 8 |
-
__all__ = [
|
| 9 |
-
"ComicPanelExtractor",
|
| 10 |
-
"Config",
|
| 11 |
-
"TextDetector",
|
| 12 |
-
"TextDetection",
|
| 13 |
-
"ImageProcessor",
|
| 14 |
-
"PanelExtractor",
|
| 15 |
-
"PanelData"
|
| 16 |
-
]
|
|
|
|
| 1 |
+
__version__ = "0.1.0"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
comic_panel_extractor/annorator_server.py
CHANGED
|
@@ -7,15 +7,14 @@ import os
|
|
| 7 |
import base64
|
| 8 |
from io import BytesIO
|
| 9 |
import shutil
|
| 10 |
-
|
| 11 |
-
current_path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
|
| 12 |
|
| 13 |
app = APIRouter()
|
| 14 |
|
| 15 |
# === Configuration ===
|
| 16 |
-
IMAGE_ROOT = os.path.join(current_path, "dataset/images")
|
| 17 |
-
LABEL_ROOT = os.path.join(current_path, "dataset/labels")
|
| 18 |
-
IMAGE_LABEL_ROOT = os.path.join(current_path, "image_labels")
|
| 19 |
|
| 20 |
CLASS_ID = 0
|
| 21 |
|
|
@@ -62,9 +61,8 @@ def load_yolo_boxes(image_path: str, label_path: str, detect: bool = False):
|
|
| 62 |
boxes = []
|
| 63 |
if detect and not os.path.exists(label_path):
|
| 64 |
from .yolo_manager import YOLOManager
|
| 65 |
-
from .utils import Config
|
| 66 |
with YOLOManager() as yolo_manager:
|
| 67 |
-
weights_path = f'{current_path}/{Config.YOLO_MODEL_NAME}.pt'
|
| 68 |
|
| 69 |
yolo_manager.load_model(weights_path)
|
| 70 |
|
|
|
|
| 7 |
import base64
|
| 8 |
from io import BytesIO
|
| 9 |
import shutil
|
| 10 |
+
from .config import Config
|
|
|
|
| 11 |
|
| 12 |
app = APIRouter()
|
| 13 |
|
| 14 |
# === Configuration ===
|
| 15 |
+
IMAGE_ROOT = os.path.join(Config.current_path, "dataset/images")
|
| 16 |
+
LABEL_ROOT = os.path.join(Config.current_path, "dataset/labels")
|
| 17 |
+
IMAGE_LABEL_ROOT = os.path.join(Config.current_path, "image_labels")
|
| 18 |
|
| 19 |
CLASS_ID = 0
|
| 20 |
|
|
|
|
| 61 |
boxes = []
|
| 62 |
if detect and not os.path.exists(label_path):
|
| 63 |
from .yolo_manager import YOLOManager
|
|
|
|
| 64 |
with YOLOManager() as yolo_manager:
|
| 65 |
+
weights_path = f'{Config.current_path}/{Config.YOLO_MODEL_NAME}.pt'
|
| 66 |
|
| 67 |
yolo_manager.load_model(weights_path)
|
| 68 |
|
comic_panel_extractor/config.py
CHANGED
|
@@ -1,27 +1,30 @@
|
|
| 1 |
from dataclasses import dataclass
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
# Path to this script's directory
|
| 5 |
-
CURRENT_DIR = Path(__file__).parent.resolve()
|
| 6 |
|
| 7 |
@dataclass
|
| 8 |
class Config:
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 25 |
|
| 26 |
def get_text_cood_file_path(config: Config):
|
| 27 |
-
|
|
|
|
| 1 |
from dataclasses import dataclass
|
| 2 |
+
import os
|
|
|
|
|
|
|
|
|
|
| 3 |
|
| 4 |
@dataclass
|
| 5 |
class Config:
|
| 6 |
+
"""Configuration settings for the comic-to-video pipeline."""
|
| 7 |
+
org_input_path: str = ""
|
| 8 |
+
input_path: str = ""
|
| 9 |
+
current_path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
|
| 10 |
+
YOLO_MODEL_NAME = os.getenv('YOLO_MODEL_NAME', 'default_model')
|
| 11 |
+
yolo_model_path: str = f'{current_path}/{YOLO_MODEL_NAME}.pt'
|
| 12 |
+
black_overlay_input_path: str = ""
|
| 13 |
+
output_folder: str = "temp_dir"
|
| 14 |
+
distance_threshold: int = 70
|
| 15 |
+
vertical_threshold: int = 30
|
| 16 |
+
text_cood_file_name: str = "detect_and_group_text.json"
|
| 17 |
+
min_text_length: int = 2
|
| 18 |
+
min_area_ratio: float = 0.05
|
| 19 |
+
min_width_ratio: float = 0.15
|
| 20 |
+
min_height_ratio: float = 0.15
|
| 21 |
+
|
| 22 |
+
# Additional parameters for BorderPanelExtractor
|
| 23 |
+
panel_filename_pattern: str = r"panel_\d+_\((\d+), (\d+), (\d+), (\d+)\)\.jpg"
|
| 24 |
+
|
| 25 |
+
"""Configuration class to manage environment variables and paths."""
|
| 26 |
+
DEFAULT_IMAGE_SIZE = 640
|
| 27 |
+
SUPPORTED_EXTENSIONS = ['jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG']
|
| 28 |
|
| 29 |
def get_text_cood_file_path(config: Config):
|
| 30 |
+
return f'{config.output_folder}/{config.text_cood_file_name}'
|
comic_panel_extractor/extractor_server.py
CHANGED
|
@@ -9,10 +9,8 @@ import shutil
|
|
| 9 |
import time
|
| 10 |
import mimetypes
|
| 11 |
|
| 12 |
-
current_path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
|
| 13 |
-
|
| 14 |
base_output_folder = "api_outputs"
|
| 15 |
-
output_folder = os.path.join(current_path, base_output_folder)
|
| 16 |
|
| 17 |
app = APIRouter()
|
| 18 |
|
|
|
|
| 9 |
import time
|
| 10 |
import mimetypes
|
| 11 |
|
|
|
|
|
|
|
| 12 |
base_output_folder = "api_outputs"
|
| 13 |
+
output_folder = os.path.join(Config.current_path, base_output_folder)
|
| 14 |
|
| 15 |
app = APIRouter()
|
| 16 |
|
comic_panel_extractor/inference.py
CHANGED
|
@@ -1,7 +1,8 @@
|
|
| 1 |
# inference.py
|
| 2 |
from yolo_manager import YOLOManager
|
| 3 |
-
from utils import
|
| 4 |
import os
|
|
|
|
| 5 |
|
| 6 |
def run_inference(weights_path: str, images_dirs, output_dir: str = 'temp_dir') -> None:
|
| 7 |
"""
|
|
|
|
| 1 |
# inference.py
|
| 2 |
from yolo_manager import YOLOManager
|
| 3 |
+
from utils import get_abs_path, get_image_paths
|
| 4 |
import os
|
| 5 |
+
from .config import Config
|
| 6 |
|
| 7 |
def run_inference(weights_path: str, images_dirs, output_dir: str = 'temp_dir') -> None:
|
| 8 |
"""
|
comic_panel_extractor/server.py
CHANGED
|
@@ -3,6 +3,7 @@ from fastapi.staticfiles import StaticFiles
|
|
| 3 |
from .extractor_server import app as extractor_app, delete_folder_if_old_or_empty, output_folder
|
| 4 |
from .annorator_server import app as annotator_app
|
| 5 |
import os
|
|
|
|
| 6 |
|
| 7 |
from fastapi import Request
|
| 8 |
from fastapi.responses import HTMLResponse
|
|
@@ -13,8 +14,7 @@ from jinja2 import Environment, FileSystemLoader, select_autoescape
|
|
| 13 |
fast_api = FastAPI()
|
| 14 |
|
| 15 |
# Mount static files ONCE
|
| 16 |
-
|
| 17 |
-
static_folder = os.path.join(current_path, "static")
|
| 18 |
fast_api.mount("/static", StaticFiles(directory=static_folder), name="static")
|
| 19 |
|
| 20 |
fast_api.include_router(extractor_app)
|
|
|
|
| 3 |
from .extractor_server import app as extractor_app, delete_folder_if_old_or_empty, output_folder
|
| 4 |
from .annorator_server import app as annotator_app
|
| 5 |
import os
|
| 6 |
+
from .config import Config
|
| 7 |
|
| 8 |
from fastapi import Request
|
| 9 |
from fastapi.responses import HTMLResponse
|
|
|
|
| 14 |
fast_api = FastAPI()
|
| 15 |
|
| 16 |
# Mount static files ONCE
|
| 17 |
+
static_folder = os.path.join(Config.current_path, "static")
|
|
|
|
| 18 |
fast_api.mount("/static", StaticFiles(directory=static_folder), name="static")
|
| 19 |
|
| 20 |
fast_api.include_router(extractor_app)
|
comic_panel_extractor/train.py
CHANGED
|
@@ -1,7 +1,98 @@
|
|
| 1 |
# train.py
|
| 2 |
-
from yolo_manager import YOLOManager
|
| 3 |
-
from utils import
|
| 4 |
import os
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 5 |
|
| 6 |
def main():
|
| 7 |
"""Main training function."""
|
|
@@ -10,7 +101,7 @@ def main():
|
|
| 10 |
yolo_manager = YOLOManager()
|
| 11 |
|
| 12 |
# Configuration
|
| 13 |
-
data_yaml_path =
|
| 14 |
|
| 15 |
if not os.path.isfile(data_yaml_path):
|
| 16 |
raise FileNotFoundError(f"❌ Dataset YAML not found: {data_yaml_path}")
|
|
@@ -37,5 +128,36 @@ def main():
|
|
| 37 |
print(f"❌ Training failed: {str(e)}")
|
| 38 |
raise
|
| 39 |
|
| 40 |
-
if __name__ == "__main__"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 41 |
main()
|
|
|
|
| 1 |
# train.py
|
| 2 |
+
from .yolo_manager import YOLOManager
|
| 3 |
+
from .utils import get_abs_path, backup_file
|
| 4 |
import os
|
| 5 |
+
from .config import Config
|
| 6 |
+
import yaml
|
| 7 |
+
import os
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
import shutil
|
| 10 |
+
|
| 11 |
+
def create_filtered_dataset(original_dataset_path, output_base_path):
|
| 12 |
+
"""
|
| 13 |
+
Create a filtered dataset with only images that have non-empty labels
|
| 14 |
+
"""
|
| 15 |
+
shutil.rmtree(f'{original_dataset_path}/filtered_dataset', ignore_errors=True)
|
| 16 |
+
original_path = Path(f'{original_dataset_path}/filtered_dataset')
|
| 17 |
+
output_path = Path(output_base_path)
|
| 18 |
+
|
| 19 |
+
# Create output directory structure
|
| 20 |
+
output_images = output_path / "images"
|
| 21 |
+
output_labels = output_path / "labels"
|
| 22 |
+
|
| 23 |
+
for split in ['train', 'val', 'test']:
|
| 24 |
+
(output_images / split).mkdir(parents=True, exist_ok=True)
|
| 25 |
+
(output_labels / split).mkdir(parents=True, exist_ok=True)
|
| 26 |
+
|
| 27 |
+
filtered_counts = {}
|
| 28 |
+
|
| 29 |
+
for split in ['train', 'val', 'test']:
|
| 30 |
+
original_images_dir = original_path / 'images' / split
|
| 31 |
+
original_labels_dir = original_path / 'labels' / split
|
| 32 |
+
|
| 33 |
+
output_images_dir = output_images / split
|
| 34 |
+
output_labels_dir = output_labels / split
|
| 35 |
+
|
| 36 |
+
if not original_images_dir.exists() or not original_labels_dir.exists():
|
| 37 |
+
print(f"Skipping {split} - source directory not found")
|
| 38 |
+
filtered_counts[split] = 0
|
| 39 |
+
continue
|
| 40 |
+
|
| 41 |
+
total_count = 0
|
| 42 |
+
copied_count = 0
|
| 43 |
+
|
| 44 |
+
# Process each image
|
| 45 |
+
for img_file in original_images_dir.glob('*'):
|
| 46 |
+
if img_file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']:
|
| 47 |
+
total_count += 1
|
| 48 |
+
label_file = original_labels_dir / f"{img_file.stem}.txt"
|
| 49 |
+
|
| 50 |
+
# Check if label file exists and has content
|
| 51 |
+
if label_file.exists():
|
| 52 |
+
with open(label_file, 'r') as f:
|
| 53 |
+
content = f.read().strip()
|
| 54 |
+
if content: # Label file has content
|
| 55 |
+
# Copy image
|
| 56 |
+
shutil.copy2(img_file, output_images_dir / img_file.name)
|
| 57 |
+
# Copy label
|
| 58 |
+
shutil.copy2(label_file, output_labels_dir / label_file.name)
|
| 59 |
+
copied_count += 1
|
| 60 |
+
else:
|
| 61 |
+
print(f"Skipping {img_file.name} - empty label file")
|
| 62 |
+
else:
|
| 63 |
+
print(f"Skipping {img_file.name} - no label file")
|
| 64 |
+
|
| 65 |
+
filtered_counts[split] = copied_count
|
| 66 |
+
print(f"{split.upper()} split: {copied_count}/{total_count} images copied")
|
| 67 |
+
|
| 68 |
+
return filtered_counts
|
| 69 |
+
|
| 70 |
+
def create_filtered_yaml(output_base_path, filtered_counts):
|
| 71 |
+
"""
|
| 72 |
+
Create the YAML file for the filtered dataset
|
| 73 |
+
"""
|
| 74 |
+
output_path = Path(output_base_path)
|
| 75 |
+
yaml_path = f'{Config.current_path}/filtered_comic.yaml'
|
| 76 |
+
|
| 77 |
+
# Create YAML structure
|
| 78 |
+
yaml_data = {
|
| 79 |
+
'names': ['panel'],
|
| 80 |
+
'nc': 1,
|
| 81 |
+
'path': str(output_path),
|
| 82 |
+
'train': str(output_path / 'images' / 'train'),
|
| 83 |
+
'val': str(output_path / 'images' / 'val')
|
| 84 |
+
}
|
| 85 |
+
|
| 86 |
+
# Only add test if it has images
|
| 87 |
+
if filtered_counts.get('test', 0) > 0:
|
| 88 |
+
yaml_data['test'] = str(output_path / 'images' / 'test')
|
| 89 |
+
|
| 90 |
+
# Write YAML file
|
| 91 |
+
with open(yaml_path, 'w') as f:
|
| 92 |
+
yaml.dump(yaml_data, f, default_flow_style=False, sort_keys=False)
|
| 93 |
+
|
| 94 |
+
print(f"\n✅ Created filtered dataset YAML: {yaml_path}")
|
| 95 |
+
return yaml_path
|
| 96 |
|
| 97 |
def main():
|
| 98 |
"""Main training function."""
|
|
|
|
| 101 |
yolo_manager = YOLOManager()
|
| 102 |
|
| 103 |
# Configuration
|
| 104 |
+
data_yaml_path = f'{Config.current_path}/filtered_comic.yaml'
|
| 105 |
|
| 106 |
if not os.path.isfile(data_yaml_path):
|
| 107 |
raise FileNotFoundError(f"❌ Dataset YAML not found: {data_yaml_path}")
|
|
|
|
| 128 |
print(f"❌ Training failed: {str(e)}")
|
| 129 |
raise
|
| 130 |
|
| 131 |
+
if __name__ == "__main__":# Configuration
|
| 132 |
+
# Configuration
|
| 133 |
+
original_dataset_path = "/home/jebineinstein/git/comic-panel-extractor/comic_panel_extractor/dataset"
|
| 134 |
+
output_base_path = "/home/jebineinstein/git/comic-panel-extractor/comic_panel_extractor"
|
| 135 |
+
|
| 136 |
+
print("🔍 Starting dataset filtering...")
|
| 137 |
+
print(f"📂 Source: {original_dataset_path}")
|
| 138 |
+
print(f"📁 Output: {output_base_path}")
|
| 139 |
+
|
| 140 |
+
# Create filtered dataset
|
| 141 |
+
filtered_counts = create_filtered_dataset(original_dataset_path, output_base_path)
|
| 142 |
+
|
| 143 |
+
# Create YAML file
|
| 144 |
+
yaml_path = create_filtered_yaml(output_base_path, filtered_counts)
|
| 145 |
+
|
| 146 |
+
# Summary
|
| 147 |
+
total_filtered = sum(filtered_counts.values())
|
| 148 |
+
print(f"\n📊 Filtering Summary:")
|
| 149 |
+
for split, count in filtered_counts.items():
|
| 150 |
+
if count > 0:
|
| 151 |
+
print(f" {split.upper()}: {count} images")
|
| 152 |
+
print(f" TOTAL: {total_filtered} images with labels")
|
| 153 |
+
|
| 154 |
+
print(f"\n🎯 Use this YAML for training: {yaml_path}")
|
| 155 |
+
|
| 156 |
+
# Display the created YAML content
|
| 157 |
+
with open(yaml_path, 'r') as f:
|
| 158 |
+
yaml_content = f.read()
|
| 159 |
+
print(f"\n📄 Generated YAML content:")
|
| 160 |
+
print("─" * 50)
|
| 161 |
+
print(yaml_content)
|
| 162 |
+
print("─" * 50)
|
| 163 |
main()
|
comic_panel_extractor/utils.py
CHANGED
|
@@ -9,6 +9,7 @@ from glob import glob
|
|
| 9 |
from typing import List, Union
|
| 10 |
from dotenv import load_dotenv
|
| 11 |
load_dotenv()
|
|
|
|
| 12 |
|
| 13 |
def remove_duplicate_boxes(boxes, compare_single=None, iou_threshold=0.7):
|
| 14 |
"""
|
|
@@ -483,57 +484,49 @@ def is_valid_panel(
|
|
| 483 |
|
| 484 |
return validity
|
| 485 |
|
| 486 |
-
|
| 487 |
-
class Config:
|
| 488 |
-
"""Configuration class to manage environment variables and paths."""
|
| 489 |
-
YOLO_MODEL_NAME = os.getenv('YOLO_MODEL_NAME', 'default_model')
|
| 490 |
-
DEFAULT_IMAGE_SIZE = 640
|
| 491 |
-
SUPPORTED_EXTENSIONS = ['jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG']
|
| 492 |
-
|
| 493 |
def get_abs_path(relative_path: str) -> str:
|
| 494 |
-
|
| 495 |
-
|
| 496 |
|
| 497 |
def get_image_paths(directories: Union[str, List[str]]) -> List[str]:
|
| 498 |
-
|
| 499 |
-
|
| 500 |
-
|
| 501 |
-
|
| 502 |
-
|
| 503 |
-
|
| 504 |
-
|
| 505 |
-
|
| 506 |
-
|
| 507 |
-
|
| 508 |
-
|
| 509 |
-
|
| 510 |
-
|
| 511 |
-
|
| 512 |
-
|
| 513 |
-
|
| 514 |
-
|
| 515 |
-
|
| 516 |
-
|
| 517 |
-
|
| 518 |
-
|
| 519 |
-
|
| 520 |
-
|
| 521 |
-
|
| 522 |
-
|
| 523 |
-
|
| 524 |
|
| 525 |
def clean_directory(directory: str, create_if_not_exists: bool = True) -> None:
|
| 526 |
-
|
| 527 |
-
|
| 528 |
-
|
| 529 |
-
|
| 530 |
-
|
| 531 |
-
os.makedirs(directory, exist_ok=True)
|
| 532 |
|
| 533 |
def backup_file(source_path: str, backup_path: str) -> str:
|
| 534 |
-
|
| 535 |
-
|
| 536 |
-
|
| 537 |
-
|
| 538 |
-
|
| 539 |
-
|
|
|
|
| 9 |
from typing import List, Union
|
| 10 |
from dotenv import load_dotenv
|
| 11 |
load_dotenv()
|
| 12 |
+
from .config import Config
|
| 13 |
|
| 14 |
def remove_duplicate_boxes(boxes, compare_single=None, iou_threshold=0.7):
|
| 15 |
"""
|
|
|
|
| 484 |
|
| 485 |
return validity
|
| 486 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 487 |
def get_abs_path(relative_path: str) -> str:
|
| 488 |
+
"""Convert relative path to absolute path."""
|
| 489 |
+
return os.path.abspath(relative_path)
|
| 490 |
|
| 491 |
def get_image_paths(directories: Union[str, List[str]]) -> List[str]:
|
| 492 |
+
"""
|
| 493 |
+
Get all image paths from given directories.
|
| 494 |
+
|
| 495 |
+
Args:
|
| 496 |
+
directories: Single directory path or list of directory paths
|
| 497 |
+
|
| 498 |
+
Returns:
|
| 499 |
+
List of image file paths
|
| 500 |
+
"""
|
| 501 |
+
if isinstance(directories, str):
|
| 502 |
+
directories = [directories]
|
| 503 |
+
|
| 504 |
+
all_images = []
|
| 505 |
+
for directory in directories:
|
| 506 |
+
abs_dir = get_abs_path(directory)
|
| 507 |
+
if not os.path.isdir(abs_dir):
|
| 508 |
+
print(f"⚠️ Warning: Skipping non-directory {abs_dir}")
|
| 509 |
+
continue
|
| 510 |
+
|
| 511 |
+
# Support multiple image extensions
|
| 512 |
+
for ext in Config.SUPPORTED_EXTENSIONS:
|
| 513 |
+
pattern = os.path.join(abs_dir, f'*.{ext}')
|
| 514 |
+
images = sorted(glob(pattern))
|
| 515 |
+
all_images.extend(images)
|
| 516 |
+
|
| 517 |
+
return list(set(all_images)) # Remove duplicates
|
| 518 |
|
| 519 |
def clean_directory(directory: str, create_if_not_exists: bool = True) -> None:
|
| 520 |
+
"""Clean directory contents or create if it doesn't exist."""
|
| 521 |
+
shutil.rmtree(directory, ignore_errors=True)
|
| 522 |
+
|
| 523 |
+
if create_if_not_exists:
|
| 524 |
+
os.makedirs(directory, exist_ok=True)
|
|
|
|
| 525 |
|
| 526 |
def backup_file(source_path: str, backup_path: str) -> str:
|
| 527 |
+
"""Backup a file to specified location."""
|
| 528 |
+
backup_path = get_abs_path(backup_path)
|
| 529 |
+
os.makedirs(os.path.dirname(backup_path), exist_ok=True)
|
| 530 |
+
shutil.copy(source_path, backup_path)
|
| 531 |
+
print(f"✅ File backed up to: {backup_path}")
|
| 532 |
+
return backup_path
|
comic_panel_extractor/yolo_manager.py
CHANGED
|
@@ -7,12 +7,6 @@ from dotenv import load_dotenv
|
|
| 7 |
|
| 8 |
load_dotenv()
|
| 9 |
|
| 10 |
-
class Config:
|
| 11 |
-
"""Configuration class to manage environment variables and paths."""
|
| 12 |
-
YOLO_MODEL_NAME = os.getenv('YOLO_MODEL_NAME', 'default_model')
|
| 13 |
-
DEFAULT_IMAGE_SIZE = 640
|
| 14 |
-
SUPPORTED_EXTENSIONS = ['jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG']
|
| 15 |
-
|
| 16 |
def get_abs_path(relative_path: str) -> str:
|
| 17 |
"""Convert relative path to absolute path."""
|
| 18 |
return os.path.abspath(relative_path)
|
|
@@ -45,14 +39,6 @@ def get_image_paths(directories: Union[str, List[str]]) -> List[str]:
|
|
| 45 |
|
| 46 |
return list(set(all_images)) # Remove duplicates
|
| 47 |
|
| 48 |
-
def clean_directory(directory: str, create_if_not_exists: bool = True) -> None:
|
| 49 |
-
"""Clean directory contents or create if it doesn't exist."""
|
| 50 |
-
if os.path.exists(directory):
|
| 51 |
-
shutil.rmtree(directory)
|
| 52 |
-
|
| 53 |
-
if create_if_not_exists:
|
| 54 |
-
os.makedirs(directory, exist_ok=True)
|
| 55 |
-
|
| 56 |
def backup_file(source_path: str, backup_path: str) -> str:
|
| 57 |
"""Backup a file to specified location."""
|
| 58 |
backup_path = get_abs_path(backup_path)
|
|
@@ -66,7 +52,8 @@ import os
|
|
| 66 |
import cv2
|
| 67 |
from ultralytics import YOLO
|
| 68 |
from typing import List, Optional, Dict, Any
|
| 69 |
-
from .utils import
|
|
|
|
| 70 |
|
| 71 |
class YOLOManager:
|
| 72 |
"""Manages YOLO model training and inference operations."""
|
|
@@ -82,7 +69,7 @@ class YOLOManager:
|
|
| 82 |
self.model = YOLO(weights_path)
|
| 83 |
else:
|
| 84 |
print("✨ Loading pretrained model 'yolo11s.pt'")
|
| 85 |
-
self.model = YOLO("
|
| 86 |
return self.model
|
| 87 |
|
| 88 |
def train(self,
|
|
@@ -102,7 +89,7 @@ class YOLOManager:
|
|
| 102 |
**kwargs: Additional training parameters
|
| 103 |
"""
|
| 104 |
run_name = run_name or self.model_name
|
| 105 |
-
checkpoint_path = f"runs/detect/{run_name}/weights/last.pt"
|
| 106 |
|
| 107 |
# Check for existing checkpoint
|
| 108 |
if resume and os.path.isfile(checkpoint_path):
|
|
@@ -122,7 +109,7 @@ class YOLOManager:
|
|
| 122 |
'name': run_name,
|
| 123 |
'device': device,
|
| 124 |
'cache': True,
|
| 125 |
-
'project': 'runs/detect',
|
| 126 |
'exist_ok': True,
|
| 127 |
'resume': resume_flag
|
| 128 |
}
|
|
|
|
| 7 |
|
| 8 |
load_dotenv()
|
| 9 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 10 |
def get_abs_path(relative_path: str) -> str:
|
| 11 |
"""Convert relative path to absolute path."""
|
| 12 |
return os.path.abspath(relative_path)
|
|
|
|
| 39 |
|
| 40 |
return list(set(all_images)) # Remove duplicates
|
| 41 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 42 |
def backup_file(source_path: str, backup_path: str) -> str:
|
| 43 |
"""Backup a file to specified location."""
|
| 44 |
backup_path = get_abs_path(backup_path)
|
|
|
|
| 52 |
import cv2
|
| 53 |
from ultralytics import YOLO
|
| 54 |
from typing import List, Optional, Dict, Any
|
| 55 |
+
from .utils import get_abs_path, clean_directory
|
| 56 |
+
from .config import Config
|
| 57 |
|
| 58 |
class YOLOManager:
|
| 59 |
"""Manages YOLO model training and inference operations."""
|
|
|
|
| 69 |
self.model = YOLO(weights_path)
|
| 70 |
else:
|
| 71 |
print("✨ Loading pretrained model 'yolo11s.pt'")
|
| 72 |
+
self.model = YOLO(f"{Config.current_path}/yolo12s.pt")
|
| 73 |
return self.model
|
| 74 |
|
| 75 |
def train(self,
|
|
|
|
| 89 |
**kwargs: Additional training parameters
|
| 90 |
"""
|
| 91 |
run_name = run_name or self.model_name
|
| 92 |
+
checkpoint_path = f"{Config.current_path}/runs/detect/{run_name}/weights/last.pt"
|
| 93 |
|
| 94 |
# Check for existing checkpoint
|
| 95 |
if resume and os.path.isfile(checkpoint_path):
|
|
|
|
| 109 |
'name': run_name,
|
| 110 |
'device': device,
|
| 111 |
'cache': True,
|
| 112 |
+
'project': f'{Config.current_path}/runs/detect',
|
| 113 |
'exist_ok': True,
|
| 114 |
'resume': resume_flag
|
| 115 |
}
|