Spaces:
Sleeping
Sleeping
Upload 3 files
Browse files- app.py +163 -0
- best-3.pt +3 -0
- requirements.txt +6 -0
app.py
ADDED
@@ -0,0 +1,163 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import gradio as gr
|
3 |
+
from PIL import Image
|
4 |
+
import os
|
5 |
+
import zipfile
|
6 |
+
|
7 |
+
# Load the YOLOv5 model (ensure the path is correct)
|
8 |
+
model = torch.hub.load(
|
9 |
+
'ultralytics/yolov5', 'custom', path='best-3.pt'
|
10 |
+
)
|
11 |
+
|
12 |
+
# Global variables to store images and labels for carousel functionality
|
13 |
+
processed_images = []
|
14 |
+
label_contents = []
|
15 |
+
processed_image_paths = []
|
16 |
+
|
17 |
+
# Temporary folder for saving processed files
|
18 |
+
TEMP_DIR = "temp_processed"
|
19 |
+
os.makedirs(TEMP_DIR, exist_ok=True)
|
20 |
+
|
21 |
+
# Function to process all uploaded images
|
22 |
+
def process_images(files):
|
23 |
+
global processed_images, label_contents, processed_image_paths
|
24 |
+
processed_images = []
|
25 |
+
label_contents = []
|
26 |
+
processed_image_paths = []
|
27 |
+
|
28 |
+
# Clear the temp directory
|
29 |
+
for f in os.listdir(TEMP_DIR):
|
30 |
+
os.remove(os.path.join(TEMP_DIR, f))
|
31 |
+
|
32 |
+
for i, file_path in enumerate(files):
|
33 |
+
# Open the image using the file path
|
34 |
+
img = Image.open(file_path)
|
35 |
+
|
36 |
+
# Run inference
|
37 |
+
results = model(img)
|
38 |
+
|
39 |
+
# Convert results to Image format for display
|
40 |
+
results.render()
|
41 |
+
if hasattr(results, 'ims'):
|
42 |
+
output_image = Image.fromarray(results.ims[0])
|
43 |
+
else:
|
44 |
+
output_image = Image.fromarray(results.imgs[0])
|
45 |
+
|
46 |
+
# Save the processed image and its labels
|
47 |
+
processed_images.append(output_image)
|
48 |
+
|
49 |
+
# Generate YOLO-format labels as a string
|
50 |
+
label_content = ""
|
51 |
+
for *box, conf, cls in results.xywh[0]:
|
52 |
+
class_id = int(cls)
|
53 |
+
x_center, y_center, width, height = box
|
54 |
+
label_content += f"{class_id} {x_center} {y_center} {width} {height}\n"
|
55 |
+
label_contents.append(label_content)
|
56 |
+
|
57 |
+
# Save the image to the temp folder
|
58 |
+
image_path = os.path.join(TEMP_DIR, f"processed_image_{i}.png")
|
59 |
+
output_image.save(image_path)
|
60 |
+
processed_image_paths.append(image_path)
|
61 |
+
|
62 |
+
# Save the label content to a text file
|
63 |
+
label_filename = f"annotation_{i}.txt"
|
64 |
+
label_path = os.path.join(TEMP_DIR, label_filename)
|
65 |
+
with open(label_path, "w") as label_file:
|
66 |
+
label_file.write(label_content)
|
67 |
+
|
68 |
+
# Return the first image and its labels
|
69 |
+
if processed_images:
|
70 |
+
return processed_images[0], label_contents[0], 0 # Start with index 0
|
71 |
+
else:
|
72 |
+
return None, "No images found.", 0
|
73 |
+
|
74 |
+
# Function to create and return the path to the ZIP file for download
|
75 |
+
def create_zip():
|
76 |
+
zip_filename = "processed_images_annotations.zip"
|
77 |
+
zip_path = os.path.join(TEMP_DIR, zip_filename)
|
78 |
+
|
79 |
+
# Remove existing ZIP file if it exists
|
80 |
+
if os.path.exists(zip_path):
|
81 |
+
os.remove(zip_path)
|
82 |
+
|
83 |
+
with zipfile.ZipFile(zip_path, 'w') as z:
|
84 |
+
# Add images and labels to the ZIP file
|
85 |
+
for image_path in processed_image_paths:
|
86 |
+
z.write(image_path, os.path.basename(image_path))
|
87 |
+
# Get index from image filename
|
88 |
+
image_filename = os.path.basename(image_path)
|
89 |
+
base_name, ext = os.path.splitext(image_filename)
|
90 |
+
index = base_name.split('_')[-1]
|
91 |
+
# Construct label filename
|
92 |
+
label_filename = f"annotation_{index}.txt"
|
93 |
+
label_path = os.path.join(TEMP_DIR, label_filename)
|
94 |
+
z.write(label_path, label_filename)
|
95 |
+
|
96 |
+
return zip_path # Return the file path as a string
|
97 |
+
|
98 |
+
# Function to navigate through images
|
99 |
+
def next_image(index):
|
100 |
+
global processed_images, label_contents
|
101 |
+
if processed_images:
|
102 |
+
index = (index + 1) % len(processed_images)
|
103 |
+
return processed_images[index], label_contents[index], index
|
104 |
+
else:
|
105 |
+
return None, "No images processed.", index
|
106 |
+
|
107 |
+
def prev_image(index):
|
108 |
+
global processed_images, label_contents
|
109 |
+
if processed_images:
|
110 |
+
index = (index - 1) % len(processed_images)
|
111 |
+
return processed_images[index], label_contents[index], index
|
112 |
+
else:
|
113 |
+
return None, "No images processed.", index
|
114 |
+
|
115 |
+
# Gradio interface
|
116 |
+
with gr.Blocks() as interface:
|
117 |
+
# Multiple file input and display area
|
118 |
+
file_input = gr.Files(label="Upload multiple image files", type="filepath")
|
119 |
+
image_display = gr.Image(label="Processed Image")
|
120 |
+
label_display = gr.Textbox(label="Label File Content")
|
121 |
+
|
122 |
+
# Buttons for carousel navigation
|
123 |
+
prev_button = gr.Button("Previous Image")
|
124 |
+
next_button = gr.Button("Next Image")
|
125 |
+
|
126 |
+
# Hidden state to store current index
|
127 |
+
current_index = gr.State(0)
|
128 |
+
|
129 |
+
# Button to download all processed images and annotations as a ZIP file
|
130 |
+
download_button = gr.Button("Prepare and Download All")
|
131 |
+
download_file = gr.File()
|
132 |
+
|
133 |
+
# Define functionality when files are uploaded
|
134 |
+
file_input.change(
|
135 |
+
process_images,
|
136 |
+
inputs=file_input,
|
137 |
+
outputs=[image_display, label_display, current_index]
|
138 |
+
)
|
139 |
+
|
140 |
+
# Define functionality for next and previous buttons
|
141 |
+
next_button.click(
|
142 |
+
next_image,
|
143 |
+
inputs=current_index,
|
144 |
+
outputs=[image_display, label_display, current_index]
|
145 |
+
)
|
146 |
+
prev_button.click(
|
147 |
+
prev_image,
|
148 |
+
inputs=current_index,
|
149 |
+
outputs=[image_display, label_display, current_index]
|
150 |
+
)
|
151 |
+
|
152 |
+
# Define functionality for the download button to zip the files and allow download
|
153 |
+
def prepare_download():
|
154 |
+
zip_path = create_zip()
|
155 |
+
return zip_path
|
156 |
+
|
157 |
+
download_button.click(
|
158 |
+
prepare_download,
|
159 |
+
outputs=download_file
|
160 |
+
)
|
161 |
+
|
162 |
+
# Launch the interface
|
163 |
+
interface.launch(share=True)
|
best-3.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:3db11edeefd171be8ce807fbbba06a253c267e39100d06bc32da890c24c65f25
|
3 |
+
size 6262179
|
requirements.txt
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
gradio
|
3 |
+
Pillow
|
4 |
+
opencv-python
|
5 |
+
numpy==1.23.4
|
6 |
+
scipy==1.9.3
|