Python code for processing datasets

Image and label

Zoom all images in a folder

import numpy as np
import glob
import os
from PIL import Image
#image_path the path to the folder where the image is located 
#new_image_path indicates the folder path for storing new images 
#new_size a list with a length of 2, where the two elements are width and height, respectively 
def img_resize(image_path,new_image_path,new_size):
#find out image_path all under .jpg/png/jpeg image path and sorting 
imgs = sorted(glob.glob(image_path + "*.jpg") )
print(len(imgs))
for img in imgs:
  image = Image.open(img)
  image=image.resize(new_size,Image.BILINEAR)
  print(image.size)
  # os.path.splitext () separate file names and extensions 
  # os.path.split () return the path and file name of the file 
  img_name=os.path.split(img)[1]
  print(img_name)
  image.save(new_image_path+img_name) #save image 

Use a threshold to convert the label graph in the folder into a binary graph

import numpy as np
import glob
import os
from PIL import Image
#this function is used to convert a label graph with only one category into a binary graph (because it uses ps the resulting image is not a binary image, and there may be grayscale transitions at the edges, unless otherwise specified labelme create label images) 
#max_value the larger value representing the binary graph is set to 1 by default, samll_value the smaller value representing the binary graph, defaults to 0 
#new_size indicates whether to scale the image, default to False if you want to scale the image, send a list with a length of 2 (width and height) 
def annotation_prepare(ann_path,new_ann_path,threshold,new_size=False,max_value=1,samll_value=0): #a binary graph used to convert a label to 1 or 255 
anns = sorted(glob.glob(ann_path + "*.jpg"))
for ann in anns:
  ann_name = os.path.split(ann)[1] #obtain image names with extensions for naming subsequent images 
  print(ann_name)
  ann = Image.open(ann)
  ann = ann.convert("L") #change to grayscale image 
  if (new_size!=False):  #new_size!=False indicates the need to scale the image, 
      w, h = new_size
      ann=ann.resize((w,h),Image.BILINEAR)
  else:
      w,h=ann.size
  ann=np.array(ann)
  #perform threshold conversion 
  new_ann=(ann>threshold)*max_value
  new_ann=new_ann+(ann>=threshold)*samll_value
  new_ann = new_ann.astype(np.uint8)  #restore to uint 8 format, images need to be saved in this format 
  new_ann = Image.fromarray(new_ann)  #convert from array format to image format, which can be used save function 
  new_ann.save(new_ann_path+ann_name) #save image 

Renaming label graphs


#used in cases where the label image has some additional suffixes than the original image name (excluding extensions), and one original image corresponds to one or more label images. excess suffixes in the label graph can be removed and replaced with a unified one _0N substitute (as originally named 10) .jpg and its corresponding label image is 10 _water_1,10_w_2. the label image can be renamed to 10 _01,10_02)。
#remove the excess original images as well, because if the dataset is messy, some original images may correspond to 0 labeled images 
#img_path the path to the folder where the original image is located 
#ann_path the path to the folder where the label image is located 
def annNameChange(img_path,ann_path,new_img_path,new_ann_path): #a binary graph used to convert a label to 1 or 255 new_ann_path,
imgs = sorted(glob.glob(img_path+"*.jpg"))
for img in imgs:
  img_save=Image.open(img)  #read img image for saving transfer path later 
  img = os.path.split(img)[1] #obtain image names with suffixes 
  img_name=os.path.splitext(img)[0]  #remove extension 
  print(img_name)
  ann_list=glob.glob(ann_path+img_name+"[_.]*")   #identify all label images corresponding to this image. [ ] represent matching [] any character in the. here, we need to modify the matching rules based on the name features of our own images 
  print(len(ann_list))
  i=1
  for ann in ann_list:
      ann = Image.open(ann) #reading images is for resaving below 
      new_img_name=new_img_path+'/'+img_name+'.jpg'
      new_ann_name=new_ann_path+'/'+img_name+"_0{}.jpg".format(i) #this is also the key point. 
      print(new_ann_name)
      i =i+1

      img_save.save(new_img_name) #resave image 
      ann.save(new_ann_name)

Renaming and saving both the original and label images

#similar to above, but here we are renaming both the original image and the label image, starting from 1 n zhang yuantu was named n.jpg and its corresponding label graph is named n_0x.jpg
def annNameChange_new(img_path,ann_path,new_img_path,new_ann_path): #a binary graph used to convert a label to 1 or 255 new_ann_path,
imgs = sorted(glob.glob(img_path+"*.jpg"))
i_name=0    #indicates that the processing is being processed i_name zhang tu, named after i_name name 
for img in imgs:
  i_name = i_name+1
  img_save=Image.open(img)  #read img image for saving transfer path later 
  img = os.path.split(img)[1] #obtain image names with suffixes 
  img_name=os.path.splitext(img)[0]  #remove extension 
  print(img_name)
  ann_list=glob.glob(ann_path+img_name+"[_.]*")   #identify all label images corresponding to this image. [ ] represent matching [] any character in 
  print(len(ann_list))
  i=1
  for ann in ann_list:
      ann = Image.open(ann) #reading images is for resaving below 
      new_img_name=new_img_path+'/{}.jpg'.format(i_name)
      new_ann_name=new_ann_path+'/{}_0{}.jpg'.format(i_name,i) #this is also the key point. 
      print(new_img_path)
      print(new_ann_name)
      i =i+1
      img_save.save(new_img_name) #resave image 
      ann.save(new_ann_name)

