Merge branch 'Experiments'

pawansharmaaaa · Mar 12, 2024 · c00a896 · c00a896
2 parents c0ff687 + 2d83ae5
commit c00a896
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 27 deletions.
diff --git a/helpers/model_loaders.py b/helpers/model_loaders.py
@@ -23,14 +23,19 @@ class ModelLoader:
     def __init__(self, restorer, weight):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.weight = weight
-        # self.wav2lip_model = self.load_wav2lip_model()
+
         if restorer == 'GFPGAN':
             self.restorer = self.load_gfpgan_model()
         elif restorer == 'RestoreFormer':
             self.restorer = self.load_restoreformer_model()
         elif restorer == 'CodeFormer':
             self.restorer = self.load_codeformer_model()
 
+        # Create Plate
+        self.plate = np.full((512, 512, 3), (128, 128, 128), dtype=np.uint8)
+
+        self.bgupsampler = None
+
     def _load(self, checkpoint_path):
         if self.device == 'cuda':
             checkpoint = torch.load(checkpoint_path)
@@ -152,15 +157,20 @@ def load_codeformer_model(self):
 
     @torch.no_grad()
     def restore_background(self, background, model_name, tile, outscale=1.0, half=False):
-        bgupsampler = self.load_realesrgan_model(model_name, tile, half=False)
-        if bgupsampler is not None:
-            background = bgupsampler.enhance(background, outscale=outscale)
+        if self.bgupsampler is None:
+            self.bgupsampler = self.load_realesrgan_model(model_name, tile, half=False)
+
+        background = self.bgupsampler.enhance(background, outscale=outscale)[0]
         return background
 
     @torch.no_grad()
     def restore_wGFPGAN(self, dubbed_face):
