From 0fdcc07422742b1296acf83caa96323ca1e2aa60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= <rafal@tluchowski.pl>
Date: Sun, 23 Feb 2025 20:05:50 +0100
Subject: [PATCH] path loader to imageprocessor

---
 core/image_processor.py       | 13 +++++++++++++
 nodes/face_fit_and_restore.py | 19 +++----------------
 nodes/face_tracker.py         |  8 +++++---
 3 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/core/image_processor.py b/core/image_processor.py
index ea8d638..8bbce34 100644
--- a/core/image_processor.py
+++ b/core/image_processor.py
@@ -460,3 +460,16 @@ class ImageProcessor:
 
         return torch.from_numpy(mask).unsqueeze(0)
 
+    # Loaders
+    @staticmethod
+    def load_image_from_path(image_path):
+        """Load and convert image from file path to tensor format."""
+        try:
+            image = Image.open(image_path)
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            return ImageProcessor.pil_to_tensor(image)
+        except Exception as e:
+            print(f"Error loading image from {image_path}: {str(e)}")
+            return None
+
diff --git a/nodes/face_fit_and_restore.py b/nodes/face_fit_and_restore.py
index f5c681c..3c3869d 100644
--- a/nodes/face_fit_and_restore.py
+++ b/nodes/face_fit_and_restore.py
@@ -92,12 +92,10 @@ class FaceFitAndRestore:
                     fp_pipe=fp_pipe
                 )
             else:  # image_sequence
-                # W trybie Restore używamy current_frame z fp_pipe
                 current_frame = (fp_pipe.get("current_frame", 0) if mode == "Restore"
                                  else image_sequence.get("current_frame", 0))
                 fp_pipe["current_frame"] = current_frame
 
-                # W trybie Restore używamy liczby klatek z fp_pipe
                 if mode == "Restore":
                     total_frames = len(fp_pipe["frames"])
                     frames_to_process = {
@@ -114,15 +112,15 @@ class FaceFitAndRestore:
                 for frame_idx in range(total_frames):
                     print(f"Processing frame {frame_idx + 1}/{total_frames}")
 
-                    # Przygotowanie obrazu wejściowego
+
                     if mode == "Fit":
                         frame_path = frames_to_process[frame_idx]
-                        frame_image = self._load_image_from_path(frame_path)
+                        frame_image = self.image_processor.load_image_from_path(frame_path)
                         if frame_image is None:
                             print(f"Warning: Could not load frame {frame_idx}")
                             continue
                     else:  # Restore
-                        frame_image = image  # Ten sam obraz dla każdej klatki w trybie Restore
+                        frame_image = image
                         frame_path = None
 
                     # Process frame
@@ -339,14 +337,3 @@ class FaceFitAndRestore:
         else:
             shape = image.shape[:2]
         return torch.zeros((1, *shape), dtype=torch.float32)
-
-    def _load_image_from_path(self, image_path):
-        """Load and convert image from file path to tensor format."""
-        try:
-            image = Image.open(image_path)
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
-            return self.image_processor.pil_to_tensor(image)
-        except Exception as e:
-            print(f"Error loading image from {image_path}: {str(e)}")
-            return None
\ No newline at end of file
diff --git a/nodes/face_tracker.py b/nodes/face_tracker.py
index 6dff0a2..8f68e92 100644
--- a/nodes/face_tracker.py
+++ b/nodes/face_tracker.py
@@ -14,12 +14,14 @@ class FaceTracker:
         return {
             "required": {
                 "image": ("IMAGE",),
+                "fp_pipe": ("DICT",),
             }
         }
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "DICT")
+    RETURN_NAMES = ("image", "fp_pipe")
     FUNCTION = "track_face"
     CATEGORY = "Face Processor"
 
-    def track_face(self, image):
-        return (image,)
\ No newline at end of file
+    def track_face(self, image, fp_pipe):
+        return image, fp_pipe
\ No newline at end of file