import torch import gradio as gr from PIL import Image import os import zipfile # Load the YOLOv5 model (ensure the path is correct) model = torch.hub.load( 'ultralytics/yolov5', 'custom', path='best-3.pt' ) # Global variables to store images and labels for carousel functionality processed_images = [] label_contents = [] processed_image_paths = [] # Temporary folder for saving processed files TEMP_DIR = "temp_processed" os.makedirs(TEMP_DIR, exist_ok=True) # Function to process all uploaded images def process_images(files): global processed_images, label_contents, processed_image_paths processed_images = [] label_contents = [] processed_image_paths = [] # Clear the temp directory for f in os.listdir(TEMP_DIR): os.remove(os.path.join(TEMP_DIR, f)) for i, file_path in enumerate(files): # Open the image using the file path img = Image.open(file_path) # Run inference results = model(img) # Convert results to Image format for display results.render() if hasattr(results, 'ims'): output_image = Image.fromarray(results.ims[0]) else: output_image = Image.fromarray(results.imgs[0]) # Save the processed image and its labels processed_images.append(output_image) # Generate YOLO-format labels as a string label_content = "" for *box, conf, cls in results.xywh[0]: class_id = int(cls) x_center, y_center, width, height = box label_content += f"{class_id} {x_center} {y_center} {width} {height}\n" label_contents.append(label_content) # Save the image to the temp folder image_path = os.path.join(TEMP_DIR, f"processed_image_{i}.png") output_image.save(image_path) processed_image_paths.append(image_path) # Save the label content to a text file label_filename = f"annotation_{i}.txt" label_path = os.path.join(TEMP_DIR, label_filename) with open(label_path, "w") as label_file: label_file.write(label_content) # Return the first image and its labels if processed_images: return processed_images[0], label_contents[0], 0 # Start with index 0 else: return None, "No images found.", 0 # Function to create and return the path to the ZIP file for download def create_zip(): zip_filename = "processed_images_annotations.zip" zip_path = os.path.join(TEMP_DIR, zip_filename) # Remove existing ZIP file if it exists if os.path.exists(zip_path): os.remove(zip_path) with zipfile.ZipFile(zip_path, 'w') as z: # Add images and labels to the ZIP file for image_path in processed_image_paths: z.write(image_path, os.path.basename(image_path)) # Get index from image filename image_filename = os.path.basename(image_path) base_name, ext = os.path.splitext(image_filename) index = base_name.split('_')[-1] # Construct label filename label_filename = f"annotation_{index}.txt" label_path = os.path.join(TEMP_DIR, label_filename) z.write(label_path, label_filename) return zip_path # Return the file path as a string # Function to navigate through images def next_image(index): global processed_images, label_contents if processed_images: index = (index + 1) % len(processed_images) return processed_images[index], label_contents[index], index else: return None, "No images processed.", index def prev_image(index): global processed_images, label_contents if processed_images: index = (index - 1) % len(processed_images) return processed_images[index], label_contents[index], index else: return None, "No images processed.", index # Gradio interface with gr.Blocks() as interface: # Welcome message and general instructions gr.Markdown("## Welcome to the Construction Image Processing App!") gr.Markdown(""" This app allows you to upload multiple construction-related images, process them using a custom YOLO model, and download the processed images along with their YOLO-format annotations. ### How to Use: 1. **Upload**: Upload multiple images for processing. 2. **View**: Navigate through the processed images using the 'Previous' and 'Next' buttons. 3. **Download**: Download all processed images and annotations as a ZIP file. Follow the instructions below to get started! """) # Step 1: Upload Instructions gr.Markdown("### Step 1: Upload Your Images") gr.Markdown("Upload multiple construction-related image files that you want to process using the model. The supported file types are JPEG and PNG.") # Multiple file input and display area file_input = gr.Files(label="Upload multiple image files", type="filepath") image_display = gr.Image(label="Processed Image") label_display = gr.Textbox(label="YOLO Annotation for the Processed Image") # Step 2: Navigation Instructions gr.Markdown("### Step 2: View Your Processed Images") gr.Markdown("After processing, use the 'Previous Image' and 'Next Image' buttons to navigate through the processed images and view their annotations.") # Buttons for carousel navigation prev_button = gr.Button("Previous Image") next_button = gr.Button("Next Image") # Hidden state to store current index current_index = gr.State(0) # Step 3: Download Instructions gr.Markdown("### Step 3: Download All Processed Images and Annotations") gr.Markdown("Click the 'Prepare and Download All' button to download a ZIP file containing all processed images and YOLO annotations.") # Button to download all processed images and annotations as a ZIP file download_button = gr.Button("Prepare and Download All") download_file = gr.File() # Define functionality when files are uploaded file_input.change( process_images, inputs=file_input, outputs=[image_display, label_display, current_index] ) # Define functionality for next and previous buttons next_button.click( next_image, inputs=current_index, outputs=[image_display, label_display, current_index] ) prev_button.click( prev_image, inputs=current_index, outputs=[image_display, label_display, current_index] ) # Define functionality for the download button to zip the files and allow download def prepare_download(): zip_path = create_zip() return zip_path download_button.click( prepare_download, outputs=download_file ) # Launch the interface interface.launch(share=True)