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,30 +13,6 @@ 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],
@@ -48,9 +24,25 @@ def make_constant_area_crop_batch(batch, **kwargs):
 
 
 class DefaultInpaintingTrainingModule(BaseInpaintingTrainingModule):
-    def __init__(self, *args, **kwargs):
+    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):
         super().__init__(*args, **kwargs)
-        self.target_size = 256  # Hoặc kích thước mong muốn khác
+        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 {}))
 
     def forward(self, batch):
         if self.training and self.rescale_size_getter is not None:
@@ -58,29 +50,6 @@ 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)
 
@@ -203,4 +172,4 @@ class DefaultInpaintingTrainingModule(BaseInpaintingTrainingModule):
             metrics['discr_adv_fake_fakes'] = fake_fakes_adv_discr_loss
             metrics.update(add_prefix_to_keys(fake_fakes_adv_metrics, 'adv_'))
 
-        return total_loss, metrics
+        return total_loss, metrics