Update app.py

app.py

@@ -54,6 +54,16 @@ def infer_headswap(source, target):
     def process_img(img, target=False):
         full_frames = np.array(img)[:, :, ::-1]
         dets = app.get(full_frames)
+        if len(dets) == 0:
+            pad_top, pad_bottom, pad_left, pad_right = (
+                full_frames.shape[0] // 2, full_frames.shape[0] // 2,
+                full_frames.shape[1] // 2, full_frames.shape[1] // 2
+            )
+            full_frames = cv2.copyMakeBorder(
+                full_frames, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=0)
+            dets = app.get(full_frames)
+            if len(dets) == 0:
+                raise ValueError(f"no head on {'target' if target else 'source'} image")
         kps = dets[0]['kps']
         wide = wide_crop_face(full_frames, kps, return_M=target)
         if target: