Upload 2 files

app.py

@@ -0,0 +1,165 @@
+import torch
+import gradio as gr
+from PIL import Image
+import os
+import zipfile
+import io
+
+# Load the YOLOv5 model (ensure the path is correct)
+model = torch.hub.load(
+    'ultralytics/yolov5', 'custom', path='yolov5/runs/train/exp/weights/best.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"
+
+# Ensure the temp directory exists
+if not os.path.exists(TEMP_DIR):
+    os.makedirs(TEMP_DIR)
+
+# 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 a ZIP file for download
+def create_zip():
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, '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)
+
+    zip_buffer.seek(0)  # Go to the start of the buffer
+    # Return the bytes of the ZIP file
+    return zip_buffer.getvalue()
+
+# 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:
+    # 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="Label File Content")
+
+    # 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)
+
+    # Button to prepare the ZIP file for download
+    prepare_download_button = gr.Button("Prepare Download")
+
+    # Download button
+    download_button = gr.DownloadButton(label="Download ZIP", visible=False)
+
+    # 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 prepare download button
+    def prepare_download():
+        data = create_zip()
+        return gr.update(value={"name": "processed_images_annotations.zip", "data": data}, visible=True)
+
+    prepare_download_button.click(
+        prepare_download,
+        outputs=download_button
+    )
+
+# Launch the interface
+interface.launch()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+gradio
+Pillow