sczhou commited on
Commit
f26e1dd
1 Parent(s): ae068c5

add global try.

Browse files
Files changed (1) hide show
  1. app.py +92 -89
app.py CHANGED
@@ -103,98 +103,101 @@ os.makedirs('output', exist_ok=True)
103
 
104
  def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
105
  """Run a single prediction on the model"""
106
- # take the default setting for the demo
107
- has_aligned = False
108
- only_center_face = False
109
- draw_box = False
110
- detection_model = "retinaface_resnet50"
111
-
112
- upscale = int(upscale) # covert type to int
113
- face_helper = FaceRestoreHelper(
114
- upscale,
115
- face_size=512,
116
- crop_ratio=(1, 1),
117
- det_model=detection_model,
118
- save_ext="png",
119
- use_parse=True,
120
- device=device,
121
- )
122
- bg_upsampler = upsampler if background_enhance else None
123
- face_upsampler = upsampler if face_upsample else None
124
-
125
- img = cv2.imread(str(image), cv2.IMREAD_COLOR)
126
-
127
- if has_aligned:
128
- # the input faces are already cropped and aligned
129
- img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
130
- face_helper.is_gray = is_gray(img, threshold=5)
131
- if face_helper.is_gray:
132
- print('Grayscale input: True')
133
- face_helper.cropped_faces = [img]
134
- else:
135
- face_helper.read_image(img)
136
- # get face landmarks for each face
137
- num_det_faces = face_helper.get_face_landmarks_5(
138
- only_center_face=only_center_face, resize=640, eye_dist_threshold=5
139
- )
140
- print(f"\tdetect {num_det_faces} faces")
141
- # align and warp each face
142
- face_helper.align_warp_face()
143
-
144
- # face restoration for each cropped face
145
- for idx, cropped_face in enumerate(face_helper.cropped_faces):
146
- # prepare data
147
- cropped_face_t = img2tensor(
148
- cropped_face / 255.0, bgr2rgb=True, float32=True
149
  )
150
- normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
151
- cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
152
-
153
- try:
154
- with torch.no_grad():
155
- output = codeformer_net(
156
- cropped_face_t, w=codeformer_fidelity, adain=True
157
- )[0]
158
- restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
159
- del output
160
- torch.cuda.empty_cache()
161
- except Exception as error:
162
- print(f"\tFailed inference for CodeFormer: {error}")
163
- restored_face = tensor2img(
164
- cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
165
- )
166
-
167
- restored_face = restored_face.astype("uint8")
168
- face_helper.add_restored_face(restored_face)
169
-
170
- # paste_back
171
- if not has_aligned:
172
- # upsample the background
173
- if bg_upsampler is not None:
174
- # Now only support RealESRGAN for upsampling background
175
- bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
176
  else:
177
- bg_img = None
178
- face_helper.get_inverse_affine(None)
179
- # paste each restored face to the input image
180
- if face_upsample and face_upsampler is not None:
181
- restored_img = face_helper.paste_faces_to_input_image(
182
- upsample_img=bg_img,
183
- draw_box=draw_box,
184
- face_upsampler=face_upsampler,
185
  )
186
- else:
187
- restored_img = face_helper.paste_faces_to_input_image(
188
- upsample_img=bg_img, draw_box=draw_box
 
 
 
 
 
 
189
  )
190
-
191
- # save restored img
192
- save_path = f'output/out.png'
193
- imwrite(restored_img, str(save_path))
194
-
195
- restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
196
- return restored_img, save_path
197
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
198
 
199
 
200
  title = "CodeFormer: Robust Face Restoration and Enhancement Network"
@@ -230,7 +233,7 @@ Redistribution and use for non-commercial purposes should follow this license.
230
 
231
  If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.
232
 
233
- ![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer)
234
  """
235
 
236
  demo = gr.Interface(
 
103
 
104
  def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
105
  """Run a single prediction on the model"""
106
+ try: # global try
107
+ # take the default setting for the demo
108
+ has_aligned = False
109
+ only_center_face = False
110
+ draw_box = False
111
+ detection_model = "retinaface_resnet50"
112
+
113
+ upscale = int(upscale) # covert type to int
114
+ face_helper = FaceRestoreHelper(
115
+ upscale,
116
+ face_size=512,
117
+ crop_ratio=(1, 1),
118
+ det_model=detection_model,
119
+ save_ext="png",
120
+ use_parse=True,
121
+ device=device,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
122
  )
123
+ bg_upsampler = upsampler if background_enhance else None
124
+ face_upsampler = upsampler if face_upsample else None
125
+
126
+ img = cv2.imread(str(image), cv2.IMREAD_COLOR)
127
+
128
+ if has_aligned:
129
+ # the input faces are already cropped and aligned
130
+ img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
131
+ face_helper.is_gray = is_gray(img, threshold=5)
132
+ if face_helper.is_gray:
133
+ print('Grayscale input: True')
134
+ face_helper.cropped_faces = [img]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
135
  else:
136
+ face_helper.read_image(img)
137
+ # get face landmarks for each face
138
+ num_det_faces = face_helper.get_face_landmarks_5(
139
+ only_center_face=only_center_face, resize=640, eye_dist_threshold=5
 
 
 
 
140
  )
141
+ print(f"\tdetect {num_det_faces} faces")
142
+ # align and warp each face
143
+ face_helper.align_warp_face()
144
+
145
+ # face restoration for each cropped face
146
+ for idx, cropped_face in enumerate(face_helper.cropped_faces):
147
+ # prepare data
148
+ cropped_face_t = img2tensor(
149
+ cropped_face / 255.0, bgr2rgb=True, float32=True
150
  )
151
+ normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
152
+ cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
153
+
154
+ try:
155
+ with torch.no_grad():
156
+ output = codeformer_net(
157
+ cropped_face_t, w=codeformer_fidelity, adain=True
158
+ )[0]
159
+ restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
160
+ del output
161
+ torch.cuda.empty_cache()
162
+ except Exception as error:
163
+ print(f"\tFailed inference for CodeFormer: {error}")
164
+ restored_face = tensor2img(
165
+ cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
166
+ )
167
+
168
+ restored_face = restored_face.astype("uint8")
169
+ face_helper.add_restored_face(restored_face)
170
+
171
+ # paste_back
172
+ if not has_aligned:
173
+ # upsample the background
174
+ if bg_upsampler is not None:
175
+ # Now only support RealESRGAN for upsampling background
176
+ bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
177
+ else:
178
+ bg_img = None
179
+ face_helper.get_inverse_affine(None)
180
+ # paste each restored face to the input image
181
+ if face_upsample and face_upsampler is not None:
182
+ restored_img = face_helper.paste_faces_to_input_image(
183
+ upsample_img=bg_img,
184
+ draw_box=draw_box,
185
+ face_upsampler=face_upsampler,
186
+ )
187
+ else:
188
+ restored_img = face_helper.paste_faces_to_input_image(
189
+ upsample_img=bg_img, draw_box=draw_box
190
+ )
191
+
192
+ # save restored img
193
+ save_path = f'output/out.png'
194
+ imwrite(restored_img, str(save_path))
195
+
196
+ restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
197
+ return restored_img, save_path
198
+ except Exception as error:
199
+ print('global exception', error)
200
+ return None, None
201
 
202
 
203
  title = "CodeFormer: Robust Face Restoration and Enhancement Network"
 
233
 
234
  If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.
235
 
236
+ ![visitors](https://visitor-badge.laobi.icu/badge?page_id=sczhou/CodeFormer)
237
  """
238
 
239
  demo = gr.Interface(