fix canny image

app_haircolor_pix2pix.py

@@ -3,6 +3,7 @@ import gradio as gr
 import time
 import torch
 import numpy as np
+import cv2
 
 from PIL import Image
 from segment_utils import(
@@ -66,7 +67,7 @@ def image_to_image(
     run_task_time = 0
     time_cost_str = ''
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-    canny_image = CannyDetector()(input_image, 384, generate_size)
+    canny_image = custom_canny_detector(input_image)
 
     cond_image = canny_image
     cond_scale = cond_scale1
@@ -88,14 +89,14 @@ def image_to_image(
 
     return generated_image, time_cost_str
 
-def make_inpaint_condition(image, image_mask):
-    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
-    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
+def custom_canny_detector(image):
+    image = np.array(image)
 
-    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
-    image[image_mask > 0.5] = -1.0  # set as masked pixel
-    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
-    image = torch.from_numpy(image)
+    low_threshold = 100
+    high_threshold = 200
+
+    image = cv2.Canny(image, low_threshold, high_threshold)
+    image = Image.fromarray(image)
     return image
 
 def get_time_cost(run_task_time, time_cost_str):