Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -3,6 +3,9 @@ from gradio_bbox_annotator import BBoxAnnotator
|
|
3 |
import json
|
4 |
import os
|
5 |
from pathlib import Path
|
|
|
|
|
|
|
6 |
|
7 |
# Define categories and their limits
|
8 |
CATEGORY_LIMITS = {
|
@@ -10,10 +13,35 @@ CATEGORY_LIMITS = {
|
|
10 |
"text": 2 # Maximum 2 text annotations per image
|
11 |
}
|
12 |
CATEGORIES = list(CATEGORY_LIMITS.keys())
|
|
|
13 |
|
14 |
class AnnotationManager:
|
15 |
def __init__(self):
|
16 |
self.annotations = {}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
17 |
|
18 |
def validate_annotations(self, bbox_data):
|
19 |
"""Validate the annotation data and return (is_valid, error_message)"""
|
@@ -60,7 +88,11 @@ class AnnotationManager:
|
|
60 |
return self.get_json_annotations(), f"❌ Error: {error_msg}"
|
61 |
|
62 |
image_path, annotations = bbox_data
|
|
|
63 |
filename = os.path.basename(image_path)
|
|
|
|
|
|
|
64 |
formatted_annotations = []
|
65 |
for ann in annotations:
|
66 |
y1, y2, x1, x2, label = ann
|
@@ -90,11 +122,11 @@ def create_interface():
|
|
90 |
annotation_mgr = AnnotationManager()
|
91 |
|
92 |
with gr.Blocks() as demo:
|
93 |
-
gr.Markdown("""
|
94 |
# Advertisement and Text Annotation Tool
|
95 |
|
96 |
**Instructions:**
|
97 |
-
1. Upload an image
|
98 |
2. Draw bounding boxes and select the appropriate label
|
99 |
3. Click 'Save Annotations' to add to the collection
|
100 |
4. Repeat for all images
|
@@ -127,10 +159,26 @@ def create_interface():
|
|
127 |
save_btn = gr.Button("Save Current Image Annotations", variant="primary")
|
128 |
clear_btn = gr.Button("Clear All Annotations", variant="secondary")
|
129 |
|
130 |
-
# Add status message
|
131 |
status_msg = gr.Markdown(label="Status")
|
132 |
|
133 |
# Event handlers
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
134 |
save_btn.click(
|
135 |
fn=annotation_mgr.add_annotation,
|
136 |
inputs=[bbox_input],
|
|
|
3 |
import json
|
4 |
import os
|
5 |
from pathlib import Path
|
6 |
+
from PIL import Image
|
7 |
+
from io import BytesIO
|
8 |
+
import tempfile
|
9 |
|
10 |
# Define categories and their limits
|
11 |
CATEGORY_LIMITS = {
|
|
|
13 |
"text": 2 # Maximum 2 text annotations per image
|
14 |
}
|
15 |
CATEGORIES = list(CATEGORY_LIMITS.keys())
|
16 |
+
MAX_SIZE = [1024, 1024] # Maximum width and height for resized images
|
17 |
|
18 |
class AnnotationManager:
|
19 |
def __init__(self):
|
20 |
self.annotations = {}
|
21 |
+
self.temp_dir = tempfile.mkdtemp() # Create temporary directory for resized images
|
22 |
+
|
23 |
+
def resize_image(self, image_path):
|
24 |
+
"""Resize image to maximum dimensions while maintaining aspect ratio"""
|
25 |
+
try:
|
26 |
+
# Read and resize image
|
27 |
+
with open(image_path, "rb") as f:
|
28 |
+
img = Image.open(BytesIO(f.read()))
|
29 |
+
img.thumbnail(MAX_SIZE, Image.Resampling.LANCZOS)
|
30 |
+
|
31 |
+
# Save resized image to temporary file
|
32 |
+
filename = os.path.basename(image_path)
|
33 |
+
temp_path = os.path.join(self.temp_dir, f"resized_{filename}")
|
34 |
+
img.save(temp_path)
|
35 |
+
|
36 |
+
return temp_path
|
37 |
+
except Exception as e:
|
38 |
+
raise ValueError(f"Error processing image: {str(e)}")
|
39 |
+
|
40 |
+
def process_image_upload(self, image_path):
|
41 |
+
"""Process uploaded image and return path to resized version"""
|
42 |
+
if not image_path:
|
43 |
+
return None
|
44 |
+
return self.resize_image(image_path)
|
45 |
|
46 |
def validate_annotations(self, bbox_data):
|
47 |
"""Validate the annotation data and return (is_valid, error_message)"""
|
|
|
88 |
return self.get_json_annotations(), f"❌ Error: {error_msg}"
|
89 |
|
90 |
image_path, annotations = bbox_data
|
91 |
+
# Use original filename (remove 'resized_' prefix)
|
92 |
filename = os.path.basename(image_path)
|
93 |
+
if filename.startswith("resized_"):
|
94 |
+
filename = filename[8:]
|
95 |
+
|
96 |
formatted_annotations = []
|
97 |
for ann in annotations:
|
98 |
y1, y2, x1, x2, label = ann
|
|
|
122 |
annotation_mgr = AnnotationManager()
|
123 |
|
124 |
with gr.Blocks() as demo:
|
125 |
+
gr.Markdown(f"""
|
126 |
# Advertisement and Text Annotation Tool
|
127 |
|
128 |
**Instructions:**
|
129 |
+
1. Upload an image (will be automatically resized to max {MAX_SIZE[0]}x{MAX_SIZE[1]})
|
130 |
2. Draw bounding boxes and select the appropriate label
|
131 |
3. Click 'Save Annotations' to add to the collection
|
132 |
4. Repeat for all images
|
|
|
159 |
save_btn = gr.Button("Save Current Image Annotations", variant="primary")
|
160 |
clear_btn = gr.Button("Clear All Annotations", variant="secondary")
|
161 |
|
162 |
+
# Add status message
|
163 |
status_msg = gr.Markdown(label="Status")
|
164 |
|
165 |
# Event handlers
|
166 |
+
def update_image(image_path):
|
167 |
+
if not image_path:
|
168 |
+
return None
|
169 |
+
try:
|
170 |
+
resized_path = annotation_mgr.process_image_upload(image_path)
|
171 |
+
return resized_path
|
172 |
+
except Exception as e:
|
173 |
+
return None
|
174 |
+
|
175 |
+
# Handle image upload and resizing
|
176 |
+
bbox_input.upload(
|
177 |
+
fn=update_image,
|
178 |
+
inputs=[bbox_input],
|
179 |
+
outputs=[bbox_input]
|
180 |
+
)
|
181 |
+
|
182 |
save_btn.click(
|
183 |
fn=annotation_mgr.add_annotation,
|
184 |
inputs=[bbox_input],
|