Spaces:
Runtime error
Runtime error
| import warnings | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from modules.control.util import HWC3, resize_image | |
| ed = None | |
| """ | |
| PFmode: bool | |
| EdgeDetectionOperator: int | |
| GradientThresholdValue: int | |
| AnchorThresholdValue: int | |
| ScanInterval: int | |
| MinPathLength: int | |
| Sigma: float | |
| SumFlag: bool | |
| NFAValidation: bool | |
| MinLineLength: int | |
| MaxDistanceBetweenTwoLines: float | |
| LineFitErrorThreshold: float | |
| MaxErrorThreshold: float | |
| """ | |
| class EdgeDetector: | |
| def __call__(self, input_image=None, pf=True, mode='edge', detect_resolution=512, image_resolution=512, output_type=None, **kwargs): | |
| global ed # pylint: disable=global-statement | |
| if ed is None: | |
| try: | |
| ed = cv2.ximgproc.createEdgeDrawing() | |
| except Exception as e: | |
| raise ImportError("Edge processor: invalid version of OpenCV found") from e | |
| params = cv2.ximgproc.EdgeDrawing.Params() | |
| params.PFmode = pf | |
| ed.setParams(params) | |
| if "img" in kwargs: | |
| warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning) | |
| input_image = kwargs.pop("img") | |
| if input_image is None: | |
| raise ValueError("input_image must be defined.") | |
| if not isinstance(input_image, np.ndarray): | |
| input_image = np.array(input_image, dtype=np.uint8) | |
| output_type = output_type or "pil" | |
| else: | |
| output_type = output_type or "np" | |
| input_image = HWC3(input_image) | |
| input_image = resize_image(input_image, detect_resolution) | |
| img_gray = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY) | |
| edges = ed.detectEdges(img_gray) | |
| if mode == 'edge': | |
| edge_map = ed.getEdgeImage(edges) | |
| else: | |
| edge_map = ed.getGradientImage(edges) | |
| edge_map = np.expand_dims(edge_map, axis=2) | |
| edge_map = cv2.cvtColor(edge_map, cv2.COLOR_GRAY2BGR).astype(np.uint8) | |
| edge_map = HWC3(edge_map) | |
| img = resize_image(input_image, image_resolution) | |
| H, W, _C = img.shape | |
| edge_map = cv2.resize(edge_map, (W, H), interpolation=cv2.INTER_LINEAR) | |
| if output_type == "pil": | |
| edge_map = Image.fromarray(edge_map) | |
| return edge_map | |