Spaces:
Running
Running
File size: 4,130 Bytes
ec6ad2f c13ce0c 951222c a1f4a1e 4ecf9fd a1f4a1e 951222c 4ecf9fd 951222c ec6ad2f 4ecf9fd fb59b22 4ecf9fd 0bb457d 4ecf9fd fb59b22 4ecf9fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
from dataclasses import dataclass
import os
import toml
from dotenv import load_dotenv
load_dotenv()
CURRENT_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__)))
CONFIG_FILE = os.path.join(CURRENT_PATH, "config.toml")
@dataclass
class Config:
"""Configuration settings for the comic-to-video pipeline."""
# Paths
current_path: str = CURRENT_PATH
config_path: str = CONFIG_FILE
# Core settings
EPOCH: int = 200
DEFAULT_IMAGE_SIZE: int = 640
BATCH: int = 10
RESUME_TRAIN: bool = True
RECREATE_DATASET: bool = True
# YOLO models
YOLO_BASE_MODEL_NAME: str = "yolo11s-seg"
YOLO_MODEL_NAME: str = "" # will be derived if empty
IMAGE_SOURCE_PATH: str = ""
YOLO_MODEL_REMOTE_URL: str = ""
# Derived paths
yolo_base_model_path: str = ""
yolo_trained_model_path: str = ""
# Pipeline parameters
org_input_path: str = ""
input_path: str = ""
black_overlay_input_path: str = ""
output_folder: str = "temp_dir"
distance_threshold: int = 70
vertical_threshold: int = 30
text_cood_file_name: str = "detect_and_group_text.json"
min_text_length: int = 2
min_area_ratio: float = 0.05
min_width_ratio: float = 0.15
min_height_ratio: float = 0.15
# BorderPanelExtractor
panel_filename_pattern: str = r"panel_\d+_\((\d+), (\d+), (\d+), (\d+)\)\.jpg"
# Constants
SUPPORTED_EXTENSIONS: tuple = ('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')
def __post_init__(self):
# Ensure absolute IMAGE_SOURCE_PATH
if self.IMAGE_SOURCE_PATH:
if not os.path.isabs(self.IMAGE_SOURCE_PATH):
self.IMAGE_SOURCE_PATH = os.path.join(self.current_path, self.IMAGE_SOURCE_PATH)
# Derive YOLO_MODEL_NAME if empty
if not self.YOLO_MODEL_NAME:
self.YOLO_MODEL_NAME = f"final_model_{self.YOLO_BASE_MODEL_NAME}"
# Derived paths
self.yolo_base_model_path = os.path.join(self.current_path, f"{self.YOLO_BASE_MODEL_NAME}.pt")
self.yolo_trained_model_path = os.path.join(self.current_path, f"{self.YOLO_MODEL_NAME}.pt")
def load_config(file_path=CONFIG_FILE) -> Config:
"""Load the latest config from TOML file and return a Config instance."""
if not os.path.exists(file_path):
raise FileNotFoundError(f"Config file not found: {file_path}")
data = toml.load(file_path)
# Convert boolean strings to actual bool
def to_bool(val):
if isinstance(val, bool):
return val
return str(val).lower() in ("1", "true", "yes")
return Config(
EPOCH=int(data.get("EPOCH", 200)),
DEFAULT_IMAGE_SIZE=int(data.get("DEFAULT_IMAGE_SIZE", 640)),
BATCH=int(data.get("BATCH", 10)),
RESUME_TRAIN=to_bool(data.get("RESUME_TRAIN", True)),
RECREATE_DATASET=to_bool(data.get("RECREATE_DATASET", True)),
YOLO_BASE_MODEL_NAME=data.get("YOLO_BASE_MODEL_NAME", "yolo11s-seg"),
YOLO_MODEL_NAME=data.get("YOLO_MODEL_NAME", ""), # derived in __post_init__
IMAGE_SOURCE_PATH=data.get("IMAGE_SOURCE_PATH", ""),
YOLO_MODEL_REMOTE_URL=data.get("YOLO_MODEL_REMOTE_URL", "")
)
def update_toml_key(key: str, value, file_path=CONFIG_FILE) -> Config:
"""Update a key in the TOML file and reload config."""
if not os.path.exists(file_path):
raise FileNotFoundError(f"Config file not found: {file_path}")
data = toml.load(file_path)
data[key] = value
with open(file_path, "w") as f:
toml.dump(data, f)
# Reload and return new Config
return load_config(file_path)
def get_text_cood_file_path(config: Config) -> str:
"""Return full path to text coordinate file."""
return os.path.join(config.output_folder, config.text_cood_file_name)
# Example usage:
if __name__ == "__main__":
# Load config
config = load_config()
print("EPOCH:", config.EPOCH)
# Update TOML key and reload
config = update_toml_key("EPOCH", 500)
print("Updated EPOCH:", config.EPOCH)
# Get text coord file path
print("Text coord path:", get_text_cood_file_path(config))
|