(model, points, max_iters, state, size, mask, lr_box)
| 69 | |
| 70 | |
| 71 | def on_drag(model, points, max_iters, state, size, mask, lr_box): |
| 72 | if len(points['handle']) == 0: |
| 73 | raise gr.Error('You must select at least one handle point and target point.') |
| 74 | if len(points['handle']) != len(points['target']): |
| 75 | raise gr.Error('You have uncompleted handle points, try to selct a target point or undo the handle point.') |
| 76 | max_iters = int(max_iters) |
| 77 | W = state['W'] |
| 78 | |
| 79 | handle_points = [torch.tensor(p, device=device).float() for p in points['handle']] |
| 80 | target_points = [torch.tensor(p, device=device).float() for p in points['target']] |
| 81 | |
| 82 | if mask.get('mask') is not None: |
| 83 | mask = Image.fromarray(mask['mask']).convert('L') |
| 84 | mask = np.array(mask) == 255 |
| 85 | |
| 86 | mask = torch.from_numpy(mask).float().to(device) |
| 87 | mask = mask.unsqueeze(0).unsqueeze(0) |
| 88 | else: |
| 89 | mask = None |
| 90 | |
| 91 | step = 0 |
| 92 | for image, W, handle_points in drag_gan(W, model['G'], |
| 93 | handle_points, target_points, mask, |
| 94 | max_iters=max_iters, lr=lr_box): |
| 95 | points['handle'] = [p.cpu().numpy().astype('int') for p in handle_points] |
| 96 | image = add_points_to_image(image, points, size=SIZE_TO_CLICK_SIZE[size]) |
| 97 | |
| 98 | state['history'].append(image) |
| 99 | step += 1 |
| 100 | yield image, state, step |
| 101 | |
| 102 | |
| 103 | def on_reset(points, image, state): |
nothing calls this directly
no test coverage detected