Take out the images of the dataset, corresponding label images, and corresponding names and place them in the list separately

#the following is my custom function, which is used to retrieve only the image for testing and retrieve the corresponding name to save the final result. used for water data set 
#the following is my custom function, which is used to retrieve only the image for testing and retrieve the corresponding name to save the final result. used for water data set 
def get_image(img_path,ann_path):
img_list = glob.glob(img_path+r'*')
mask_list=glob.glob(ann_path+r'*')
print(len(img_list))
print('mask_list',len(mask_list))
imgs = []
names=[]
masks=[]
for img_path in img_list:
  names.append(os.path.basename(img_path))
  img=Image.open(img_path)
  imgs.append(img)
for mask_path in mask_list:
  mask = Image.open(mask_path)  #read images, read 
  masks.append(mask)
return names,imgs,masks

Simultaneously resize the img and corresponding mask to the same size and save it


"""used to simultaneously transport img corresponding to mask resize to unified size"""
def resize_image(args):
  img_names,imgs=get_image(args.imgs_dir) #get_image in the previous text 
  mask_names,masks=get_image(args.masks_dir)
  assert len(imgs)==len(masks), 'len(imgs_dir) and len(masks_dir) must be the same'
  for i in range(len(imgs)):
      re_img = imgs[i].resize((args.w, args.h), Image.BICUBIC) # args.w,args.h refers to width and height 
      re_mask= masks[i].resize((args.w, args.h), Image.BICUBIC) #after interpolation, it needs to go through threshold processing again 
      re_label_1,re_label_255=make_label(re_mask,args.Threshold)  # make_label in the following text 
      re_img.save(os.path.join(args.re_img_dir, img_names[i]))
      re_label_1.save(os.path.join(args.re_mask_dir_new, mask_names[i]))
      re_label_255.save(os.path.join(args.re_mask_dir, mask_names[i]))
#this function is used to convert mask after threshold processing, the images are returned with grayscale values of 1 and 255, respectively 
def make_label(mask,Threshold):
mask = mask.convert("L")  #convert open images to grayscale images 
mask = np.array(mask).astype(int)  #convert the opened image into an array and change the type from uint 8 becomes int 32, because pixel addition is required below, use uint 8 will exceed the range 
label_1 = (mask > Threshold) *1  #select a threshold and change the pixel to 1. if there are multiple classes, the others should be changed to 2, 3, 4, etc 
label_255=label_1*255
# label2 = np.zeros((512,512))
# label_1=label1+label2
label_1 = label_1.astype(np.uint8)  #restore to uint 8 format, images need to be saved in this format 
label_255 = label_255.astype(np.uint8)  #restore to uint 8 format, images need to be saved in this format 
# label_1=np.where(mask > Threshold, 1, 0)
# label_255 = label_1 * 255
label_1 = Image.fromarray(label_1)  #convert from array format to image format, which can be used save function 
label_255= Image.fromarray(label_255)
return label_1,label_255

Used to divide a dataset into training, validation, and testing sets and store them (randomly proportional)

import random
def split_dataset(img_path,ann_path,train_dir,val_dir,test_dir=None,is_split_test=False):
img_names,imgs,masks=get_image(img_path,ann_path)  # get_image in the previous text, there are 
	#create some directories 
train_imgs_dir=os.path.join(train_dir,'imgs')
train_anns_dir = os.path.join(train_dir,'anns')
val_imgs_dir=os.path.join(val_dir,'imgs')
val_anns_dir = os.path.join(val_dir,'anns')
mkdir(train_imgs_dir)
mkdir(train_anns_dir)
mkdir(val_imgs_dir)
mkdir(val_anns_dir)
if test_dir:
  test_imgs_dir = os.path.join(train_dir, 'imgs')
  test_anns_dir = os.path.join(train_dir, 'anns')
  mkdir(test_imgs_dir)
  mkdir(test_anns_dir)
assert len(masks) == len(imgs), 'len(imgs_dir) and len(masks_dir) must be the same'
indexs = [i for i in range(len(imgs))]
print(len(imgs))
random.shuffle(indexs)  #scrambling indexes 
print(indexs)
print(type(indexs))
if is_split_test:  #if you want to divide the test set, the ratio is 6 :2:2
  cnt_train = int(round(len(imgs) * 0.6, 0))  #rounding, 0 represents retaining 0 decimal points 
  print(cnt_train)
  cnt_val = int(round(len(imgs) * 0.2, 0))  #because indexing is required, cast to integer 
  print(cnt_val)
  cnt_test = int(round(len(imgs) * 0.2, 0))
  #the following uses list slicing to obtain train/val/test index of data 
  train_indexs = indexs[0:cnt_train]
  print(train_indexs)
  val_indexs = indexs[cnt_train:(cnt_train + cnt_val):1]
  test_indexs = indexs[(cnt_train + cnt_val):]
