Spaces:
Sleeping
Sleeping
| import glob | |
| import os | |
| import cv2 | |
| import gradio as gr | |
| def center_crop_mod64(img): | |
| """Crop image to largest area that's divisible by 64""" | |
| h, w = img.shape[:2] | |
| shortside, longside = min(w, h), max(w, h) | |
| longside_crop = (longside // 64) * 64 | |
| left = (w - longside_crop) // 2 if w > h else 0 | |
| top = (h - longside_crop) // 2 if h > w else 0 | |
| right = left + longside_crop if w > h else shortside | |
| bottom = top + longside_crop if h > w else shortside | |
| return img[top:bottom, left:right] | |
| def load_example_images(): | |
| """Load all images from images/ directory""" | |
| image_dir = os.path.join(os.path.dirname(__file__), "images") | |
| if not os.path.exists(image_dir): | |
| return [] | |
| extensions = ['*.png', '*.jpg', '*.jpeg', '*.webp'] | |
| examples = [] | |
| for ext in extensions: | |
| examples.extend(glob.glob(os.path.join(image_dir, ext))) | |
| examples.extend(glob.glob(os.path.join(image_dir, ext.upper()))) | |
| return sorted(examples) | |
| def apply_edge_detection(image_rgb, algorithm, apply_crop, convert2grayscale, params): | |
| """Apply selected edge detection algorithm with specified parameters""" | |
| if image_rgb is None: | |
| return None | |
| # Apply cropping if requested | |
| if apply_crop: | |
| img_processed = center_crop_mod64(image_rgb) | |
| else: | |
| img_processed = image_rgb.copy() | |
| # Convert to grayscale if requested | |
| if convert2grayscale and len(img_processed.shape) == 3: | |
| img_input = cv2.cvtColor(cv2.cvtColor(img_processed, cv2.COLOR_RGB2GRAY), cv2.COLOR_GRAY2RGB) | |
| else: | |
| img_input = img_processed | |
| ret = img_input.copy() | |
| # set parameters | |
| # Detect edges | |
| ed = cv2.ximgproc.createEdgeDrawing() | |
| ed.setParams(params) | |
| ed.detectEdges(img_input) | |
| if algorithm == "Edges": | |
| ret = ed.getEdgeImage() | |
| return ret | |
| elif algorithm == "Lines": | |
| lines = ed.detectLines() | |
| # draw lines | |
| if lines is not None: | |
| for line in lines: | |
| x1, y1, x2, y2 = map(int, line[0]) | |
| cv2.line(ret, (x1, y1), (x2, y2), (0, 255, 0), 1) | |
| return ret | |
| elif algorithm == "Ellipses": | |
| ellipsess = ed.detectEllipses() | |
| # draw ellipses | |
| if ellipsess is not None: | |
| for ellipses in ellipsess: | |
| ellipse = ellipses[0] | |
| center = (int(ellipse[0]), int(ellipse[1])) | |
| if ellipse[2] == 0: # Ellipse | |
| axes = (int(ellipse[3]), int(ellipse[4])) | |
| angle = ellipse[5] | |
| cv2.ellipse(ret, center, axes, angle, 0, 360, (0, 255, 0), 1) | |
| else: # Circle | |
| radius = int(ellipse[2]) | |
| cv2.circle(ret, center, radius, (0, 255, 0), 1) | |
| return ret | |
| return img_input | |
| # Create Gradio interface | |
| with gr.Blocks(title="Edge Drawing") as app: | |
| examples = load_example_images() | |
| gr.Markdown("# Enhanced Edge Detection using [Edge Drawing](https://github.com/CihanTopal/ED_Lib) with [opencv-contrib-python](https://docs.opencv.org/4.x/d1/d1c/classcv_1_1ximgproc_1_1EdgeDrawing.html)") | |
| # Image row - Input and Output side by side | |
| with gr.Row(): | |
| input_image = gr.Image(value=examples[0] if examples else None , label="Input Image" , type="numpy") | |
| output_image = gr.Image(value=None , label="Detection Result" , type="numpy") | |
| # Controls row - Processing Options on left, Parameters on right | |
| with gr.Row(): | |
| # Processing Options | |
| with gr.Column(): | |
| with gr.Group(): | |
| gr.Markdown("### Processing Options") | |
| algorithm_radio = gr.Radio(label="Detection mode", value="Edges", choices=["Edges", "Lines", "Ellipses"]) | |
| paramterfree_checkbox = gr.Checkbox(label ="Parameter Free Mode" , value=True , info="Auto-determine optimal parameters") | |
| convert2grayscale_checkbox = gr.Checkbox(label ="Convert to Grayscale" , value=False , info="Force conversion to grayscale prior to edge detection (note that Edge Drawing has native support for color images)") | |
| crop_checkbox = gr.Checkbox(label ="Apply center crop mod-64" , value=False , info="Crop to largest area divisible by 64 (used for Stable Diffusion)") | |
| # EdgeDrawing Parameters | |
| with gr.Column(): | |
| with gr.Group(): | |
| gr.Markdown("### Edge Drawing Parameters") | |
| # Parameter controls group (disabled by default) | |
| with gr.Group(visible=False) as param_group: | |
| with gr.Column(): | |
| with gr.Row(): | |
| gradient_operator = gr.Radio(choices=["PREWITT", "SOBEL", "SCHARR", "LSD"], value="PREWITT", label="Gradient Operator", info="indicates the operator used for gradient calculation. Default value is PREWITT") | |
| gradient_threshold = gr.Slider(minimum= 1 , maximum= 100 , value= 20 , step= 1 , label="Gradient Threshold" , scale= 1 , info="Threshold value of gradiential difference between pixels. Used to create gradient image. Default value is 20" ) | |
| with gr.Row(): | |
| anchor_threshold = gr.Slider(minimum= 0 , maximum= 255 , value= 0 , step= 1 , label="Anchor Threshold" , scale= 1 , info="Threshold value used to select anchor points. Default value is 0" ) | |
| min_path_length = gr.Slider(minimum= 1 , maximum= 50 , value= 10 , step= 1 , label="Min Path Length" , scale= 1 , info="Minimum connected pixels length processed to create an edge segment. Default value is 10" ) | |
| with gr.Row(): | |
| min_line_length = gr.Slider(minimum= -1 , maximum= 100 , value= -1 , step= 1 , label="Min Line Length" , scale= 1 , info="Minimum line length to detect. Default value is -1" ) | |
| line_fit_error_threshold = gr.Slider(minimum= 0.1 , maximum= 10.0 , value= 1.0 , step= 0.1 , label="Line Fit Error Threshold" , scale= 1 , info="Default value is 1.0" ) | |
| max_distance_between_two_lines = gr.Slider(minimum= 1.0 , maximum= 20.0 , value= 6.0 , step= 0.1 , label="Max Distance Between Two Lines" , scale= 1 , info="Default value is 6.0" ) | |
| with gr.Row(): | |
| max_error_threshold = gr.Slider(minimum= 0.1 , maximum= 5.0 , value= 1.3 , step= 0.1 , label="Max Error Threshold" , scale= 1 , info="Default value is 1.3" ) | |
| scan_interval = gr.Slider(minimum= 1 , maximum= 10 , value= 1 , step= 1 , label="Scan Interval" , scale= 1 , info="Default value is 1" ) | |
| with gr.Row(): | |
| sigma = gr.Slider(minimum= 0.1 , maximum= 5.0 , value= 1.0 , step= 0.1 , label="Sigma" , scale= 1 , info="Sigma value for internal GaussianBlur() function. Default value is 1.0" ) | |
| nfa_validation = gr.Checkbox(label="NFA Validation" , value=True , info="Indicates if NFA (Number of False Alarms) algorithm will be used for line and ellipse validation. Default value is true") | |
| sum_flag = gr.Checkbox(label="Sum Flag" , value=True , info="Default value is true") | |
| # Update function | |
| def update_output( | |
| image_rgb, algorithm, apply_crop, conv2grayscale, pf_mode, | |
| gradient_operator, anchor_threshold, gradient_threshold, min_path_length, min_line_length, line_fit_error_threshold, max_distance_between_two_lines, max_error_threshold, nfa_validation, scan_interval, sigma, sum_flag | |
| ): | |
| operator_map = { | |
| "PREWITT" : cv2.ximgproc.EdgeDrawing_PREWITT, | |
| "SOBEL" : cv2.ximgproc.EdgeDrawing_SOBEL, | |
| "SCHARR" : cv2.ximgproc.EdgeDrawing_SCHARR, | |
| "LSD" : cv2.ximgproc.EdgeDrawing_LSD, | |
| } | |
| params = cv2.ximgproc.EdgeDrawing.Params() | |
| params.PFmode = pf_mode | |
| params.EdgeDetectionOperator = operator_map.get(gradient_operator, cv2.ximgproc.EdgeDrawing_PREWITT) | |
| if not pf_mode: | |
| params.AnchorThresholdValue = anchor_threshold | |
| params.GradientThresholdValue = gradient_threshold | |
| params.MinPathLength = min_path_length | |
| params.MinLineLength = min_line_length | |
| params.LineFitErrorThreshold = line_fit_error_threshold | |
| params.MaxDistanceBetweenTwoLines = max_distance_between_two_lines | |
| params.MaxErrorThreshold = max_error_threshold | |
| params.NFAValidation = nfa_validation | |
| params.ScanInterval = scan_interval | |
| params.Sigma = sigma | |
| params.SumFlag = sum_flag | |
| ret = apply_edge_detection(image_rgb, algorithm, apply_crop, conv2grayscale, params) | |
| return ret | |
| # Function to toggle parameter group visibility | |
| def toggle_params(pf_mode): | |
| return gr.Group(visible=not pf_mode) | |
| # Set up event handlers | |
| inputs = [ | |
| input_image, | |
| algorithm_radio, | |
| crop_checkbox, | |
| convert2grayscale_checkbox, | |
| paramterfree_checkbox, | |
| gradient_operator, | |
| anchor_threshold, | |
| gradient_threshold, | |
| min_path_length, | |
| min_line_length, | |
| line_fit_error_threshold, | |
| max_distance_between_two_lines, | |
| max_error_threshold, | |
| nfa_validation, | |
| scan_interval, | |
| sigma, | |
| sum_flag, | |
| ] | |
| # Update output when any input changes | |
| for inp in inputs: | |
| inp.change(fn=update_output, inputs=inputs, outputs=output_image) | |
| # Toggle parameter group when parameter free mode changes | |
| paramterfree_checkbox.change(fn=toggle_params, inputs=paramterfree_checkbox, outputs=param_group) | |
| # Apply edge detection immediately on startup | |
| app.load(fn=update_output, inputs=inputs, outputs=output_image) | |
| # Examples | |
| if examples: | |
| gr.Examples(examples=examples, inputs=input_image, label="Example Images" ) | |
| if __name__ == "__main__": | |
| app.launch() | |