Cosmic-Ali/augmented_img_generator

1# External libraries2import streamlit as st3import numpy as np4import cv25 6#Python modules7import warnings8warnings.filterwarnings('ignore')9import os10from zipfile import ZipFile11import zlib     # for compression="zipfile.ZIP_DEFLATED"12import shutil13import gc       # for garbage collection14gc.enable()       15 16 17 18st.title("Augmented Image Generator")19 20 21# Transformation Functions:22 23def Translation(img,n=5):24    # for i in range(n):25        tx=np.random.randint(2,45)26        ty=np.random.randint(1,40)27        tm=np.array([[1,0,tx],[0,1,ty]],dtype=np.float32)28        tr_img=cv2.warpAffine(img,tm,dsize=(img.shape[1],img.shape[0]))29        return tr_img30        # save_image_and_show(trans_img)31 32 33def Rotation(img,n=5):34    # for i in range(n):35        angle=np.random.choice([90, 180, -90])36        scale=np.round(np.random.uniform(1,1.6),2)37        r_x=img.shape[0]//238        r_y=img.shape[1]//239        rm=cv2.getRotationMatrix2D((r_y,r_x),angle,scale)40        ro_img=cv2.warpAffine(img,rm,(img.shape[1],img.shape[0]))41        return ro_img42        # save_image_and_show(r_img)43 44 45def Scaling(img,n=5):46    # for i in range(n):47        sx=np.round(np.random.uniform(1,1.6),2)48        sy=np.round(np.random.uniform(1,1.6),2)49        tx=np.random.randint(1,16)50        ty=np.random.randint(1,12)51        sm=np.array([[sx,0,tx],[0,sy,ty]],dtype=np.float32)52        sc_img=cv2.warpAffine(img,sm,(img.shape[1],img.shape[0]))53        return sc_img54        # save_image_and_show(s_img)55 56 57def Shearing(img,n=5):58    # for i in range(n):59        sx=np.round(np.random.uniform(1,1.5),2)60        shx=np.round(np.random.uniform(0,0.4))61        tx=np.random.randint(1,15)62        shy=np.round(np.random.uniform(0,0.4))63        sy=np.round(np.random.uniform(1,1.2),2)64        ty=np.random.randint(1,20)65        shm=np.array([[sx,shx,tx],[shy,sy,ty]],dtype=np.float32)66        sh_img=cv2.warpAffine(img,shm,(img.shape[1],img.shape[0]))67        return sh_img68        # save_image_and_show(img_shear)69 70 71def Cropping(img,n=5):72    # for i in range(n):73        x_1=np.random.randint(20,60)74        y_1=np.random.randint(110,160)75        x_2=np.random.randint(20,60)76        y_2=np.random.randint(110,160)77        cr_img=img[x_1:y_1,x_2:y_2]78        return cr_img79        # save_image_and_show(img_crop)80 81def Flip_h(img):82    fliphor_img = cv2.flip(img, 1) 83    return fliphor_img84 85def Flip_v(img):86    flipver_img = cv2.flip(img, 0)87    return flipver_img88 89 90# Function for combining and applying transformations on a single image91def combined_trans(trans_types,img):92    trans_img = img93    for i in trans_types:94        trans_img = eval(f"{i}(trans_img)")95    return trans_img96 97 98# Taking Files from the User99# img = cv2.imread("/Users/ali/Desktop/data_science/ML/computer_vision/Apple-logo-1977.jpg")    #example image100files = st.file_uploader("Upload image",type=["jpg","png","zip"],accept_multiple_files=True)    101# Future Improvements: Add functionality to accept compressed zip files and then extract them.102 103# Reading/Converting files to cv2.image()/array 104images = []105for i,file in enumerate(files):106 107    # 1) converting to a 1D uint8 NumPy array108    np_arr = np.frombuffer(file.read(), dtype=np.uint8)     #file.read() is to read bytes from the UploadedFile109 110    # 2) decoding into an OpenCV image (BGR color by default)111    img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)112 113    # Appending image arrays to images list114    images.append(img)115 116    files[i] = None     # Clearing files simultaneous to the conversion117 118# Deleting files from RAM119del files120# ;del file;del img;del np_arr  121 122 123 124with st.form(key='transformation'):125 126    trans_types = st.pills("**Select transformations**",options=["Translation","Rotation","Shearing","Cropping","Scaling","Flip_h","Flip_v"],selection_mode="multi")127 128    preview = st.form_submit_button("Preview")129    if preview:130        c1,c2 = st.columns(2)131        for i in range(2):132            with c1:133                st.image(combined_trans(trans_types,cv2.cvtColor(images[0],cv2.COLOR_BGR2RGB)))  #converting image defualt bgr to rgb because st displays rgb134            with c2:135                st.image(combined_trans(trans_types,cv2.cvtColor(images[0],cv2.COLOR_BGR2RGB)))  #converting image defualt bgr to rgb because st displays rgb136        st.info("Previewing 4 transformed images with the selected combination")137 138# with st.container():139#     trans = st.pills("Select transformations",options=["Translation","Rotation","Shearing","Cropping","Scaling"],selection_mode="multi")140#     def combined_trans(selected):141#         trans_img = img142#         for i in selected:143#             trans_img = eval(f"{i}(trans_img)")144#         return trans_img        145#     st.image(combined_trans(trans))146 147 148if len(trans_types)>0:149 150    with st.popover("Click to view selected transformations"):151        for n,i in enumerate(trans_types,start=1):152            st.write(f"{n}.{i}")153        st.write("(Press 'Preview' to confirm change in selection)")154 155    156    with st.form(key='download'):157        img_count = st.slider("Select number of images to download",1,50)158        confirm = st.form_submit_button("Confirm")159        160    if confirm:161 162        if not os.path.exists("data/augmented_images"):163            os.mkdir("data/augmented_images")164        # Creating a folder, and writing images files to it one by one165 166        for i,img in enumerate(images):167            for n in range(img_count):168                try:169                    cv2.imwrite(f"data/augmented_images/{i}{n}.jpg",combined_trans(trans_types,img))170                except AttributeError:  # instead of exception handling, write a condition using regex to extract only the acceptable files171                    continue172 173        # Zipping/archiving the augmented_images folder174        shutil.make_archive("data/augmented_images","zip",base_dir="data/augmented_images")175 176        with open("data/augmented_images.zip",'rb') as f:177            download_button = st.download_button("Download images",f,file_name="augmented_images.zip")178 179  180        os.remove("data/augmented_images.zip")181        shutil.rmtree("data/augmented_images")182        #Deleting the files from dir after user download completes183 184 185 186 187# improvements and bug fixes:188 189# 1. Display the total number of images that are going to be downloaded. 190#    And change the download slider text to "Select the number of augmented images per given image"191# 2. Instead of exception handling, write a condition using regex to extract only the acceptable files192# 3. Write a st.spinner to display loading animation when the images are loading193# 4. Fix the cropping function 194