Update saicinpainting/training/trainers/default.py

saicinpainting/training/trainers/default.py

@@ -1,5 +1,5 @@
 import logging
-
+import cv2
 import torch
 import torch.nn.functional as F
 from omegaconf import OmegaConf
@@ -13,6 +13,30 @@ from saicinpainting.utils import add_prefix_to_keys, get_ramp
 
 LOGGER = logging.getLogger(__name__)
 
+def resize_to_square(image, target_size):
+    h, w = image.shape[:2]
+    if h == w:
+        return cv2.resize(image, (target_size, target_size))
+    
+    dif = h if h > w else w
+    interpolation = cv2.INTER_AREA if dif > target_size else cv2.INTER_CUBIC
+    
+    x_pos = (dif - w) // 2
+    y_pos = (dif - h) // 2
+    
+    if len(image.shape) == 2:
+        mask = np.zeros((dif, dif), dtype=image.dtype)
+        mask[y_pos:y_pos+h, x_pos:x_pos+w] = image
+    else:
+        mask = np.zeros((dif, dif, image.shape[2]), dtype=image.dtype)
+        mask[y_pos:y_pos+h, x_pos:x_pos+w, :] = image
+    
+    return cv2.resize(mask, (target_size, target_size), interpolation=interpolation)
+
+# Sử dụng
+target_size = 256
+resized_frame = resize_to_square(frame, target_size)
+
 
 def make_constant_area_crop_batch(batch, **kwargs):
     crop_y, crop_x, crop_height, crop_width = make_constant_area_crop_params(img_height=batch['image'].shape[2],
@@ -24,25 +48,9 @@ def make_constant_area_crop_batch(batch, **kwargs):
 
 
 class DefaultInpaintingTrainingModule(BaseInpaintingTrainingModule):
-    def __init__(self, *args, concat_mask=True, rescale_scheduler_kwargs=None, image_to_discriminator='predicted_image',
-                 add_noise_kwargs=None, noise_fill_hole=False, const_area_crop_kwargs=None,
-                 distance_weighter_kwargs=None, distance_weighted_mask_for_discr=False,
-                 fake_fakes_proba=0, fake_fakes_generator_kwargs=None,
-                 **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.concat_mask = concat_mask
-        self.rescale_size_getter = get_ramp(**rescale_scheduler_kwargs) if rescale_scheduler_kwargs is not None else None
-        self.image_to_discriminator = image_to_discriminator
-        self.add_noise_kwargs = add_noise_kwargs
-        self.noise_fill_hole = noise_fill_hole
-        self.const_area_crop_kwargs = const_area_crop_kwargs
-        self.refine_mask_for_losses = make_mask_distance_weighter(**distance_weighter_kwargs) \
-            if distance_weighter_kwargs is not None else None
-        self.distance_weighted_mask_for_discr = distance_weighted_mask_for_discr
-
-        self.fake_fakes_proba = fake_fakes_proba
-        if self.fake_fakes_proba > 1e-3:
-            self.fake_fakes_gen = FakeFakesGenerator(**(fake_fakes_generator_kwargs or {}))
+        self.target_size = 256  # Hoặc kích thước mong muốn khác
 
     def forward(self, batch):
         if self.training and self.rescale_size_getter is not None:
@@ -50,6 +58,29 @@ class DefaultInpaintingTrainingModule(BaseInpaintingTrainingModule):
             batch['image'] = F.interpolate(batch['image'], size=cur_size, mode='bilinear', align_corners=False)
             batch['mask'] = F.interpolate(batch['mask'], size=cur_size, mode='nearest')
 
+        # Thêm đoạn code resize ở đây
+        resized_images = []
+        resized_masks = []
+        for img, mask in zip(batch['image'], batch['mask']):
+            # Chuyển từ tensor sang numpy array
+            img_np = img.permute(1, 2, 0).cpu().numpy()
+            mask_np = mask.squeeze().cpu().numpy()
+            
+            # Resize
+            img_resized = resize_to_square(img_np, self.target_size)
+            mask_resized = resize_to_square(mask_np, self.target_size)
+            
+            # Chuyển lại thành tensor
+            img_resized = torch.from_numpy(img_resized).permute(2, 0, 1).float().to(img.device)
+            mask_resized = torch.from_numpy(mask_resized).unsqueeze(0).float().to(mask.device)
+            
+            resized_images.append(img_resized)
+            resized_masks.append(mask_resized)
+        
+        batch['image'] = torch.stack(resized_images)
+        batch['mask'] = torch.stack(resized_masks)
+
+        # Tiếp tục với phần còn lại của phương thức forward
         if self.training and self.const_area_crop_kwargs is not None:
             batch = make_constant_area_crop_batch(batch, **self.const_area_crop_kwargs)