image to image function (#3)

* image to image function * add images * Update README.md * Add files via upload * Update README.md * Update README.md * Delete orig.webm * Delete mark.webm * Delete khabenskii.webm * Delete elon.webm * Add files via upload * Add files via upload * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Add files via upload * Update README.md * Add files via upload * preprocess vgg dataset * Update README.md * Update preprocess_vgg.py * Update README.md * Add files via upload * Add files via upload * Delete khabenskii.jpg * Add files via upload * Update video_processing.py * Update video_processing.py * Add files via upload * Add files via upload Co-authored-by: danyache <[email protected]>
ai-forever · Dec 27, 2021 · 27f81c7 · 27f81c7
1 parent 8c974cd
commit 27f81c7
Show file tree

Hide file tree

Showing 16 changed files with 95 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -1 +1,52 @@
-# sber-swap
+# sber-swap
+
+## Results 
+![](/examples/images/example1.png)
+
+![](/examples/images/example2.png)
+
+## Video Swap
+<div>
+<img src="/examples/videos/orig.webp" width="360"/>
+<img src="/examples/videos/elon.webp" width="360"/>
+<img src="/examples/videos/khabenskii.webp" width="360"/>
+<img src="/examples/videos/mark.webp" width="360"/>
+<div/>
+
+## Installation
+
+1. Clone this repository
+  ```bash
+  git clone https://github.com/Danyache/sber-swap.git
+  cd sber-swap
+  git submodule init
+  git submodule update
+  ```
+2. Install dependent packages
+  ```bash
+  pip install -r requirements.txt
+  ```
+3. Download weights
+  ```bash
+  sh download_models.sh
+  ```
+## Usage
+  1. Colab Demo
+  2. Face Swap On Video
+  > python inference.py 
+  3. Face Swap On Image
+  > python inference.py --target_path examples/images/beckham.jpg --image_to_image True
+  
+## Training
+
+We also provide the training code for face swap model as follows:
+  1. Download [VGGFace2 Dataset](https://www.robots.ox.ac.uk/~vgg/data/vgg_face/).
+  2. Crop and align faces with out detection model.
+  > python preprocess_vgg.py --path_to_dataset ./VggFace2/VGG-Face2/data/preprocess_train --save_path ./VggFace2-crop
+  3. Start training. 
+  > python train.py
+  
+### Tips:
+  1. For first epochs we suggest not to use eye detection loss
+  2. In case of finetuning model you can variate losses coefficients to make result look more like source identity, or vice versa, save features and attributes of target face
+
diff --git a/examples/images/beckham.jpg b/examples/images/beckham.jpg
diff --git a/examples/images/example1.png b/examples/images/example1.png
diff --git a/examples/images/example2.png b/examples/images/example2.png
diff --git a/examples/images/murakami.jpg b/examples/images/murakami.jpg
diff --git a/examples/results/result.mp4 b/examples/results/result.mp4
diff --git a/examples/results/result.png b/examples/results/result.png
diff --git a/examples/videos/dance.mp4 b/examples/videos/dance.mp4
diff --git a/examples/videos/elon.webp b/examples/videos/elon.webp
diff --git a/examples/videos/khabenskii.webp b/examples/videos/khabenskii.webp
diff --git a/examples/videos/mark.webp b/examples/videos/mark.webp
diff --git a/examples/videos/nggyup.mp4 b/examples/videos/nggyup.mp4
diff --git a/examples/videos/orig.webp b/examples/videos/orig.webp
diff --git a/preprocess_vgg.py b/preprocess_vgg.py
@@ -0,0 +1,41 @@
+import os
+import sys
+import cv2
+import argparse
+import torch
+import numpy as np
+from insightface_func.face_detect_crop_single import Face_detect_crop
+from pathlib import Path
+from tqdm import tqdm
+
+def main(args):
+    app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+    app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
+    crop_size = 224
+
+    dirs = os.listdir(args.path_to_dataset)
+    for i in tqdm(range(len(dirs))):
+        d = os.path.join(args.path_to_dataset, dirs[i])
+        dir_to_save = os.path.join(args.save_path, dirs[i])
+        Path(dir_to_save).mkdir(parents=True, exist_ok=True)
+
+        image_names = os.listdir(d)
+        for image_name in image_names:
+            try:
+                image_path = os.path.join(d, image_name)
+                image = cv2.imread(image_path)
+                cropped_image, _ = app.get(image, crop_size)
+                cv2.imwrite(os.path.join(dir_to_save, image_name), cropped_image[0])
+            except:
+                pass
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--path_to_dataset', default='./VggFace2/VGG-Face2/data/preprocess_train', type=str)
+    parser.add_argument('--save_path', default='./VggFace2-crop', type=str)
+
+    args = parser.parse_args()
+
+    main(args)
diff --git a/utils/inference/image_processing.py b/utils/inference/image_processing.py
@@ -83,4 +83,3 @@ def show_images(target, swap):
     ax2.axis('off')
     ax1.imshow(target[:,:,::-1])
     ax2.imshow(swap[:,:,::-1])
-
diff --git a/utils/inference/video_processing.py b/utils/inference/video_processing.py
@@ -16,14 +16,6 @@
 import kornia
 
 
-# def video_is_less_than(path_to_video: str, seconds=10.) -> bool:
-
-#     cap = cv2.VideoCapture(path_to_video)
-#     fps, frames = cap.get(cv2.CAP_PROP_FPS), cap.get(cv2.CAP_PROP_FRAME_COUNT)
-
-#     return frames < fps*seconds
-
-
 def add_audio_from_another_video(video_with_sound: str, 
                                  video_without_sound: str, 
                                  audio_name: str, 
@@ -220,7 +212,7 @@ def __getitem__(self, idx):
 
     def __len__(self):
         return len(self.frames_list)
-    
+
 
 def face_enhancement(final_frames: List[np.ndarray], model) -> List[np.ndarray]:
     enhanced_frames_all = []
@@ -244,4 +236,4 @@ def face_enhancement(final_frames: List[np.ndarray], model) -> List[np.ndarray]:
                 ff_i+=1
         enhanced_frames_all.append(enhanced_frames)
 
-    return enhanced_frames_all
+    return enhanced_frames_all
Original file line number	Diff line number	Diff line change
Expand Up		@@ -83,4 +83,3 @@ def show_images(target, swap):
		ax2.axis('off')
		ax1.imshow(target[:,:,::-1])
		ax2.imshow(swap[:,:,::-1])