THP2903 commited on
Commit
b0c81df
·
verified ·
1 Parent(s): c3edd3f

Upload folder using huggingface_hub

Browse files
Files changed (5) hide show
  1. Dockerfile +43 -0
  2. GUI.py +217 -0
  3. README.md +259 -12
  4. main.py +3 -14
  5. requirements.txt +1 -0
Dockerfile ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM nvidia/cuda:11.1-base-ubuntu20.04
2
+
3
+ RUN apt update && DEBIAN_FRONTEND=noninteractive apt install git bzip2 wget unzip python3-pip python3-dev cmake libgl1-mesa-dev python-is-python3 libgtk2.0-dev -yq
4
+ ADD . /app
5
+ WORKDIR /app
6
+ RUN cd Face_Enhancement/models/networks/ &&\
7
+ git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch &&\
8
+ cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm . &&\
9
+ cd ../../../
10
+
11
+ RUN cd Global/detection_models &&\
12
+ git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch &&\
13
+ cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm . &&\
14
+ cd ../../
15
+
16
+ RUN cd Face_Detection/ &&\
17
+ wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 &&\
18
+ bzip2 -d shape_predictor_68_face_landmarks.dat.bz2 &&\
19
+ cd ../
20
+
21
+ RUN cd Face_Enhancement/ &&\
22
+ wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Face_Enhancement/checkpoints.zip &&\
23
+ unzip checkpoints.zip &&\
24
+ cd ../ &&\
25
+ cd Global/ &&\
26
+ wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Global/checkpoints.zip &&\
27
+ unzip checkpoints.zip &&\
28
+ rm -f checkpoints.zip &&\
29
+ cd ../
30
+
31
+ RUN pip3 install numpy
32
+
33
+ RUN pip3 install dlib
34
+
35
+ RUN pip3 install -r requirements.txt
36
+
37
+ RUN git clone https://github.com/NVlabs/SPADE.git
38
+
39
+ RUN cd SPADE/ && pip3 install -r requirements.txt
40
+
41
+ RUN cd ..
42
+
43
+ CMD ["python3", "run.py"]
GUI.py ADDED
@@ -0,0 +1,217 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+ import PySimpleGUI as sg
4
+ import os.path
5
+ import argparse
6
+ import os
7
+ import sys
8
+ import shutil
9
+ from subprocess import call
10
+
11
+ def modify(image_filename=None, cv2_frame=None):
12
+
13
+ def run_cmd(command):
14
+ try:
15
+ call(command, shell=True)
16
+ except KeyboardInterrupt:
17
+ print("Process interrupted")
18
+ sys.exit(1)
19
+
20
+ parser = argparse.ArgumentParser()
21
+ parser.add_argument("--input_folder", type=str,
22
+ default= image_filename, help="Test images")
23
+ parser.add_argument(
24
+ "--output_folder",
25
+ type=str,
26
+ default="./output",
27
+ help="Restored images, please use the absolute path",
28
+ )
29
+ parser.add_argument("--GPU", type=str, default="-1", help="0,1,2")
30
+ parser.add_argument(
31
+ "--checkpoint_name", type=str, default="Setting_9_epoch_100", help="choose which checkpoint"
32
+ )
33
+ parser.add_argument("--with_scratch",default="--with_scratch" ,action="store_true")
34
+ opts = parser.parse_args()
35
+
36
+ gpu1 = opts.GPU
37
+
38
+ # resolve relative paths before changing directory
39
+ opts.input_folder = os.path.abspath(opts.input_folder)
40
+ opts.output_folder = os.path.abspath(opts.output_folder)
41
+ if not os.path.exists(opts.output_folder):
42
+ os.makedirs(opts.output_folder)
43
+
44
+ main_environment = os.getcwd()
45
+
46
+ # Stage 1: Overall Quality Improve
47
+ print("Running Stage 1: Overall restoration")
48
+ os.chdir("./Global")
49
+ stage_1_input_dir = opts.input_folder
50
+ stage_1_output_dir = os.path.join(
51
+ opts.output_folder, "stage_1_restore_output")
52
+ if not os.path.exists(stage_1_output_dir):
53
+ os.makedirs(stage_1_output_dir)
54
+
55
+ if not opts.with_scratch:
56
+ stage_1_command = (
57
+ "python test.py --test_mode Full --Quality_restore --test_input "
58
+ + stage_1_input_dir
59
+ + " --outputs_dir "
60
+ + stage_1_output_dir
61
+ + " --gpu_ids "
62
+ + gpu1
63
+ )
64
+ run_cmd(stage_1_command)
65
+ else:
66
+
67
+ mask_dir = os.path.join(stage_1_output_dir, "masks")
68
+ new_input = os.path.join(mask_dir, "input")
69
+ new_mask = os.path.join(mask_dir, "mask")
70
+ stage_1_command_1 = (
71
+ "python detection.py --test_path "
72
+ + stage_1_input_dir
73
+ + " --output_dir "
74
+ + mask_dir
75
+ + " --input_size full_size"
76
+ + " --GPU "
77
+ + gpu1
78
+ )
79
+ stage_1_command_2 = (
80
+ "python test.py --Scratch_and_Quality_restore --test_input "
81
+ + new_input
82
+ + " --test_mask "
83
+ + new_mask
84
+ + " --outputs_dir "
85
+ + stage_1_output_dir
86
+ + " --gpu_ids "
87
+ + gpu1
88
+ )
89
+ run_cmd(stage_1_command_1)
90
+ run_cmd(stage_1_command_2)
91
+
92
+ # Solve the case when there is no face in the old photo
93
+ stage_1_results = os.path.join(stage_1_output_dir, "restored_image")
94
+ stage_4_output_dir = os.path.join(opts.output_folder, "final_output")
95
+ if not os.path.exists(stage_4_output_dir):
96
+ os.makedirs(stage_4_output_dir)
97
+ for x in os.listdir(stage_1_results):
98
+ img_dir = os.path.join(stage_1_results, x)
99
+ shutil.copy(img_dir, stage_4_output_dir)
100
+
101
+ print("Finish Stage 1 ...")
102
+ print("\n")
103
+
104
+ # Stage 2: Face Detection
105
+
106
+ print("Running Stage 2: Face Detection")
107
+ os.chdir(".././Face_Detection")
108
+ stage_2_input_dir = os.path.join(stage_1_output_dir, "restored_image")
109
+ stage_2_output_dir = os.path.join(
110
+ opts.output_folder, "stage_2_detection_output")
111
+ if not os.path.exists(stage_2_output_dir):
112
+ os.makedirs(stage_2_output_dir)
113
+ stage_2_command = (
114
+ "python detect_all_dlib.py --url " + stage_2_input_dir +
115
+ " --save_url " + stage_2_output_dir
116
+ )
117
+ run_cmd(stage_2_command)
118
+ print("Finish Stage 2 ...")
119
+ print("\n")
120
+
121
+ # Stage 3: Face Restore
122
+ print("Running Stage 3: Face Enhancement")
123
+ os.chdir(".././Face_Enhancement")
124
+ stage_3_input_mask = "./"
125
+ stage_3_input_face = stage_2_output_dir
126
+ stage_3_output_dir = os.path.join(
127
+ opts.output_folder, "stage_3_face_output")
128
+ if not os.path.exists(stage_3_output_dir):
129
+ os.makedirs(stage_3_output_dir)
130
+ stage_3_command = (
131
+ "python test_face.py --old_face_folder "
132
+ + stage_3_input_face
133
+ + " --old_face_label_folder "
134
+ + stage_3_input_mask
135
+ + " --tensorboard_log --name "
136
+ + opts.checkpoint_name
137
+ + " --gpu_ids "
138
+ + gpu1
139
+ + " --load_size 256 --label_nc 18 --no_instance --preprocess_mode resize --batchSize 4 --results_dir "
140
+ + stage_3_output_dir
141
+ + " --no_parsing_map"
142
+ )
143
+ run_cmd(stage_3_command)
144
+ print("Finish Stage 3 ...")
145
+ print("\n")
146
+
147
+ # Stage 4: Warp back
148
+ print("Running Stage 4: Blending")
149
+ os.chdir(".././Face_Detection")
150
+ stage_4_input_image_dir = os.path.join(
151
+ stage_1_output_dir, "restored_image")
152
+ stage_4_input_face_dir = os.path.join(stage_3_output_dir, "each_img")
153
+ stage_4_output_dir = os.path.join(opts.output_folder, "final_output")
154
+ if not os.path.exists(stage_4_output_dir):
155
+ os.makedirs(stage_4_output_dir)
156
+ stage_4_command = (
157
+ "python align_warp_back_multiple_dlib.py --origin_url "
158
+ + stage_4_input_image_dir
159
+ + " --replace_url "
160
+ + stage_4_input_face_dir
161
+ + " --save_url "
162
+ + stage_4_output_dir
163
+ )
164
+ run_cmd(stage_4_command)
165
+ print("Finish Stage 4 ...")
166
+ print("\n")
167
+
168
+ print("All the processing is done. Please check the results.")
169
+
170
+ # --------------------------------- The GUI ---------------------------------
171
+
172
+ # First the window layout...
173
+
174
+ images_col = [[sg.Text('Input file:'), sg.In(enable_events=True, key='-IN FILE-'), sg.FileBrowse()],
175
+ [sg.Button('Modify Photo', key='-MPHOTO-'), sg.Button('Exit')],
176
+ [sg.Image(filename='', key='-IN-'), sg.Image(filename='', key='-OUT-')],]
177
+ # ----- Full layout -----
178
+ layout = [[sg.VSeperator(), sg.Column(images_col)]]
179
+
180
+ # ----- Make the window -----
181
+ window = sg.Window('Bringing-old-photos-back-to-life', layout, grab_anywhere=True)
182
+
183
+ # ----- Run the Event Loop -----
184
+ prev_filename = colorized = cap = None
185
+ while True:
186
+ event, values = window.read()
187
+ if event in (None, 'Exit'):
188
+ break
189
+
190
+ elif event == '-MPHOTO-':
191
+ try:
192
+ n1 = filename.split("/")[-2]
193
+ n2 = filename.split("/")[-3]
194
+ n3 = filename.split("/")[-1]
195
+ filename= str(f"./{n2}/{n1}")
196
+ modify(filename)
197
+
198
+ global f_image
199
+ f_image = f'./output/final_output/{n3}'
200
+ image = cv2.imread(f_image)
201
+ window['-OUT-'].update(data=cv2.imencode('.png', image)[1].tobytes())
202
+
203
+ except:
204
+ continue
205
+
206
+ elif event == '-IN FILE-': # A single filename was chosen
207
+ filename = values['-IN FILE-']
208
+ if filename != prev_filename:
209
+ prev_filename = filename
210
+ try:
211
+ image = cv2.imread(filename)
212
+ window['-IN-'].update(data=cv2.imencode('.png', image)[1].tobytes())
213
+ except:
214
+ continue
215
+
216
+ # ----- Exit program -----
217
+ window.close()
README.md CHANGED
@@ -1,12 +1,259 @@
1
- ---
2
- title: Phục hồi ảnh & mô tả triết học
3
- emoji: 🖼️
4
- colorFrom: blue
5
- colorTo: green
6
- sdk: gradio
7
- sdk_version: 5.31.0
8
- app_file: app.py
9
- pinned: false
10
- ---
11
-
12
- Upload ảnh cũ bị hư để phục hồi và nhận mô tả triết học từ AI.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Old Photo Restoration (Official PyTorch Implementation)
2
+
3
+ <img src='imgs/0001.jpg'/>
4
+
5
+ ### [Project Page](http://raywzy.com/Old_Photo/) | [Paper (CVPR version)](https://arxiv.org/abs/2004.09484) | [Paper (Journal version)](https://arxiv.org/pdf/2009.07047v1.pdf) | [Pretrained Model](https://hkustconnect-my.sharepoint.com/:f:/g/personal/bzhangai_connect_ust_hk/Em0KnYOeSSxFtp4g_dhWdf0BdeT3tY12jIYJ6qvSf300cA?e=nXkJH2) | [Colab Demo](https://colab.research.google.com/drive/1NEm6AsybIiC5TwTU_4DqDkQO0nFRB-uA?usp=sharing) | [Replicate Demo & Docker Image](https://replicate.ai/zhangmozhe/bringing-old-photos-back-to-life) :fire:
6
+
7
+ **Bringing Old Photos Back to Life, CVPR2020 (Oral)**
8
+
9
+ **Old Photo Restoration via Deep Latent Space Translation, TPAMI 2022**
10
+
11
+ [Ziyu Wan](http://raywzy.com/)<sup>1</sup>,
12
+ [Bo Zhang](https://www.microsoft.com/en-us/research/people/zhanbo/)<sup>2</sup>,
13
+ [Dongdong Chen](http://www.dongdongchen.bid/)<sup>3</sup>,
14
+ [Pan Zhang](https://panzhang0212.github.io/)<sup>4</sup>,
15
+ [Dong Chen](https://www.microsoft.com/en-us/research/people/doch/)<sup>2</sup>,
16
+ [Jing Liao](https://liaojing.github.io/html/)<sup>1</sup>,
17
+ [Fang Wen](https://www.microsoft.com/en-us/research/people/fangwen/)<sup>2</sup> <br>
18
+ <sup>1</sup>City University of Hong Kong, <sup>2</sup>Microsoft Research Asia, <sup>3</sup>Microsoft Cloud AI, <sup>4</sup>USTC
19
+
20
+ <!-- ## Notes of this project
21
+ The code originates from our research project and the aim is to demonstrate the research idea, so we have not optimized it from a product perspective. And we will spend time to address some common issues, such as out of memory issue, limited resolution, but will not involve too much in engineering problems, such as speedup of the inference, fastapi deployment and so on. **We welcome volunteers to contribute to this project to make it more usable for practical application.** -->
22
+
23
+ ## :sparkles: News
24
+ **2022.3.31**: Our new work regarding old film restoration will be published in CVPR 2022. For more details, please refer to the [project website](http://raywzy.com/Old_Film/) and [github repo](https://github.com/raywzy/Bringing-Old-Films-Back-to-Life).
25
+
26
+ The framework now supports the restoration of high-resolution input.
27
+
28
+ <img src='imgs/HR_result.png'>
29
+
30
+ Training code is available and welcome to have a try and learn the training details.
31
+
32
+ You can now play with our [Colab](https://colab.research.google.com/drive/1NEm6AsybIiC5TwTU_4DqDkQO0nFRB-uA?usp=sharing) and try it on your photos.
33
+
34
+ ## Requirement
35
+ The code is tested on Ubuntu with Nvidia GPUs and CUDA installed. Python>=3.6 is required to run the code.
36
+
37
+ ## Installation
38
+
39
+ Clone the Synchronized-BatchNorm-PyTorch repository for
40
+
41
+ ```
42
+ cd Face_Enhancement/models/networks/
43
+ git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
44
+ cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm .
45
+ cd ../../../
46
+ ```
47
+
48
+ ```
49
+ cd Global/detection_models
50
+ git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
51
+ cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm .
52
+ cd ../../
53
+ ```
54
+
55
+ Download the landmark detection pretrained model
56
+
57
+ ```
58
+ cd Face_Detection/
59
+ wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
60
+ bzip2 -d shape_predictor_68_face_landmarks.dat.bz2
61
+ cd ../
62
+ ```
63
+
64
+ Download the pretrained model, put the file `Face_Enhancement/checkpoints.zip` under `./Face_Enhancement`, and put the file `Global/checkpoints.zip` under `./Global`. Then unzip them respectively.
65
+
66
+ ```
67
+ cd Face_Enhancement/
68
+ wget https://github.com/microsoft/Bringing-Old-Photos-Back-to-Life/releases/download/v1.0/face_checkpoints.zip
69
+ unzip face_checkpoints.zip
70
+ cd ../
71
+ cd Global/
72
+ wget https://github.com/microsoft/Bringing-Old-Photos-Back-to-Life/releases/download/v1.0/global_checkpoints.zip
73
+ unzip global_checkpoints.zip
74
+ cd ../
75
+ ```
76
+
77
+ Install dependencies:
78
+
79
+ ```bash
80
+ pip install -r requirements.txt
81
+ ```
82
+
83
+ ## :rocket: How to use?
84
+
85
+ **Note**: GPU can be set 0 or 0,1,2 or 0,2; use -1 for CPU
86
+
87
+ ### 1) Full Pipeline
88
+
89
+ You could easily restore the old photos with one simple command after installation and downloading the pretrained model.
90
+
91
+ For images without scratches:
92
+
93
+ ```
94
+ python run.py --input_folder [test_image_folder_path] \
95
+ --output_folder [output_path] \
96
+ --GPU 0
97
+ ```
98
+
99
+ For scratched images:
100
+
101
+ ```
102
+ python run.py --input_folder [test_image_folder_path] \
103
+ --output_folder [output_path] \
104
+ --GPU 0 \
105
+ --with_scratch
106
+ ```
107
+
108
+ **For high-resolution images with scratches**:
109
+
110
+ ```
111
+ python run.py --input_folder [test_image_folder_path] \
112
+ --output_folder [output_path] \
113
+ --GPU 0 \
114
+ --with_scratch \
115
+ --HR
116
+ ```
117
+
118
+ Note: Please try to use the absolute path. The final results will be saved in `./output_path/final_output/`. You could also check the produced results of different steps in `output_path`.
119
+
120
+ ### 2) Scratch Detection
121
+
122
+ Currently we don't plan to release the scratched old photos dataset with labels directly. If you want to get the paired data, you could use our pretrained model to test the collected images to obtain the labels.
123
+
124
+ ```
125
+ cd Global/
126
+ python detection.py --test_path [test_image_folder_path] \
127
+ --output_dir [output_path] \
128
+ --input_size [resize_256|full_size|scale_256]
129
+ ```
130
+
131
+ <img src='imgs/scratch_detection.png'>
132
+
133
+ ### 3) Global Restoration
134
+
135
+ A triplet domain translation network is proposed to solve both structured degradation and unstructured degradation of old photos.
136
+
137
+ <p align="center">
138
+ <img src='imgs/pipeline.PNG' width="50%" height="50%"/>
139
+ </p>
140
+
141
+ ```
142
+ cd Global/
143
+ python test.py --Scratch_and_Quality_restore \
144
+ --test_input [test_image_folder_path] \
145
+ --test_mask [corresponding mask] \
146
+ --outputs_dir [output_path]
147
+
148
+ python test.py --Quality_restore \
149
+ --test_input [test_image_folder_path] \
150
+ --outputs_dir [output_path]
151
+ ```
152
+
153
+ <img src='imgs/global.png'>
154
+
155
+
156
+ ### 4) Face Enhancement
157
+
158
+ We use a progressive generator to refine the face regions of old photos. More details could be found in our journal submission and `./Face_Enhancement` folder.
159
+
160
+ <p align="center">
161
+ <img src='imgs/face_pipeline.jpg' width="60%" height="60%"/>
162
+ </p>
163
+
164
+
165
+ <img src='imgs/face.png'>
166
+
167
+ > *NOTE*:
168
+ > This repo is mainly for research purpose and we have not yet optimized the running performance.
169
+ >
170
+ > Since the model is pretrained with 256*256 images, the model may not work ideally for arbitrary resolution.
171
+
172
+ ### 5) GUI
173
+
174
+ A user-friendly GUI which takes input of image by user and shows result in respective window.
175
+
176
+ #### How it works:
177
+
178
+ 1. Run GUI.py file.
179
+ 2. Click browse and select your image from test_images/old_w_scratch folder to remove scratches.
180
+ 3. Click Modify Photo button.
181
+ 4. Wait for a while and see results on GUI window.
182
+ 5. Exit window by clicking Exit Window and get your result image in output folder.
183
+
184
+ <img src='imgs/gui.PNG'>
185
+
186
+ ## How to train?
187
+
188
+ ### 1) Create Training File
189
+
190
+ Put the folders of VOC dataset, collected old photos (e.g., Real_L_old and Real_RGB_old) into one shared folder. Then
191
+ ```
192
+ cd Global/data/
193
+ python Create_Bigfile.py
194
+ ```
195
+ Note: Remember to modify the code based on your own environment.
196
+
197
+ ### 2) Train the VAEs of domain A and domain B respectively
198
+
199
+ ```
200
+ cd ..
201
+ python train_domain_A.py --use_v2_degradation --continue_train --training_dataset domain_A --name domainA_SR_old_photos --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 100 --no_html --gpu_ids 0,1,2,3 --self_gen --nThreads 4 --n_downsample_global 3 --k_size 4 --use_v2 --mc 64 --start_r 1 --kl 1 --no_cgan --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder]
202
+
203
+ python train_domain_B.py --continue_train --training_dataset domain_B --name domainB_old_photos --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 120 --no_html --gpu_ids 0,1,2,3 --self_gen --nThreads 4 --n_downsample_global 3 --k_size 4 --use_v2 --mc 64 --start_r 1 --kl 1 --no_cgan --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder]
204
+ ```
205
+ Note: For the --name option, please ensure your experiment name contains "domainA" or "domainB", which will be used to select different dataset.
206
+
207
+ ### 3) Train the mapping network between domains
208
+
209
+ Train the mapping without scratches:
210
+ ```
211
+ python train_mapping.py --use_v2_degradation --training_dataset mapping --use_vae_which_epoch 200 --continue_train --name mapping_quality --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 80 --no_html --gpu_ids 0,1,2,3 --nThreads 8 --load_pretrainA [ckpt_of_domainA_SR_old_photos] --load_pretrainB [ckpt_of_domainB_old_photos] --l2_feat 60 --n_downsample_global 3 --mc 64 --k_size 4 --start_r 1 --mapping_n_block 6 --map_mc 512 --use_l1_feat --niter 150 --niter_decay 100 --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder]
212
+ ```
213
+
214
+
215
+ Traing the mapping with scraches:
216
+ ```
217
+ python train_mapping.py --no_TTUR --NL_res --random_hole --use_SN --correlation_renormalize --training_dataset mapping --NL_use_mask --NL_fusion_method combine --non_local Setting_42 --use_v2_degradation --use_vae_which_epoch 200 --continue_train --name mapping_scratch --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 36 --no_html --gpu_ids 0,1,2,3 --nThreads 8 --load_pretrainA [ckpt_of_domainA_SR_old_photos] --load_pretrainB [ckpt_of_domainB_old_photos] --l2_feat 60 --n_downsample_global 3 --mc 64 --k_size 4 --start_r 1 --mapping_n_block 6 --map_mc 512 --use_l1_feat --niter 150 --niter_decay 100 --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder] --irregular_mask [absolute_path_of_mask_file]
218
+ ```
219
+
220
+ Traing the mapping with scraches (Multi-Scale Patch Attention for HR input):
221
+ ```
222
+ python train_mapping.py --no_TTUR --NL_res --random_hole --use_SN --correlation_renormalize --training_dataset mapping --NL_use_mask --NL_fusion_method combine --non_local Setting_42 --use_v2_degradation --use_vae_which_epoch 200 --continue_train --name mapping_Patch_Attention --label_nc 0 --loadSize 256 --fineSize 256 --dataroot [your_data_folder] --no_instance --resize_or_crop crop_only --batchSize 36 --no_html --gpu_ids 0,1,2,3 --nThreads 8 --load_pretrainA [ckpt_of_domainA_SR_old_photos] --load_pretrainB [ckpt_of_domainB_old_photos] --l2_feat 60 --n_downsample_global 3 --mc 64 --k_size 4 --start_r 1 --mapping_n_block 6 --map_mc 512 --use_l1_feat --niter 150 --niter_decay 100 --outputs_dir [your_output_folder] --checkpoints_dir [your_ckpt_folder] --irregular_mask [absolute_path_of_mask_file] --mapping_exp 1
223
+ ```
224
+
225
+
226
+ ## Citation
227
+
228
+ If you find our work useful for your research, please consider citing the following papers :)
229
+
230
+ ```bibtex
231
+ @inproceedings{wan2020bringing,
232
+ title={Bringing Old Photos Back to Life},
233
+ author={Wan, Ziyu and Zhang, Bo and Chen, Dongdong and Zhang, Pan and Chen, Dong and Liao, Jing and Wen, Fang},
234
+ booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
235
+ pages={2747--2757},
236
+ year={2020}
237
+ }
238
+ ```
239
+
240
+ ```bibtex
241
+ @article{wan2020old,
242
+ title={Old Photo Restoration via Deep Latent Space Translation},
243
+ author={Wan, Ziyu and Zhang, Bo and Chen, Dongdong and Zhang, Pan and Chen, Dong and Liao, Jing and Wen, Fang},
244
+ journal={arXiv preprint arXiv:2009.07047},
245
+ year={2020}
246
+ }
247
+ ```
248
+
249
+ If you are also interested in the legacy photo/video colorization, please refer to [this work](https://github.com/zhangmozhe/video-colorization).
250
+
251
+ ## Maintenance
252
+
253
+ This project is currently maintained by Ziyu Wan and is for academic research use only. If you have any questions, feel free to contact raywzy@gmail.com.
254
+
255
+ ## License
256
+
257
+ The codes and the pretrained model in this repository are under the MIT license as specified by the LICENSE file. We use our labeled dataset to train the scratch detection model.
258
+
259
+ This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
main.py CHANGED
@@ -11,7 +11,7 @@
11
  # ]
12
 
13
  # API_URL = "https://generativelanguage.googleapis.com/v1beta"
14
-
15
  # WORKING_MODELS = [
16
  # "models/gemini-1.5-flash-latest",
17
  # "models/gemini-1.5-flash-001",
@@ -200,18 +200,7 @@ def try_models_on_image(image_path):
200
  print("\n🚫 All models failed or all API keys quota exceeded.")
201
  return "" # Return empty string if all attempts fail
202
 
203
- # if __name__ == "__main__":
204
- # image_path = "/home/nhattan05022003/coding/SEM_8/MLN_111/photo_restoration/output_img_folder/final_output/monalisa.png"
205
- # description = try_models_on_image(image_path)
206
- # print(f"Final description: {description}")
207
-
208
  if __name__ == "__main__":
209
- import sys
210
-
211
- if len(sys.argv) < 2:
212
- print("❌ Không truyền đường dẫn ảnh.")
213
- sys.exit(1)
214
-
215
- image_path = sys.argv[1]
216
  description = try_models_on_image(image_path)
217
- print(description)
 
11
  # ]
12
 
13
  # API_URL = "https://generativelanguage.googleapis.com/v1beta"
14
+
15
  # WORKING_MODELS = [
16
  # "models/gemini-1.5-flash-latest",
17
  # "models/gemini-1.5-flash-001",
 
200
  print("\n🚫 All models failed or all API keys quota exceeded.")
201
  return "" # Return empty string if all attempts fail
202
 
 
 
 
 
 
203
  if __name__ == "__main__":
204
+ image_path = "/home/nhattan05022003/coding/SEM_8/MLN_111/photo_restoration/output_img_folder/final_output/monalisa.png"
 
 
 
 
 
 
205
  description = try_models_on_image(image_path)
206
+ print(f"Final description: {description}")
requirements.txt CHANGED
@@ -1,5 +1,6 @@
1
  torch
2
  torchvision
 
3
  scikit-image
4
  easydict
5
  PyYAML
 
1
  torch
2
  torchvision
3
+ dlib
4
  scikit-image
5
  easydict
6
  PyYAML