Update app.py

app.py

@@ -69,11 +69,10 @@ COLORS = [
 def process(query_img, state,
             cfg_path='configs/demo_b.py'):
     cfg = Config.fromfile(cfg_path)
+    width, height, _ = state['original_support_image'].shape
     kp_src_np = np.array(state['kp_src']).copy().astype(np.float32)
-    kp_src_np[:, 0] = kp_src_np[:,
-                      0] / 128. * cfg.model.encoder_config.img_size
-    kp_src_np[:, 1] = kp_src_np[:,
-                      1] / 128. * cfg.model.encoder_config.img_size
+    kp_src_np[:, 0] = kp_src_np[:,0] / (width // 2) * cfg.model.encoder_config.img_size
+    kp_src_np[:, 1] = kp_src_np[:,1] / (height // 2) * cfg.model.encoder_config.img_size
     kp_src_np = np.flip(kp_src_np, 1).copy()
     kp_src_tensor = torch.tensor(kp_src_np).float()
     preprocess = transforms.Compose([
@@ -83,7 +82,7 @@ def process(query_img, state,
                    cfg.model.encoder_config.img_size)])
 
     if len(state['skeleton']) == 0:
-        skeleton = [(0, 0)]
+        state['skeleton'] = [(0, 0)]
 
     support_img = preprocess(state['original_support_image']).flip(0)[None]
     np_query = np.array(query_img)[:, :, ::-1].copy()
@@ -296,7 +295,8 @@ with gr.Blocks() as demo:
         state['skeleton'].clear()
         state['kp_src'].clear()
         state['original_support_image'] = np.array(support_img)[:, :, ::-1].copy()
-        support_img = support_img.resize((128, 128), Image.Resampling.LANCZOS)
+        width, height = support_img.size
+        support_img = support_img.resize((width // 2, width // 2), Image.Resampling.LANCZOS)
         return support_img, support_img, state