else:  #only the training set and validation set are divided, and 7 are used here :3
  cnt_train = int(round(len(imgs) * 0.7, 0))  #rounding, 0 represents retaining 0 decimal points 
  cnt_val = int(round(len(imgs) * 0.3, 0))  #because indexing is required, cast to integer 
  train_indexs = indexs[0:cnt_train]
  print(train_indexs)
  val_indexs = indexs[cnt_train:]
	#save image 
for i in train_indexs:
  imgs[i].save(os.path.join(train_imgs_dir, img_names[i]))
  # masks[i].save(os.path.join(args.train_dir, 'annotation', mask_names[i]))
  masks[i].save(os.path.join(train_anns_dir,img_names[i]))
for i in val_indexs:
  imgs[i].save(os.path.join(val_imgs_dir,img_names[i]))
  # masks[i].save(os.path.join(args.val_dir, 'annotation', mask_names[i]))
  masks[i].save(os.path.join(val_anns_dir, img_names[i]))
if is_split_test:
  for i in test_indexs:
      imgs[i].save(os.path.join(test_imgs_dir, img_names[i]))
      # masks[i].save(os.path.join(args.val_dir, 'annotation', mask_names[i]))
      masks[i].save(os.path.join(test_anns_dir, img_names[i]))
def mkdir(path):   #used to create a directory 
#introducing modules 
#remove the first blank space 
path = path.strip()
#remove tail  symbol 
path = path.rstrip("\")
#determine if the path exists 
#existence True
#not present False
isExists = os.path.exists(path)
#judging the result 
if not isExists:
  #if it does not exist, create a directory 
  #create directory operation function 
  os.makedirs(path)
  return True
else:
  #if the directory exists, do not create it and prompt that the directory already exists 
  return False

Semantic segmentation dataset data augmentation

import Augmentor
from glob import glob
import os
import random
from PIL import Image
import cv2
import numpy as np
def augmentor(image_path,label_path,num):
#num indicates the number of images you want to generate 
p = Augmentor.Pipeline(image_path)
# Point to a directory containing ground truth data.
# Images with the same file names will be added as ground truth data
# and augmented in parallel to the original data.
p.ground_truth(label_path)
# Add operations to the pipeline as normal:
p.rotate(probability=0.5, max_left_rotation=10, max_right_rotation=10)
p.flip_left_right(probability=0.5)
p.zoom_random(probability=0.5, percentage_area=0.8) #randomly zoom in and out, but maintain the original size 
p.flip_top_bottom(probability=0.5)
p.rotate_random_90(probability=0.5)
p.random_distortion(probability=0.2, grid_width=4, grid_height=4, magnitude=8)  #indicates random distortion and some places 
p.sample(num)

Merge label images

'''for a Groundtruth label each object separately in a single image. this function is used for merging mask to a picture'''
#the following is my custom function, which is used to retrieve only images for testing and retrieve corresponding names to save the final result. used for water data set 
import glob
import os
import numpy as np
from PIL import Image
def merge(imgpath_list):
mask=Image.open(imgpath_list[0])
mask = mask.convert('L')  #make the label graph a channel 
mask=np.array(mask)
# print(mask.shape)
for _mask in imgpath_list[1:]:
  _mask = Image.open(_mask)
  _mask = _mask.convert('L')
  _mask = np.array(_mask)
  mask = np.maximum(mask, _mask)  #merge two images 
mask=Image.fromarray(mask)
return mask
def merge_mask(img_path,ann_path,new_annpath):
img_list = glob.glob(img_path+'*')
print(len(img_list))
for img_path in img_list:
  name=os.path.basename(img_path)
  ann_img=os.path.join(ann_path,name)
  mask_list=glob.glob(ann_img[0:-4]+'*')  #develop matching rules to obtain img_path all corresponding Mask chart 
  print(name,mask_list)
  if(len(mask_list)>1):
      mask=merge(mask_list)
  else:
      mask=Image.open(mask_list[0])
  mask.save(os.path.join(new_annpath,name))

Use PIL library to determine image size and rotate

img=Image.open(img_path)
h,w=img.size
if h>w:
	img=img.transpose(Image.ROTATE_90)
#if in a dataset, all images are of the same size, but some are vertical and some are horizontal, the following code can be used to unify others 
import os
from PIL import Image
cascade_mask_rcnn_water_train_imgs_dir=r'F:Datasettraincolor_img'
cascade_mask_rcnn_water_train_marks_dir=r'F:JunDatasettrainbinary_img'
cascade_mask_rcnn_water_val_imgs_dir=r'F:JunDatasetvalcolor_img_all'
cascade_mask_rcnn_water_val_marks_dir=r'F:JunDatasetvalbinary_img_all'
def rotate_image(img_path,save_path):
imgs_path = os.listdir(img_path)
for r in imgs_path:
  img=os.path.join(img_path,r)
  img=Image.open(img)
  w, h = img.size
  if w>h:
      img = img.transpose(Image.ROTATE_90)
  img.save(os.path.join(save_path,r))

Related articles