Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| from PIL import Image | |
| from streamlit_drawable_canvas import st_canvas | |
| import os | |
| def image_annoter(): | |
| st.title("Image Annoter for YOLO") | |
| st.write("Enter image files folder Location.") | |
| # Upload video | |
| file_location = st.text_input("Enter File folder Location",None) | |
| if file_location is not None: | |
| # Folder containing your images | |
| label_folder = file_location | |
| # Get a list of all images in the folder | |
| image_files = [f for f in os.listdir(label_folder) if f.endswith(('png', 'jpg', 'jpeg'))] | |
| # Initialize session state to keep track of the current image index | |
| if 'current_index' not in st.session_state: | |
| st.session_state.current_index = 0 | |
| # Set the downscaling factor | |
| downscale_factor = 0.5 # Adjust the downscale factor as needed | |
| # Navigation buttons | |
| col1, col2, col3 = st.columns([1, 2, 1]) | |
| if col1.button("Previous"): | |
| if st.session_state.current_index > 0: | |
| st.session_state.current_index -= 1 | |
| if col3.button("Next"): | |
| if st.session_state.current_index < len(image_files) - 1: | |
| st.session_state.current_index += 1 | |
| # Display the current image | |
| current_image_file = image_files[st.session_state.current_index] | |
| st.write(f"Annotating: {current_image_file}") | |
| image_path = os.path.join(label_folder, current_image_file) | |
| image = Image.open(image_path) | |
| # Downscale the image for the canvas | |
| scaled_width = int(image.width * downscale_factor) | |
| scaled_height = int(image.height * downscale_factor) | |
| scaled_image = image.resize((scaled_width, scaled_height)) | |
| # Display the image on the canvas | |
| canvas_result = st_canvas( | |
| fill_color="rgba(255, 0, 0, 0.3)", # Fill color for the bounding box | |
| stroke_width=3, | |
| stroke_color="#ff0000", | |
| background_image=scaled_image, | |
| height=scaled_height, | |
| width=scaled_width, | |
| drawing_mode="rect", | |
| key=current_image_file | |
| ) | |
| # Save annotations | |
| if st.button("Save Annotation"): | |
| if canvas_result.json_data is not None: | |
| # Extract coordinates of the bounding box | |
| for obj in canvas_result.json_data["objects"]: | |
| if obj["type"] == "rect": | |
| left = obj["left"] / downscale_factor | |
| top = obj["top"] / downscale_factor | |
| width = obj["width"] / downscale_factor | |
| height = obj["height"] / downscale_factor | |
| x_center = left + width / 2 | |
| y_center = top + height / 2 | |
| # Normalize the coordinates (YOLO format) | |
| x_center /= image.width | |
| y_center /= image.height | |
| width /= image.width | |
| height /= image.height | |
| # Save the annotation to a .txt file | |
| annotation_path = os.path.join(label_folder, current_image_file.replace('.jpg', '.txt').replace('.png', '.txt')) | |
| with open(annotation_path, 'w') as f: | |
| f.write(f"0 {x_center} {y_center} {width} {height}\n") | |
| st.success(f"Annotation saved for {current_image_file}") | |