【源码】基于Python编程语言的图像生成项目

项目简介

本项目借助输入道路与河川的草图，生成对应的真实影像。基于PITI（Pretraining is All You Need for Image-to-Image Translation）模型开发，以生成影像与真实影像间的FID score作为评分标准，在相关竞赛中荣获第二名。

项目的主要特性和功能

1. 草图到图像生成：依据输入的草图生成对应的真实影像。
2. 多任务训练：分别对河道和道路生成任务进行模型训练。
3. 超分辨率模型：把生成的64x64大小的图像转换为256x256的高分辨率图像。
4. 模型微调：提供多个阶段的模型微调策略，涵盖encoder和decoder的微调。
5. 模型推理：支持从训练好的模型中生成图像，同时提供图像筛选和格式转换功能。

安装使用步骤

环境准备

1. 操作系统：Linux - 64
2. 硬件要求：
3. CPU: 12 核心 Intel(R) Xeon(R) CPU @ 2.20GHz
4. GPU: Nvidia L4 (23 GB)
5. Python版本：3.7.11

环境安装

bash conda env create -f environment.yml

数据准备

1. 下载数据集：将比赛数据集放置于dataset/路径下，包含训练数据和测试数据。
2. 数据预处理： bash python preprocess/aicup_preprocess.py --root dataset/Training_dataset --save_dir dataset --test_size 0.1

模型训练

1. 下载预训练权重：在项目路径下创建pretrained/文件夹，并下载所有预训练权重放置于此。
2. 训练sketch - to - image模型：
3. 河道模型： bash bash sketch_finetune_base.sh river
4. 道路模型： bash bash sketch_finetune_base.sh road
5. 训练upsample模型：
6. 河道模型： bash bash sketch_finetune_upsample.sh river
7. 道路模型： bash bash sketch_finetune_upsample.sh road

模型推理

1. 下载模型权重：从提供的链接下载所有模型权重。
2. 生成河流影像： bash python inference.py --encoder_ckpt path/to/river/sketch/to/image/encoder --decoder_ckpt path/to/river/sketch/to/image/decoder --sr_model_path path/to/river/upsample/model --image "dataset/Private_testing_dataset/label_img/PRI_RO*.png" --save_dir ./generate_images --num_samples 8
3. 生成道路影像： bash python inference.py --encoder_ckpt path/to/road/sketch/to/image/encoder --decoder_ckpt path/to/road/sketch/to/image/decoder --sr_model_path path/to/road/upsample/model --image "dataset/Private_testing_dataset/label_img/PRI_RI*.png" --save_dir ./generate_images --num_samples 8
4. 筛选合适的影像： bash python optimize_fid.py --real_dir dataset/Training_dataset/img --gen_dir ./generate_images --save_dir ./select_images --batch_size 256 --num_sample 8 --freq 2 --model resnet34
5. 转换为提交格式： bash python convert_to_submit.py --image_dir select_images --save_dir ./submit

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】