skyatgit commited on
Commit
4fa545c
1 Parent(s): e797135

Make the comparison effect of the results more obvious

Browse files

Make the comparison effect of the results more obvious
Simultaneously fixing the issue of slight image misalignment in the results caused by different sizes of image and image_preds [0]

Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -77,9 +77,9 @@ def predict(image, resolution='1024x1024'):
77
  preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
78
  image_preds = []
79
  for image, pred in zip(images, preds):
80
- image_preds.append(
81
- cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
82
- )
83
 
84
  return image, image_preds[0]
85
 
 
77
  preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
78
  image_preds = []
79
  for image, pred in zip(images, preds):
80
+ image = image.resize(pred.shape[::-1])
81
+ pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
82
+ image_preds.append((pred * image).astype(np.uint8))
83
 
84
  return image, image_preds[0]
85