|
from pathlib import Path |
|
from typing import Optional |
|
from tqdm import tqdm |
|
|
|
import hydra |
|
from omegaconf import DictConfig |
|
import numpy as np |
|
|
|
from src.simswap import SimSwap |
|
from src.DataManager.ImageDataManager import ImageDataManager |
|
from src.DataManager.VideoDataManager import VideoDataManager |
|
from src.DataManager.utils import imread_rgb |
|
|
|
|
|
class Application: |
|
def __init__(self, config: DictConfig): |
|
|
|
id_image_path = Path(config.data.id_image) |
|
specific_id_image_path = Path(config.data.specific_id_image) |
|
att_image_path = Path(config.data.att_image) |
|
att_video_path = Path(config.data.att_video) |
|
output_dir = Path(config.data.output_dir) |
|
|
|
assert id_image_path.exists(), f"Can't find {id_image_path} file!" |
|
|
|
self.id_image: Optional[np.ndarray] = imread_rgb(id_image_path) |
|
self.specific_id_image: Optional[np.ndarray] = ( |
|
imread_rgb(specific_id_image_path) |
|
if specific_id_image_path and specific_id_image_path.is_file() |
|
else None |
|
) |
|
|
|
self.att_image: Optional[ImageDataManager] = None |
|
if att_image_path and (att_image_path.is_file() or att_image_path.is_dir()): |
|
self.att_image: Optional[ImageDataManager] = ImageDataManager( |
|
src_data=att_image_path, output_dir=output_dir |
|
) |
|
|
|
self.att_video: Optional[VideoDataManager] = None |
|
if att_video_path and att_video_path.is_file(): |
|
self.att_video: Optional[VideoDataManager] = VideoDataManager( |
|
src_data=att_video_path, output_dir=output_dir, clean_work_dir=config.data.clean_work_dir |
|
) |
|
|
|
assert not (self.att_video and self.att_image), "Only one attribute source can be used!" |
|
|
|
self.data_manager = self.att_video if self.att_video else self.att_image |
|
|
|
self.model = SimSwap( |
|
config=config.pipeline, |
|
id_image=self.id_image, |
|
specific_image=self.specific_id_image, |
|
) |
|
|
|
def run(self): |
|
for _ in tqdm(range(len(self.data_manager))): |
|
|
|
att_img = self.data_manager.get() |
|
|
|
output = self.model(att_img) |
|
|
|
self.data_manager.save(output) |
|
|
|
|
|
@hydra.main(config_path="configs/", config_name="run_image.yaml") |
|
def main(config: DictConfig): |
|
|
|
app = Application(config) |
|
|
|
app.run() |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