-        dubbed_face = cv2.resize(dubbed_face.astype(np.uint8) / 255., (512, 512), interpolation=cv2.INTER_LANCZOS4)
-        dubbed_face_t = img2tensor(dubbed_face, bgr2rgb=True, float32=True)
+        dubbed_face = cv2.resize(dubbed_face.astype(np.uint8), (300, 300), interpolation=cv2.INTER_LANCZOS4)
+        bkg = self.plate.copy()
+        bkg[106:406, 106:406] = dubbed_face
+        dubbed_face = bkg
+
+        dubbed_face_t = img2tensor(dubbed_face / 255., bgr2rgb=True, float32=True)
         normalize(dubbed_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
         dubbed_face_t = dubbed_face_t.unsqueeze(0).to(self.device)
 
@@ -172,12 +182,18 @@ def restore_wGFPGAN(self, dubbed_face):
             restored_face = tensor2img(dubbed_face_t.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
 
         restored_face = restored_face.astype(np.uint8)
+        restored_face = self.restore_background(restored_face, 'RealESRGAN_x2plus', tile=512, half=False)
+        restored_face = restored_face[106:406, 106:406]
         return restored_face
 
     @torch.no_grad()
     def restore_wRF(self, dubbed_face):
-        dubbed_face = cv2.resize(dubbed_face.astype(np.uint8) / 255., (512, 512), interpolation=cv2.INTER_LANCZOS4)
-        dubbed_face_t = img2tensor(dubbed_face, bgr2rgb=True, float32=True)
+        dubbed_face = cv2.resize(dubbed_face.astype(np.uint8), (300, 300), interpolation=cv2.INTER_LANCZOS4)
+        bkg = self.plate.copy()
+        bkg[106:406, 106:406] = dubbed_face
+        dubbed_face = bkg
+
+        dubbed_face_t = img2tensor(dubbed_face / 255., bgr2rgb=True, float32=True)
         normalize(dubbed_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
         dubbed_face_t = dubbed_face_t.unsqueeze(0).to(self.device)
 
@@ -189,12 +205,18 @@ def restore_wRF(self, dubbed_face):
             restored_face = tensor2img(dubbed_face_t.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
 
         restored_face = restored_face.astype(np.uint8)
+        restored_face = self.restore_background(restored_face, 'RealESRGAN_x2plus', tile=512, half=False)
+        restored_face = restored_face[106:406, 106:406]
         return restored_face
 
     @torch.no_grad()
     def restore_wCodeFormer(self, dubbed_face):
-        dubbed_face = cv2.resize(dubbed_face.astype(np.uint8) / 255., (512, 512), interpolation=cv2.INTER_LANCZOS4)
-        dubbed_face_t = img2tensor(dubbed_face, bgr2rgb=True, float32=True)
+        dubbed_face = cv2.resize(dubbed_face.astype(np.uint8), (300, 300), interpolation=cv2.INTER_LANCZOS4)
+        bkg = self.plate.copy()
+        bkg[106:406, 106:406] = dubbed_face
+        dubbed_face = bkg
+
+        dubbed_face_t = img2tensor(dubbed_face / 255., bgr2rgb=True, float32=True)
         normalize(dubbed_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
         dubbed_face_t = dubbed_face_t.unsqueeze(0).to(self.device)
 
@@ -209,4 +231,6 @@ def restore_wCodeFormer(self, dubbed_face):
             restored_face = tensor2img(dubbed_face_t, rgb2bgr=True, min_max=(-1, 1))
 
         restored_face = restored_face.astype(np.uint8)
+        restored_face = self.restore_background(restored_face, 'RealESRGAN_x2plus', tile=512, half=False)
+        restored_face = restored_face[106:406, 106:406]
         return restored_face
diff --git a/helpers/preprocess_mp.py b/helpers/preprocess_mp.py
@@ -630,7 +630,6 @@ def paste_back_black_bg(self, processed_face, aligned_bbox, full_frame, ml):
         bbox = np.asarray(([aligned_bbox[0], aligned_bbox[1]], [aligned_bbox[0]+aligned_bbox[2], aligned_bbox[1]+aligned_bbox[3]]))
         processed_ready = np.zeros_like(full_frame)
         try:
-            processed_face = ml.restore_background(processed_face, 'RealESRGAN_x2plus', 400)[0]
             processed_ready[bbox[0,1]:bbox[1,1], bbox[0,0]:bbox[1,0]] = processed_face
         except IndexError as e:
             gr.Warning(f"Failed to paste face back onto full frame: {e}")

diff --git a/opt.ipynb b/opt.ipynb
@@ -394,22 +394,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Feed to model:\n",
-    "images_list = []\n",
-    "mels_list = []\n",
-    "extracted_faces_list = []\n",
-    "face_masks_list = []\n",
-    "inv_masks_list = []\n",
-    "cropped_faces_list = []\n",
-    "frame_batch_list = []\n",
-    "mel_batch_list = []\n",
-    "dubbed_faces_list = []\n",
-    "restored_faces_list = []\n",
-    "resized_restored_faces_list = []\n",
-    "pasted_ready_faces_list = []\n",
-    "ready_to_paste_list = []\n",
-    "restored_images_list = []\n",
-    "upscaled_bg_list = []"
+    "import os\n",
+    "\n",
+    "# Create a directory inside the data directory\n",
+    "data_directory = \"data\"\n",
+    "list_directory = os.path.join(data_directory, \"lists\")\n",
+    "os.makedirs(list_directory, exist_ok=True)\n",
+    "\n",
+    "# Create directories with the name of each list\n",
+    "lists = [\"images_list\", \"mels_list\", \"extracted_faces_list\", \"face_masks_list\", \"inv_masks_list\", \"cropped_faces_list\", \"frame_batch_list\", \"mel_batch_list\", \"dubbed_faces_list\", \"restored_faces_list\", \"resized_restored_faces_list\", \"pasted_ready_faces_list\", \"ready_to_paste_list\", \"restored_images_list\", \"upscaled_bg_list\"]\n",
+    "\n",
+    "for list_name in lists:\n",
+    "    list_path = os.path.join(list_directory, list_name)\n",
+    "    os.makedirs(list_path, exist_ok=True)\n"
    ]
   },
   {