Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import numpy as np | |
| import pandas as pd | |
| import os | |
| import zipfile | |
| import cv2 | |
| st.markdown(""" | |
| <style> | |
| /* Set full-page width and remove Streamlit default padding */ | |
| .stApp { | |
| background-color: #d0e8ff; | |
| } | |
| .main .block-container { | |
| padding-top: 2rem; | |
| padding-right: 2rem; | |
| padding-left: 2rem; | |
| padding-bottom: 2rem; | |
| max-width: 100%; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # Main Heading | |
| st.markdown('<h2 style="text-align: center;"><span style="color:#003366;">Generate Augmented Images </span></h2>', unsafe_allow_html=True) | |
| st.markdown("<br><br>", unsafe_allow_html=True) | |
| #st.markdown('<h4 style="text-align: center;"><span style="color:#003366;">Upload the image </span></h4>', unsafe_allow_html=True) | |
| # Upload image | |
| uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"]) | |
| image_count=st.number_input("Number of Images to be generated",1,50,placeholder="Type a number...") | |
| st.write("Select the Affine Transformation to be applied:") | |
| trans = st.checkbox("Translation") | |
| rot = st.checkbox("Rotation") | |
| scal = st.checkbox("Scaling") | |
| shear=st.checkbox("Shearing") | |
| crop=st.checkbox("Crop") | |
| transformation=[] | |
| btn=st.button("Preview Image") | |
| #sub=st.button("Preview Image") | |
| def clear_folder(folder_name="output_images"): | |
| if os.path.exists(folder_name): | |
| for file in os.listdir(folder_name): | |
| file_path=os.path.join(folder_name,file) | |
| if os.path.isfile(file_path): | |
| os.remove(file_path) | |
| else: | |
| os.makedirs(folder_name) | |
| def save_image_and_show(image, folder_name="output_images"): | |
| if not os.path.exists(folder_name): | |
| os.makedirs(folder_name) | |
| image_name = f"image_{len(os.listdir(folder_name)) + 1}.jpg" | |
| image_path = os.path.join(folder_name, image_name) | |
| # Save the image | |
| cv2.imwrite(image_path, image) | |
| if btn: | |
| # Showing the success message | |
| st.success(f"Image saved to {image_path} (local path)") | |
| # Show the image to the user (not the path) | |
| st.image(image, caption="Saved Image", channels="BGR") | |
| def zip_file(folder_name="output_images"): | |
| zip_filename = f"{folder_name}.zip" | |
| with zipfile.ZipFile(zip_filename, 'w') as zipf: | |
| for root, dirs, files in os.walk(folder_name): | |
| for file in files: | |
| filepath = os.path.join(root, file) | |
| arcname = os.path.relpath(filepath, folder_name) | |
| zipf.write(filepath, arcname) | |
| # Read zip for download | |
| with open(zip_filename, "rb") as f: | |
| st.download_button( | |
| label="Download All Images as ZIP", | |
| data=f, | |
| file_name=zip_filename, | |
| mime="application/zip" | |
| ) | |
| def translation(img,n=5): | |
| for i in range(n): | |
| tx=np.random.randint(2,50) | |
| ty=np.random.randint(1,50) | |
| tm=np.array([[1,0,tx],[0,1,ty]],dtype=np.float32) | |
| trans_img=cv2.warpAffine(img,tm,dsize=(img.shape[1],img.shape[0]),borderMode=cv2.BORDER_REFLECT) | |
| #img_trans_rgb = cv2.cvtColor(trans_img, cv2.COLOR_BGR2RGB) | |
| save_image_and_show(trans_img) | |
| #st.image(img_trans_rgb, caption="Translation Applied", use_container_width=True) | |
| def rotation(img,n=5): | |
| for i in range(n): | |
| angle=np.random.choice([30,45,90, 180, -90,-30,-45]) | |
| scale=np.round(np.random.uniform(1,1.6),2) | |
| r_x=img.shape[0]//2 | |
| r_y=img.shape[1]//2 | |
| rm=cv2.getRotationMatrix2D((r_y,r_x),angle,scale) | |
| r_img=cv2.warpAffine(img,rm,(img.shape[1],img.shape[0])) | |
| #img_rot_rgb = cv2.cvtColor(r_img, cv2.COLOR_BGR2RGB) | |
| save_image_and_show(r_img) | |
| #st.image(img_rot_rgb, caption="Rotation Applied", use_container_width=True) | |
| def scaling(img,n=5): | |
| for i in range(n): | |
| sx=np.random.uniform(1,2) | |
| sy=np.random.uniform(1,2) | |
| tx=np.random.randint(1,16) | |
| ty=np.random.randint(1,12) | |
| sm=np.array([[sx,0,tx],[0,sy,ty]],dtype=np.float32) | |
| s_img=cv2.warpAffine(img,sm,(img.shape[1],img.shape[0])) | |
| #img_scal_rgb=cv2.cvtColor(s_img,cv2.COLOR_BGR2RGB) | |
| save_image_and_show(s_img) | |
| #st.image(img_scal_rgb, caption="Scaling Applied", use_container_width=True) | |
| def shearing(img,n=5): | |
| for i in range(n): | |
| sx=np.round(np.random.uniform(1,1.5),2) | |
| shx=np.round(np.random.uniform(0,0.5)) | |
| tx=np.random.randint(1,10) | |
| shy=np.round(np.random.uniform(0,0.5)) | |
| sy=np.round(np.random.uniform(1,1.2),2) | |
| ty=np.random.randint(1,15) | |
| shm=np.array([[sx,shx,tx],[shy,sy,ty]],dtype=np.float32) | |
| img_shear=cv2.warpAffine(img,shm,(img.shape[1],img.shape[0]),borderMode=cv2.BORDER_REFLECT) | |
| save_image_and_show(img_shear) | |
| #img_shear_rgb=cv2.cvtColor(img_all,cv2.COLOR_BGR2RGB) | |
| #st.image(img_shear_rgb, caption="Shearing Applied", use_container_width=True) | |
| def cropping(img,n=5): | |
| for i in range(n): | |
| x_1=np.random.randint(20,60) | |
| y_1=np.random.randint(110,160) | |
| x_2=np.random.randint(20,60) | |
| y_2=np.random.randint(110,160) | |
| img_crop=img[x_1:y_1,x_2:y_2] | |
| save_image_and_show(img_crop) | |
| # img_crop_rgb=cv2.cvtColor(img_crop,cv2.COLOR_BGR2RGB) | |
| # st.image(img_crop_rgb, caption="Shearing Applied", use_container_width=True) | |
| # def save_image(image, folder_name="output_images"): | |
| # if not os.path.exists(folder_name): | |
| # os.makedirs(folder_name) | |
| # image_name = f"image_{len(os.listdir(folder_name)) + 1}.jpg" | |
| # image_path = os.path.join(folder_name, image_name) | |
| # cv2.imwrite(image_path, image) | |
| # st.success(f"Image saved to {image_path}") | |
| if uploaded_file is not None: | |
| # Converting the image into array | |
| img_array=np.asarray(bytearray(uploaded_file.read()),dtype=np.uint8) | |
| # Decode or converting the array back to image array | |
| img=cv2.imdecode(img_array,cv2.IMREAD_COLOR) | |
| # # Convert BGR to RGB as streamlit expects RGB colorspace | |
| # img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| transformation = [] | |
| if trans: | |
| transformation.append("Translation") | |
| if rot: | |
| transformation.append("Rotation") | |
| if scal: | |
| transformation.append("Scaling") | |
| if shear: | |
| transformation.append("Shearing") | |
| if crop: | |
| transformation.append("Crop") | |
| if transformation: | |
| clear_folder() | |
| # per_trans = max(1,(image_count // len(transformation))) | |
| # trans_count=[per_trans+1 if (per_trans * len(transformation))!= image_count else per_trans][0] | |
| num_trans = len(transformation) | |
| base = image_count // num_trans | |
| remainder = image_count % num_trans | |
| # Create a list like [base+1, base, base, ...] to ensure sum == image_count | |
| counts = [base + 1 if i < remainder else base for i in range(num_trans)] | |
| # if "Translation" in transformation: | |
| # translation(img,trans_count) | |
| # if "Rotation" in transformation: | |
| # rotation(img, per_trans) | |
| # if "Scaling" in transformation: | |
| # scaling(img, per_trans) | |
| # if "Shearing" in transformation: | |
| # shearing(img, per_trans) | |
| # if "Crop" in transformation: | |
| # cropping(img, per_trans) | |
| for t_name, t_count in zip(transformation, counts): | |
| if t_name == "Translation": | |
| translation(img, t_count) | |
| elif t_name == "Rotation": | |
| rotation(img, t_count) | |
| elif t_name == "Scaling": | |
| scaling(img, t_count) | |
| elif t_name == "Shearing": | |
| shearing(img, t_count) | |
| elif t_name == "Crop": | |
| cropping(img, t_count) | |
| zip_file() | |
| #st.download_button("Download File",img_rgb) | |
| # # Displaying image in streamlit | |
| # st.image(img_rgb, caption="Uploaded Image (OpenCV)", use_container_width=True) | |