inf all

comic_panel_extractor/create_dataset.py

@@ -3,7 +3,8 @@ import shutil
 import random
 from pathlib import Path
 from dotenv import load_dotenv
-from tqdm import tqdm  # <-- import tqdm
+from tqdm import tqdm
+from .config import Config
 
 load_dotenv()
 SOURCE_PATHS = os.getenv('SOURCE_PATH')
@@ -14,17 +15,12 @@ if not SOURCE_PATHS:
 # Split by comma and strip whitespace
 source_paths = [Path(p.strip()) for p in SOURCE_PATHS.split(',')]
 
-images_dir = Path('images')
-dataset_dir = Path('dataset')
+images_dir = Path(f'{Config.current_path}/images')
+dataset_dir = Path(f'{Config.current_path}/dataset')
 
 image_exts = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}
 label_exts = {'.txt'}
 
-# Clear images_dir
-if images_dir.exists():
-    shutil.rmtree(images_dir)
-images_dir.mkdir(parents=True)
-
 # Copy images from all source paths with tqdm progress
 for source_path in source_paths:
     if not source_path.exists():
@@ -76,7 +72,7 @@ splits = {
     'test': all_images[val_end:]
 }
 
-label_src_dir = Path('image_labels')
+label_src_dir = Path(f'{Config.current_path}/image_labels')
 
 # Move/copy images and labels to their split folders with tqdm
 for split, files in splits.items():

comic_panel_extractor/inference.py

@@ -1,6 +1,6 @@
 # inference.py
-from yolo_manager import YOLOManager
-from utils import get_abs_path, get_image_paths
+from .yolo_manager import YOLOManager
+from .utils import get_abs_path, get_image_paths
 import os
 from .config import Config
 
@@ -50,5 +50,24 @@ def main():
     
     run_inference(weights_path, images_dirs, './temp_dir')
 
+def annotate_all_image():
+    with YOLOManager() as yolo_manager:
+        weights_path = Config.yolo_trained_model_path
+        yolo_manager.load_model(weights_path)
+        IMAGE_ROOT = os.path.join(Config.current_path, "dataset/images")
+        IMAGE_LABEL_ROOT = os.path.join(Config.current_path, "image_labels")
+        for root, _, files in os.walk(IMAGE_ROOT):
+            for file in sorted(files):
+                if file.lower().endswith((".jpg", ".jpeg", ".png")):
+                    name, ext = os.path.splitext(file)
+                    save_txt_path = os.path.join(IMAGE_LABEL_ROOT, f'{name}.txt')  # YOLO label txt
+                    if not os.path.exists(save_txt_path):
+                        image_path = os.path.join(root, file)
+                        yolo_manager.annotate_images(
+                            image_paths=[image_path],
+                            output_dir=IMAGE_LABEL_ROOT,
+                            save_image=False
+                        )
+
 if __name__ == "__main__":
-    main()
+    annotate_all_image()