Make the comparison effect of the results more obvious (#3)

- Make the comparison effect of the results more obvious (4fa545c5a3bfba54abaa5e0862805b17047cb758)


Co-authored-by: sky <[email protected]>

app.py

@@ -77,9 +77,9 @@ def predict(image, resolution='1024x1024'):
         preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
     image_preds = []
     for image, pred in zip(images, preds):
-        image_preds.append(
-            cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
-        )
+        image = image.resize(pred.shape[::-1])
+        pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
+        image_preds.append((pred * image).astype(np.uint8))
 
     return image, image_preds[0]