Image_Augmentation / Introduction.py
Pasham123's picture
Update Introduction.py
8af8edf verified
import streamlit as st
from PIL import Image
import random
import os
import zipfile
import io
# Directory to save images
SAVE_DIR = r"C:\Users\syam0\Downloads\Image App"
# Ensure the directory exists
if not os.path.exists(SAVE_DIR):
os.makedirs(SAVE_DIR)
# Augmentation functions with user inputs
def scale_image(image, count, scale_factor=None):
if scale_factor is None:
scale_factor = random.uniform(0.8, 1.2)
return [image.resize((int(image.width * scale_factor), int(image.height * scale_factor))) for _ in range(count)]
def translate_image(image, count, axis, translation_range=(-40, 40)):
translations = []
for _ in range(count):
if axis == "x-axis":
translation_x = random.randint(*translation_range)
translation_y = 0
elif axis == "y-axis":
translation_x = 0
translation_y = random.randint(*translation_range)
else:
translation_x = random.randint(*translation_range)
translation_y = random.randint(*translation_range)
translated_image = image.transform(image.size, Image.AFFINE, (1, 0, translation_x, 0, 1, translation_y))
translations.append(translated_image)
return translations
def crop_image(image, count, crop_margin=None):
if crop_margin is None:
crop_margin = random.randint(5, 20)
return [image.crop((crop_margin, crop_margin, image.width - crop_margin, image.height - crop_margin)) for _ in range(count)]
def rotate_image(image, count, angle=None):
if angle is None:
angle = random.randint(0, 360)
return [image.rotate(angle) for _ in range(count)]
def shear_image(image, count, shear_factor=None):
if shear_factor is None:
shear_factor = random.uniform(-0.5, 0.5)
return [
image.transform(image.size, Image.AFFINE, (1, shear_factor, 0, shear_factor, 1, 0))
for _ in range(count)
]
# Function to save images to the specified directory
def save_images(images, prefix):
saved_files = []
for i, img in enumerate(images):
filename = os.path.join(SAVE_DIR, f"{prefix}_image_{i+1}.png")
img.save(filename, format="PNG")
saved_files.append(filename)
return saved_files
# Function to zip the images
def create_zip(saved_files):
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w") as zipf:
for file in saved_files:
zipf.write(file, os.path.basename(file))
zip_buffer.seek(0)
return zip_buffer
# Custom CSS for background and styling
st.markdown(
"""
<style>
body {
background: linear-gradient(#ffaa80, #ff5500, #b33c00); /* Gradient background from yellow to red */
font-family: 'Arial', sans-serif;
color: #333;
}
.stApp {
background: linear-gradient(#66d9ff, #d966ff, #66ff66); /* Gradient background for the entire app */
}
.stFileUploader {
display: flex;
justify-content: center;
align-items: center;
padding: 30px;
border: 2px dashed #4CAF50; /* Green dashed border for a fresh look */
background-color: rgba(76, 175, 80, 0.1); /* Light green background on hover */
border-radius: 8px;
cursor: pointer;
}
.stFileUploader:hover {
background-color: rgba(76, 175, 80, 0.2); /* Slightly darker green on hover */
}
</style>
""",
unsafe_allow_html=True
)
# Streamlit app
st.title("Image Augmentation App with Gradient Background")
# Image uploader with custom styling
uploaded_image = st.file_uploader("Drag and Drop Your Image", type=["jpg", "jpeg", "png"], label_visibility="collapsed")
if uploaded_image:
image = Image.open(uploaded_image)
st.image(image, caption="Uploaded Image", use_container_width=True) # Fixed the deprecation warning by using `use_container_width=True`
# User inputs for augmentation
st.write("### Generate Augmented Images")
num_augments = st.number_input("How many augmented images would you like to generate?", min_value=1, max_value=20, value=5)
# Additional inputs for specific techniques
rotate_angle = st.number_input("Enter the angle for rotation (0-360):", min_value=0, max_value=360, value=random.randint(0, 360))
scale_factor = st.number_input("Enter the scaling factor (0.8-1.2):", min_value=0.8, max_value=1.2, value=random.uniform(0.8, 1.2))
crop_margin = st.number_input("Enter the crop margin (5-20):", min_value=5, max_value=20, value=random.randint(5, 20))
shear_factor = st.number_input("Enter the shear factor (-0.5 to 0.5):", min_value=-0.5, max_value=0.5, value=random.uniform(-0.5, 0.5))
translate_axis = st.selectbox("Select axis for translation:", ["x-axis", "y-axis", "both axes"])
# Generate and save all augmentations
if st.button("Generate All Techniques"):
all_augmentations = scale_image(image, num_augments, scale_factor) + \
translate_image(image, num_augments, translate_axis) + \
crop_image(image, num_augments, crop_margin) + \
rotate_image(image, num_augments, rotate_angle) + \
shear_image(image, num_augments, shear_factor)
# Save images to the specified folder
st.write("Generating and Saving Images...")
saved_files = save_images(all_augmentations, "combined")
# Display generated images
st.write("Generated Augmented Images:")
for img in all_augmentations:
st.image(img, use_container_width=True)
# Create a zip file
zip_buffer = create_zip(saved_files)
# Provide a download button for the zip file
st.download_button(
label="Download Augmented Images (ZIP)",
data=zip_buffer,
file_name="augmented_images.zip",
mime="application/zip"
)
# Generate and save technique-specific augmentations
st.write("### Select Specific Technique")
technique = st.selectbox("Choose an augmentation technique:",
["Scale", "Translate", "Crop", "Rotate", "Shear"])
technique_count = st.number_input(f"How many images do you want for {technique}?", min_value=1, max_value=20, value=5)
if st.button(f"Generate and Save {technique} Images"):
if technique == "Scale":
specific_augmentations = scale_image(image, technique_count, scale_factor)
elif technique == "Translate":
specific_augmentations = translate_image(image, technique_count, translate_axis)
elif technique == "Crop":
specific_augmentations = crop_image(image, technique_count, crop_margin)
elif technique == "Rotate":
specific_augmentations = rotate_image(image, technique_count, rotate_angle)
elif technique == "Shear":
specific_augmentations = shear_image(image, technique_count, shear_factor)
# Save images to the specified folder
st.write(f"Generating and Saving {technique} Images...")
saved_files = save_images(specific_augmentations, technique.lower())
# Display generated images
st.write(f"Generated {technique} Augmented Images:")
for img in specific_augmentations:
st.image(img, use_container_width=True)
# Create a zip file
zip_buffer = create_zip(saved_files)
# Provide a download button for the zip file
st.download_button(
label=f"Download {technique} Augmented Images (ZIP)",
data=zip_buffer,
file_name=f"{technique.lower()}_augmented_images.zip",
mime="application/zip"
)