Create app.py

app.py

@@ -0,0 +1,188 @@
+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)