import gradio as gr import os import shutil from typing import Optional from .base import BaseUI from config.path_config import ( DEMOS_DIR as DEMO_IMAGES_DIR, VISUALIZATIONS_DIR, SUPPORTED_IMAGE_EXTENSIONS, SUPPORTED_CONFIG_EXTENSIONS, ensure_directories_exist, get_working_image_path, get_working_config_path ) from endpoints.converter import run_and_save_pipeline class ImageProcessor(BaseUI): """Image Processor UI component""" def __init__(self): super().__init__() self.demo_images = self._load_demo_images() # UI components self.demo_gallery = None self.image_upload = None self.yaml_upload = None self.process_btn = None self.output_gallery = None self.output_message = None # State variables self.current_image = None self.working_image = None self.current_yaml = None self.working_yaml = None self.status_text = None def _load_demo_images(self): """Load demo images from the demos directory""" try: demo_images = [os.path.join(DEMO_IMAGES_DIR, f) for f in os.listdir(DEMO_IMAGES_DIR) if f.endswith(SUPPORTED_IMAGE_EXTENSIONS)] if not demo_images: print("No demo images found in the demos directory. Please add some images to", DEMO_IMAGES_DIR) return [] return demo_images except Exception as e: print(f"Error loading demo images: {e}") return [] def create_interface(self) -> gr.TabItem: """Create the Image Processor interface""" with gr.TabItem("Image Processor") as tab: gr.Markdown("## Sketch to PNML Converter") # State variables for the app self.current_image = self.create_state_var("current_image", None) self.working_image = self.create_state_var("working_image", None) self.current_yaml = self.create_state_var("current_yaml", None) self.working_yaml = self.create_state_var("working_yaml", None) self.status_text = self.create_state_var("status_text", "") # Demo section at the top with gr.Row(): with gr.Column(): gr.Markdown("### Demo Images") if self.demo_images: self.demo_gallery = gr.Gallery( value=self.demo_images, columns=5, object_fit="contain", height="150px", show_label=False ) else: gr.Markdown(f"*No demo images available. Place image files in the '{DEMO_IMAGES_DIR}' directory.*") # Input and output sections below with gr.Row(): # Left column for inputs with gr.Column(): # Image upload section gr.Markdown("### Or Upload Your Own Image") self.image_upload = self.create_file_upload( "Upload Image", ["image"] ) # YAML config upload section gr.Markdown("### Upload YAML Configuration") self.yaml_upload = self.create_file_upload( "Upload YAML Config", list(SUPPORTED_CONFIG_EXTENSIONS) ) # Process button self.process_btn = self.create_button("Process Files", variant="primary", size="lg") # Right column for outputs with gr.Column(): gr.Markdown("### Output Visualizations") self.output_gallery = gr.Gallery( label="Processing Results", columns=1, height=400, show_label=True ) self.output_message = gr.Markdown("") self.setup_event_handlers() return tab def setup_event_handlers(self): """Setup event handlers for the Image Processor""" # Event handlers if self.demo_images and self.demo_gallery: self.demo_gallery.select( self._select_demo_image, inputs=[self.current_yaml, self.working_yaml], outputs=[self.current_image, self.working_image, self.current_yaml, self.working_yaml, self.status_text] ) self.image_upload.change( self._upload_image, inputs=[self.image_upload, self.current_yaml, self.working_yaml], outputs=[self.current_image, self.working_image, self.current_yaml, self.working_yaml, self.status_text] ) self.yaml_upload.change( self._upload_yaml, inputs=[self.yaml_upload, self.current_image, self.working_image], outputs=[self.current_image, self.working_image, self.current_yaml, self.working_yaml, self.status_text] ) self.process_btn.click( self._process_files, inputs=[self.working_image, self.working_yaml], outputs=[self.output_gallery, self.output_message] ) # Update status message when it changes self.status_text.change(lambda x: x, inputs=[self.status_text], outputs=[self.output_message]) def _save_as_working_image(self, image_path): """Save the selected image as the working image""" # Ensure the extension is preserved when copying _, ext = os.path.splitext(image_path) working_image_path = get_working_image_path(ext) shutil.copy(image_path, working_image_path) return working_image_path def _save_as_working_yaml(self, yaml_path): """Save the uploaded YAML as the working configuration""" _, ext = os.path.splitext(yaml_path) working_yaml_path = get_working_config_path(ext) shutil.copy(yaml_path, working_yaml_path) return working_yaml_path def _clear_visualizations(self): """Clear the visualizations directory""" # Create the directory if it doesn't exist if not os.path.exists(VISUALIZATIONS_DIR): os.makedirs(VISUALIZATIONS_DIR, exist_ok=True) # Remove existing files for file in os.listdir(VISUALIZATIONS_DIR): file_path = os.path.join(VISUALIZATIONS_DIR, file) if os.path.isfile(file_path): os.remove(file_path) def _process_files(self, image_path, yaml_path): """Process the uploaded files and run the pipeline""" if image_path is None or yaml_path is None: return None, "Both image and YAML configuration file are required." # Clear previous visualization results self._clear_visualizations() # Run the pipeline try: run_and_save_pipeline(config_path=yaml_path, image_path=image_path) # Get visualization results vis_files = [os.path.join(VISUALIZATIONS_DIR, f) for f in os.listdir(VISUALIZATIONS_DIR) if f.endswith('.png')] if not vis_files: return None, "No visualization results were generated." return vis_files, "Processing completed successfully!" except Exception as e: return None, self.handle_error(e, "during processing") def _select_demo_image(self, evt: gr.SelectData, current_yaml=None, working_yaml=None): """Handle selection of a demo image""" image_path = self.demo_images[evt.index] if image_path: working_image_path = self._save_as_working_image(image_path) status = "Demo image selected." if working_yaml: status += " Ready to process!" else: status += " Please upload a YAML configuration file." return image_path, working_image_path, current_yaml, working_yaml, status return None, None, current_yaml, working_yaml, "" def _upload_image(self, file_obj, current_yaml=None, working_yaml=None): """Handle upload of a custom image""" if file_obj is not None: # Save the uploaded file to working image _, ext = os.path.splitext(file_obj.name) working_image_path = get_working_image_path(ext) shutil.copy(file_obj.name, working_image_path) status = "Custom image uploaded." if working_yaml: status += " Ready to process!" else: status += " Please upload a YAML configuration file." return file_obj.name, working_image_path, current_yaml, working_yaml, status return None, None, current_yaml, working_yaml, "" def _upload_yaml(self, file_obj, current_image, working_image): """Handle upload of a YAML file""" if file_obj is not None: # Save the uploaded YAML as the working config yaml_path = file_obj.name working_yaml_path = self._save_as_working_yaml(yaml_path) if working_image is not None: return current_image, working_image, yaml_path, working_yaml_path, "Both files uploaded. Ready to process!" else: return None, None, yaml_path, working_yaml_path, "YAML file uploaded. Please select or upload an image." return current_image, working_image, None, None, ""