refine

comic_panel_extractor/annorator_server.py

@@ -75,7 +75,7 @@ def load_yolo_annotations(image_path: str, label_path: str, detect: bool = False
         if detect and not os.path.exists(label_path):
             from .yolo_manager import YOLOManager
             with YOLOManager() as yolo_manager:
-                weights_path = f'{Config.current_path}/{Config.YOLO_MODEL_NAME}.pt'
+                weights_path = Config.yolo_trained_model_path
                 yolo_manager.load_model(weights_path)
                 _, label_path = yolo_manager.annotate_images(
                     image_paths=[image_path],

comic_panel_extractor/config.py

@@ -10,8 +10,10 @@ class Config:
 	org_input_path: str = ""
 	input_path: str = ""
 	current_path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
-	YOLO_MODEL_NAME = os.getenv('YOLO_MODEL_NAME', 'default_model')
-	yolo_model_path: str = f'{current_path}/{YOLO_MODEL_NAME}.pt'
+	YOLO_BASE_MODEL_NAME = os.getenv('YOLO_BASE_MODEL_NAME', 'yolo11s-seg')
+	yolo_base_model_path: str = f'{current_path}/{YOLO_BASE_MODEL_NAME}.pt'
+	YOLO_MODEL_NAME = f'{os.getenv('YOLO_MODEL_NAME', 'comic_panel')}_{YOLO_BASE_MODEL_NAME}'
+	yolo_trained_model_path: str = f'{current_path}/{YOLO_MODEL_NAME}.pt'
 	black_overlay_input_path: str = ""
 	output_folder: str = "temp_dir"
 	distance_threshold: int = 70

comic_panel_extractor/inference.py

@@ -41,7 +41,7 @@ def run_inference(weights_path: str, images_dirs, output_dir: str = 'temp_dir')
 
 def main():
     """Main inference function."""
-    weights_path = f'{Config.YOLO_MODEL_NAME}.pt'
+    weights_path = Config.yolo_trained_model_path
     images_dirs = [
         './dataset/images/train',
         './dataset/images/val', 

comic_panel_extractor/llm_panel_extractor.py

@@ -16,17 +16,16 @@ class LLMPanelExtractor:
 		self.config = config or Config()
 
 		# Check if YOLO model exists; if not, download it to the specified path
-		yolo_model_path = self.config.yolo_model_path
-		if not os.path.exists(yolo_model_path):
+		if not os.path.exists(self.config.yolo_base_model_path):
 			url = "https://huggingface.co/mosesb/best-comic-panel-detection/resolve/main/best.pt"
-			print(f"Downloading YOLO model to {yolo_model_path}...")
+			print(f"Downloading YOLO model to {self.config.yolo_base_model_path}...")
 			response = requests.get(url)
 			response.raise_for_status()  # Raise an error if the download fails
-			with open(yolo_model_path, "wb") as f:
+			with open(self.config.yolo_base_model_path, "wb") as f:
 				f.write(response.content)
 			print("YOLO model downloaded successfully.")
 
-		self.yolo_model = YOLO(yolo_model_path)
+		self.yolo_model = YOLO(self.config.yolo_base_model_path)
 		os.makedirs(self.config.output_folder, exist_ok=True)
 
 	def extract_bounding_boxes(self, detection_result_boxes):

comic_panel_extractor/train.py

@@ -185,7 +185,7 @@ def main():
         
         # Backup best weights
         weights_path = yolo_manager.get_best_weights_path()
-        backup_path = f'{Config.current_path}/{Config.YOLO_MODEL_NAME}.pt'
+        backup_path = Config.yolo_trained_model_path
         backup_file(weights_path, backup_path)
         
         print("🎉 Training completed successfully!")

comic_panel_extractor/yolo_manager.py

@@ -54,6 +54,8 @@ from ultralytics import YOLO
 from typing import List, Optional, Dict, Any
 from .utils import get_abs_path, clean_directory
 from .config import Config
+from dotenv import load_dotenv
+load_dotenv()
 
 class YOLOManager:
     """Manages YOLO model training and inference operations."""
@@ -68,8 +70,8 @@ class YOLOManager:
             print(f"📦 Loading model from: {weights_path}")
             self.model = YOLO(weights_path)
         else:
-            print("✨ Loading pretrained model 'yolov12s-seg.pt'")
-            self.model = YOLO(f"{Config.current_path}/yolov12s-seg.pt")
+            print(f"✨ Loading pretrained model '{Config.yolo_base_model_path}'")
+            self.model = YOLO(f"{Config.yolo_base_model_path}")
         return self.model
     
     def train(self, 
@@ -104,7 +106,7 @@ class YOLOManager:
         train_params = {
             'data': data_yaml_path,
             'imgsz': Config.DEFAULT_IMAGE_SIZE,
-            'epochs': 200,
+            'epochs': Config.YOLO_BASE_MODEL_NAME,
             'batch': 10,
             'name': run_name,
             'device': device,