|
import torch |
|
import gradio as gr |
|
from PIL import Image |
|
import os |
|
import zipfile |
|
|
|
|
|
model = torch.hub.load( |
|
'ultralytics/yolov5', 'custom', path='best-3.pt' |
|
) |
|
|
|
|
|
processed_images = [] |
|
label_contents = [] |
|
processed_image_paths = [] |
|
|
|
|
|
TEMP_DIR = "temp_processed" |
|
os.makedirs(TEMP_DIR, exist_ok=True) |
|
|
|
|
|
def process_images(files): |
|
global processed_images, label_contents, processed_image_paths |
|
processed_images = [] |
|
label_contents = [] |
|
processed_image_paths = [] |
|
|
|
|
|
for f in os.listdir(TEMP_DIR): |
|
os.remove(os.path.join(TEMP_DIR, f)) |
|
|
|
for i, file_path in enumerate(files): |
|
|
|
img = Image.open(file_path) |
|
|
|
|
|
results = model(img) |
|
|
|
|
|
results.render() |
|
if hasattr(results, 'ims'): |
|
output_image = Image.fromarray(results.ims[0]) |
|
else: |
|
output_image = Image.fromarray(results.imgs[0]) |
|
|
|
|
|
processed_images.append(output_image) |
|
|
|
|
|
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) |
|
|
|
|
|
image_path = os.path.join(TEMP_DIR, f"processed_image_{i}.png") |
|
output_image.save(image_path) |
|
processed_image_paths.append(image_path) |
|
|
|
|
|
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) |
|
|
|
|
|
if processed_images: |
|
return processed_images[0], label_contents[0], 0 |
|
else: |
|
return None, "No images found.", 0 |
|
|
|
|
|
def create_zip(): |
|
zip_filename = "processed_images_annotations.zip" |
|
zip_path = os.path.join(TEMP_DIR, zip_filename) |
|
|
|
|
|
if os.path.exists(zip_path): |
|
os.remove(zip_path) |
|
|
|
with zipfile.ZipFile(zip_path, 'w') as z: |
|
|
|
for image_path in processed_image_paths: |
|
z.write(image_path, os.path.basename(image_path)) |
|
|
|
image_filename = os.path.basename(image_path) |
|
base_name, ext = os.path.splitext(image_filename) |
|
index = base_name.split('_')[-1] |
|
|
|
label_filename = f"annotation_{index}.txt" |
|
label_path = os.path.join(TEMP_DIR, label_filename) |
|
z.write(label_path, label_filename) |
|
|
|
return zip_path |
|
|
|
|
|
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 |
|
|
|
|
|
with gr.Blocks() as interface: |
|
|
|
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! |
|
""") |
|
|
|
|
|
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.") |
|
|
|
|
|
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") |
|
|
|
|
|
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.") |
|
|
|
|
|
prev_button = gr.Button("Previous Image") |
|
next_button = gr.Button("Next Image") |
|
|
|
|
|
current_index = gr.State(0) |
|
|
|
|
|
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.") |
|
|
|
|
|
download_button = gr.Button("Prepare and Download All") |
|
download_file = gr.File() |
|
|
|
|
|
file_input.change( |
|
process_images, |
|
inputs=file_input, |
|
outputs=[image_display, label_display, current_index] |
|
) |
|
|
|
|
|
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] |
|
) |
|
|
|
|
|
def prepare_download(): |
|
zip_path = create_zip() |
|
return zip_path |
|
|
|
download_button.click( |
|
prepare_download, |
|
outputs=download_file |
|
) |
|
|
|
|
|
interface.launch(share=True) |
|
|